" + html.escape(text) + "" + + async def map(self) -> ProjectionInputs: + """Produce file entities for strict conversion without any provider network calls.""" + data = self.record.payload + if self.record.deleted_at is not None or self.record.completeness not in { + "complete", + "partial", + }: + raise ValueError("Gmail projection requires an active captured source record") + if data.get("id") != self.record.identity.native_id or not data.get("threadId"): + raise ValueError("Gmail projection identity is inconsistent") + body = await self.render(data["payload"], "/payload") + local = await write_blob(body.encode("utf-8"), self.directory, suffix=".html") + entity = GmailMessageEntity.from_api(data, thread_id=data["threadId"], breadcrumbs=[]) + entity.local_path = str(local) + entity.size = len(body.encode("utf-8")) + internal = datetime.fromtimestamp(int(data["internalDate"]) / 1000, timezone.utc) + entity.internal_date = internal + entity.internal_timestamp = internal + if entity.sent_at.tzinfo is None: + entity.sent_at = entity.sent_at.replace(tzinfo=timezone.utc) + for field in ("to", "cc", "bcc"): + addresses = getaddresses([_header(data["payload"], field)]) + setattr(entity, field, [formataddr(address) for address in addresses if address[1]]) + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, + key="/payload/body", + kind="body", + media_type="text/html", + extension=".html", + charset_recoveries=tuple(self.charset_recoveries), + ), + entity=entity, + ), + *self.attachments, + ) + ) + + +async def map_gmail( + record: SourceRecord, + storage: StorageBackend, + directory: Path, +) -> ProjectionInputs: + """Map immutable Gmail capture into existing message and attachment entities.""" + return await GmailProjection(record, storage, directory).map() diff --git a/backend/airweave/domains/entities/canonical/linear_projection.py b/backend/airweave/domains/entities/canonical/linear_projection.py new file mode 100644 index 0000000000..374729606a --- /dev/null +++ b/backend/airweave/domains/entities/canonical/linear_projection.py @@ -0,0 +1,152 @@ +"""Offline Linear views of native originals and verified immutable files.""" + +import mimetypes +from pathlib import Path +from urllib.parse import urlsplit +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field + +from airweave.domains.entities.canonical.blob_materializer import read_blob, write_blob +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.storage.protocols import StorageBackend +from airweave.platform.entities._base import BaseEntity, Breadcrumb +from airweave.platform.entities.linear import ( + LinearAttachmentEntity, + LinearCommentEntity, + LinearIssueEntity, + LinearLinkedAttachmentEntity, +) + + +class _Named(BaseModel): + model_config = ConfigDict(extra="ignore") + id: UUID + name: str + + +class _IssueContext(BaseModel): + model_config = ConfigDict(extra="ignore") + id: UUID + identifier: str = Field(min_length=1) + title: str + url: str + team: _Named + project: _Named | None = None + + +class _Original(BaseModel): + model_config = ConfigDict(extra="ignore") + id: UUID + createdAt: AwareDatetime + updatedAt: AwareDatetime + url: str + + +class _Child(_Original): + issue: _IssueContext + + +class _Comment(_Child): + body: str + + +class _Attachment(_Child): + title: str + + +def _breadcrumbs(issue: _IssueContext) -> list[Breadcrumb]: + return [ + Breadcrumb(entity_id=str(issue.id), name=issue.identifier, entity_type="LinearIssueEntity") + ] + + +async def map_linear( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> tuple[BaseEntity, ...]: + """Map only retained observations, rejecting identity drift or missing native context.""" + original = _Original.model_validate(record.payload) + if str(original.id) != record.identity.native_id: + raise ValueError("Linear record identity differs from captured original") + kind = record.identity.record_type + if kind == "issue": + _IssueContext.model_validate(record.payload) + if record.identity.container_id is not None: + raise ValueError("Linear issue unexpectedly has a parent container") + issue = LinearIssueEntity.from_api(record.payload) + issue.web_url_value = original.url + return (issue,) + if kind not in {"comment", "attachment"}: + raise ValueError("Unsupported Linear original kind") + child = _Child.model_validate(record.payload) + context = child.issue + if str(context.id) != record.identity.container_id: + raise ValueError("Linear child parent differs from canonical container") + breadcrumbs = _breadcrumbs(context) + if kind == "comment": + _Comment.model_validate(record.payload) + entity = LinearCommentEntity.from_api( + record.payload, + issue_id=str(context.id), + issue_identifier=context.identifier, + breadcrumbs=breadcrumbs, + team_id=str(context.team.id), + team_name=context.team.name, + project_id=str(context.project.id) if context.project else None, + project_name=context.project.name if context.project else None, + ) + entity.web_url_value = original.url + return (entity,) + return await _attachment(record, storage, directory) + + +async def _attachment( + record: SourceRecord, + storage: StorageBackend, + directory: Path, +) -> tuple[BaseEntity, ...]: + attachment = _Attachment.model_validate(record.payload) + context = attachment.issue + breadcrumbs = _breadcrumbs(context) + if not record.blobs: + if record.completeness == "complete": + raise ValueError("Linear attachment claims complete content without retained bytes") + return ( + LinearLinkedAttachmentEntity( + attachment_id=str(attachment.id), + title=attachment.title, + issue_identifier=context.identifier, + target_url=attachment.url, + content_coverage="Link metadata only; file content is not retained", + breadcrumbs=breadcrumbs, + created_at=attachment.createdAt, + updated_at=attachment.updatedAt, + ), + ) + if len(record.blobs) != 1 or record.blobs[0].source_path != "/url": + raise ValueError("Linear attachment requires one exact URL blob reference") + ref = record.blobs[0] + content = await read_blob(record, ref, storage) + suffix = mimetypes.guess_extension(ref.media_type or "") + if not suffix: + suffix = Path(urlsplit(attachment.url).path).suffix.lower() + if not suffix: + raise ValueError("Retained Linear file has no known format for projection") + local_path = await write_blob(content, directory, suffix=suffix) + return ( + LinearAttachmentEntity( + attachment_id=str(attachment.id), + issue_id=str(context.id), + issue_identifier=context.identifier, + title=attachment.title, + url=attachment.url, + size=len(content), + file_type=suffix.lstrip("."), + mime_type=ref.media_type, + local_path=str(local_path), + breadcrumbs=breadcrumbs, + web_url_value=attachment.url, + created_at=attachment.createdAt, + updated_at=attachment.updatedAt, + ), + ) diff --git a/backend/airweave/domains/entities/canonical/mime_text.py b/backend/airweave/domains/entities/canonical/mime_text.py new file mode 100644 index 0000000000..c7163f277b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/mime_text.py @@ -0,0 +1,29 @@ +"""Strict MIME text decoding with explicit, bounded recovery evidence.""" + +from pydantic import BaseModel, ConfigDict + +from airweave.domains.entities.canonical.extraction_models import CharsetRecovery + + +class DecodedText(BaseModel): + """Derived text and any charset correction; original bytes remain authoritative.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + text: str + recovery: CharsetRecovery | None = None + + +def decode_mime_text(content: bytes, charset: str, source_path: str) -> DecodedText: + """Honor the charset first; recover only bytes that are strictly valid UTF-8. + + Unknown codecs and content invalid in both encodings still fail. No detector, + replacement characters or discarded bytes can turn a failure into success. + """ + try: + return DecodedText(text=content.decode(charset, errors="strict")) + except UnicodeDecodeError: + text = content.decode("utf-8", errors="strict") + return DecodedText( + text=text, + recovery=CharsetRecovery(source_path=source_path, from_charset=charset), + ) diff --git a/backend/airweave/domains/entities/canonical/models.py b/backend/airweave/domains/entities/canonical/models.py new file mode 100644 index 0000000000..1c0bb36f68 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/models.py @@ -0,0 +1,90 @@ +"""Canonical record results. Journal snapshots are historical, reads are current.""" + +from datetime import datetime +from typing import Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, JsonValue + +from airweave.domains.entities.canonical.coverage_models import CaptureCoverage +from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity + + +class SourceRecord(BaseModel): + """Durable source state, independent of search projection success.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + id: UUID + sync_id: UUID + identity: RecordIdentity + parent: RecordIdentity | None = None + revision: int + payload: dict[str, JsonValue] + content_access: Literal["available", "unavailable"] = "available" + payload_schema_version: int + capture_hash: str + content_hash: str | None + completeness: str + observed_at: datetime + source_created_at: datetime | None + source_updated_at: datetime | None + deleted_at: datetime | None + removal_reason: str | None + blobs: tuple[BlobReference, ...] + indexed_revision: int | None + indexed_pipeline_version: int | None + + +class IndexedRecordRead(SourceRecord): + """Current original plus independently attested current extraction coverage.""" + + extraction: ExtractionCoverage | None = None + + +class ObservedChange(BaseModel): + """Snapshot at commit time; never a claim to every upstream transient event.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + sequence: int + kind: Literal["upsert", "delete"] + record: SourceRecord + + +class CaptureResult(BaseModel): + """Journal events created by this commit and its final per-sync position.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + changes: tuple[ObservedChange, ...] + sequence: int + unchanged: int + + +class ReconcileResult(BaseModel): + """Repeat while more remains; only then may the source commit its checkpoint.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + capture: CaptureResult + has_more: bool + + +class ChangePage(BaseModel): + """All committed changes after a position, capped at a fixed high watermark.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + changes: tuple[ObservedChange, ...] + next_sequence: int + high_watermark: int + has_more: bool + + +class RecordPage(BaseModel): + """Committed records with truthful traversal semantics.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + capture: CaptureCoverage | None = None + records: tuple[SourceRecord, ...] + next_cursor: str | None + has_more: bool + consistency: Literal["live"] = "live" + order: Literal["id_asc"] = "id_asc" diff --git a/backend/airweave/domains/entities/canonical/notion_projection.py b/backend/airweave/domains/entities/canonical/notion_projection.py new file mode 100644 index 0000000000..9d1e6cce65 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/notion_projection.py @@ -0,0 +1,258 @@ +"""Offline, record-local Notion search text; no page assembly or provider access.""" + +from typing import Literal +from uuid import UUID + +from pydantic import ( + AwareDatetime, + BaseModel, + ConfigDict, + HttpUrl, + JsonValue, + StrictBool, + TypeAdapter, + field_validator, +) + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import ProjectionMappingError +from airweave.platform.entities.notion import NotionOriginalEntity + + +class _Native(BaseModel): + model_config = ConfigDict(extra="ignore") + + +class _RichText(_Native): + type: Literal["text", "mention", "equation"] + plain_text: str + + +class _Original(_Native): + id: UUID + object: Literal["page", "database", "data_source", "block"] + created_time: AwareDatetime + last_edited_time: AwareDatetime + in_trash: StrictBool + url: str | None = None + + @field_validator("url") + @classmethod + def native_url(cls, value: str | None) -> str | None: + if value is not None: + parsed = TypeAdapter(HttpUrl).validate_python(value) + if ( + parsed.scheme != "https" + or parsed.host not in {"app.notion.com", "notion.so", "www.notion.so"} + or parsed.username is not None + or parsed.password is not None + or parsed.port not in (None, 443) + ): + raise ProjectionMappingError("Notion original URL is outside audited native hosts") + return value + + +class _Property(_Native): + type: str + title: list[_RichText] | None = None + + +class _Page(_Original): + properties: dict[str, _Property] + + +class _Database(_Original): + title: list[_RichText] + description: list[_RichText] + + +class _SchemaProperty(_Native): + type: str + + +class _DataSource(_Database): + properties: dict[str, _SchemaProperty] + + +class _Parent(_Native): + type: Literal["page_id", "block_id"] + page_id: UUID | None = None + block_id: UUID | None = None + + +class _Block(_Original): + type: str + parent: _Parent + has_children: StrictBool + + +class _Text(_Native): + rich_text: list[_RichText] + + +class _Code(_Text): + language: str + caption: list[_RichText] + + +class _Todo(_Text): + checked: StrictBool + + +class _TableRow(_Native): + cells: list[list[_RichText]] + + +class _Reference(_Native): + title: str + + +class _Equation(_Native): + expression: str + + +class _Table(_Native): + table_width: int + has_column_header: StrictBool + has_row_header: StrictBool + + +class _Synced(_Native): + synced_from: dict[str, str] | None + + +def _text(items: list[_RichText]) -> str: + # Native fragments include their own spacing. Mentions use only retained labels. + return "".join(item.plain_text for item in items) + + +def _block_text(record: SourceRecord, block: _Block) -> tuple[str, str, str]: + parent = record.parent + native_parent = ( + block.parent.page_id if block.parent.type == "page_id" else block.parent.block_id + ) + expected_kind = "page" if block.parent.type == "page_id" else "block" + if ( + parent is None + or parent.record_type != expected_kind + or native_parent is None + or parent.native_id != str(native_parent) + or parent.container_id is not None + ): + raise ProjectionMappingError("Notion block parent differs from captured authority") + kind = block.type + content = record.payload.get(kind) + coverage = "This block only; child blocks, comments and file contents are not included." + if kind in { + "paragraph", + "heading_1", + "heading_2", + "heading_3", + "bulleted_list_item", + "numbered_list_item", + "quote", + "toggle", + "callout", + }: + text = _text(_Text.model_validate(content).rich_text) + elif kind == "to_do": + todo = _Todo.model_validate(content) + text = ("[x] " if todo.checked else "[ ] ") + _text(todo.rich_text) + elif kind == "code": + code = _Code.model_validate(content) + text = "\n".join( + part + for part in ( + code.language, + _text(code.rich_text), + _text(code.caption), + ) + if part + ) + elif kind == "table_row": + text = " | ".join(_text(cell) for cell in _TableRow.model_validate(content).cells) + elif kind == "equation": + text = _Equation.model_validate(content).expression + else: + return _structure(kind, content) + title = text.splitlines()[0][:160] if text else kind.replace("_", " ").capitalize() + return title, text, coverage + + +def _structure(kind: str, content: JsonValue) -> tuple[str, str, str]: + """Render known structure/reference labels without inventing their body.""" + if kind in {"child_page", "child_database"}: + label = _Reference.model_validate(content).title + return ( + label or "Untitled reference", + "", + "Reference label only; linked content is separate.", + ) + elif kind == "synced_block": + synced = _Synced.model_validate(content) + if synced.synced_from is not None: + raise ProjectionMappingError( + "Synced block reference content is not independently retained" + ) + return "Synced block", "", "Container only; child blocks are separate." + elif kind == "table": + _Table.model_validate(content) + return "Table", "", "Table container only; rows are separate." + elif kind in {"column", "column_list", "divider", "table_of_contents"}: + _Native.model_validate(content) + return kind.replace("_", " ").capitalize(), "", "Structure only; no body text." + else: + raise ProjectionMappingError("Notion block type has no supported text projection") + + +def map_notion(record: SourceRecord) -> tuple[NotionOriginalEntity, ...]: + """Map one retained native observation; missing/unknown content stays pending.""" + native = _Original.model_validate(record.payload) + kind = record.identity.record_type + if ( + str(native.id) != record.identity.native_id + or native.object != kind + or record.identity.container_id is not None + ): + raise ProjectionMappingError("Notion original identity differs from captured identity") + if native.in_trash: + raise ProjectionMappingError("Trashed Notion content cannot be projected") + if kind == "block": + title, text, coverage = _block_text(record, _Block.model_validate(record.payload)) + else: + if record.parent is not None: + raise ProjectionMappingError("Notion root has an unexpected capture parent") + coverage = ( + "Metadata only; page body, full properties, comments " + "and file contents are not included." + ) + if kind == "page": + page = _Page.model_validate(record.payload) + titles = [prop.title for prop in page.properties.values() if prop.type == "title"] + if len(titles) != 1 or titles[0] is None: + raise ProjectionMappingError("Notion page requires its retained title property") + title, text = _text(titles[0]), "" + elif kind == "database": + database = _Database.model_validate(record.payload) + title, text = _text(database.title), _text(database.description) + elif kind == "data_source": + data = _DataSource.model_validate(record.payload) + title = _text(data.title) + schema = "\n".join(f"{name} ({prop.type})" for name, prop in data.properties.items()) + text = "\n".join(part for part in (_text(data.description), schema) if part) + else: + raise ProjectionMappingError("Unsupported Notion original kind") + title = title or "Untitled " + kind.replace("_", " ") + return ( + NotionOriginalEntity( + native_id=record.identity.native_id, + original_kind=kind, + title=title, + text=text, + native_url=native.url, + content_coverage=coverage, + breadcrumbs=[], + created_at=native.created_time, + updated_at=native.last_edited_time, + ), + ) diff --git a/backend/airweave/domains/entities/canonical/notion_property_projection.py b/backend/airweave/domains/entities/canonical/notion_property_projection.py new file mode 100644 index 0000000000..d545319beb --- /dev/null +++ b/backend/airweave/domains/entities/canonical/notion_property_projection.py @@ -0,0 +1,300 @@ +"""Offline rendering of complete, bounded native Notion property archives.""" + +from typing import Literal +from uuid import UUID + +from pydantic import ( + AwareDatetime, + BaseModel, + ConfigDict, + Field, + FiniteFloat, + JsonValue, + StrictBool, + StrictInt, + TypeAdapter, +) + +from airweave.domains.entities.canonical.blob_materializer import read_blob +from airweave.domains.entities.canonical.extraction_models import ExtractionPart +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_inputs import ProjectionInput, ProjectionInputs +from airweave.domains.entities.canonical.projection_mappers import ProjectionMappingError +from airweave.domains.storage.protocols import StorageBackend +from airweave.platform.entities.notion import NotionOriginalEntity + +MAX_PROPERTY_PROJECTION_BYTES = 16 * 1024 * 1024 + + +class _Header(BaseModel): + model_config = ConfigDict(extra="forbid") + format_version: Literal[1] + page_id: UUID + property_id: str = Field(min_length=1) + notion_version: Literal["2026-03-11"] + page_last_edited_time: AwareDatetime + + +class _Archive(_Header): + responses: list[dict[str, JsonValue]] = Field(min_length=1) + + +class _Manifest(_Header): + name: str + property: dict[str, JsonValue] + response_count: int = Field(ge=1) + value_status: Literal["available", "unsupported"] + + +class _Native(BaseModel): + model_config = ConfigDict(extra="ignore") + + +class _Value(_Native): + id: str + type: str + + +class _Item(_Value): + object: Literal["property_item"] + + +class _Status(_Native): + type: Literal["complete"] + + +class _Page(_Native): + object: Literal["list"] + type: Literal["property_item"] + results: list[dict[str, JsonValue]] + property_item: dict[str, JsonValue] + has_more: StrictBool + next_cursor: str | None + request_status: _Status | None = None + + +class _Label(_Native): + name: str + + +class _Reference(_Native): + id: UUID + name: str | None = None + + +class _Date(_Native): + start: str + end: str | None = None + time_zone: str | None = None + + +class _Rich(_Native): + type: Literal["text", "mention", "equation"] + plain_text: str + + +class _Unique(_Native): + number: StrictInt | None + prefix: str | None + + +class _Unsupported(Exception): + """Known native values whose searchable interpretation is not implemented.""" + + +_LIST_TYPES = {"title", "rich_text", "people", "relation", "rollup"} +_SUPPORTED = _LIST_TYPES | { + "select", + "status", + "multi_select", + "files", + "created_by", + "last_edited_by", + "date", + "checkbox", + "boolean", + "number", + "url", + "email", + "phone_number", + "string", + "created_time", + "last_edited_time", + "unique_id", + "formula", +} + + +def _known(kind: str) -> None: + if kind not in _SUPPORTED: + if kind in {"verification", "button", "place"}: + raise _Unsupported + raise ProjectionMappingError("Unknown Notion property value type") + + +def _render(kind: str, value: JsonValue) -> str: # noqa: C901 + _known(kind) + if value is None: + if kind in {"number", "date", "select", "status", "url", "email", "phone_number", "string"}: + return "" + raise ProjectionMappingError("Notion property has an invalid null value") + if kind in {"title", "rich_text"}: + return _Rich.model_validate(value).plain_text + if kind in {"select", "status"}: + return _Label.model_validate(value).name + if kind in {"multi_select", "files"}: + return ", ".join(item.name for item in TypeAdapter(list[_Label]).validate_python(value)) + if kind in {"people", "relation", "created_by", "last_edited_by"}: + ref = _Reference.model_validate(value) + return f"{ref.name} ({ref.id})" if ref.name else str(ref.id) + if kind == "date": + date = _Date.model_validate(value) + return ( + date.start + + (f" – {date.end}" if date.end else "") + + (f" ({date.time_zone})" if date.time_zone else "") + ) + if kind in {"checkbox", "boolean"}: + return "true" if TypeAdapter(StrictBool).validate_python(value) else "false" + if kind == "number": + return str(TypeAdapter(StrictInt | FiniteFloat).validate_python(value, strict=True)) + if kind in {"url", "email", "phone_number", "string", "created_time", "last_edited_time"}: + return TypeAdapter(str).validate_python(value, strict=True) + if kind == "unique_id": + unique = _Unique.model_validate(value) + return (f"{unique.prefix}-" if unique.prefix else "") + ( + str(unique.number) if unique.number is not None else "" + ) + if kind in {"formula", "rollup"}: + result = TypeAdapter(dict[str, JsonValue]).validate_python(value) + subtype = TypeAdapter(str).validate_python(result.get("type")) + if subtype == "incomplete": + raise ProjectionMappingError("Notion property calculation is incomplete") + if subtype == "unsupported": + raise _Unsupported + if subtype not in {"string", "number", "boolean", "date"}: + raise _Unsupported + if subtype not in result: + raise ProjectionMappingError("Notion calculation result is missing") + return _render(subtype, result[subtype]) + if kind in {"verification", "button", "place"}: + raise _Unsupported + raise ProjectionMappingError("Unknown Notion property value type") + + +def _pagination(page: _Page, last: bool, cursors: set[str]) -> None: + if page.has_more == last or (page.next_cursor is None) != last: + raise ProjectionMappingError("Notion property archive pagination is incomplete") + if page.next_cursor is not None: + if not page.next_cursor or page.next_cursor in cursors: + raise ProjectionMappingError("Notion property archive repeats a cursor") + cursors.add(page.next_cursor) + if last and page.property_item.get("next_url") is not None: + raise ProjectionMappingError("Notion terminal property retains continuation") + + +def _values( + archive: _Archive, prop: _Value +) -> tuple[list[dict[str, JsonValue]], dict[str, JsonValue]]: + values: list[dict[str, JsonValue]] = [] + metadata = {} + cursors: set[str] = set() + for index, raw in enumerate(archive.responses): + last = index == len(archive.responses) - 1 + if raw.get("object") == "property_item": + if len(archive.responses) != 1 or prop.type in _LIST_TYPES: + raise ProjectionMappingError("Notion single property has multiple responses") + item = _Item.model_validate(raw) + metadata = raw + values.append(raw) + else: + if prop.type not in _LIST_TYPES: + raise ProjectionMappingError("Notion scalar property returned a list") + page = _Page.model_validate(raw) + metadata = page.property_item + item = _Value.model_validate(metadata) + _pagination(page, last, cursors) + for raw_item in page.results: + parsed = _Item.model_validate(raw_item) + if prop.type != "rollup" and (parsed.id != prop.id or parsed.type != prop.type): + raise ProjectionMappingError("Notion property contains another identity") + values.extend(page.results) + if item.id != prop.id or item.type != prop.type: + raise ProjectionMappingError("Notion property metadata identity mismatch") + return values, metadata + + +def _text(archive: _Archive, prop: _Value, status: str) -> str: + _known(prop.type) + values, metadata = _values(archive, prop) + if prop.type in {"formula", "rollup"}: + calculation = TypeAdapter(dict[str, JsonValue]).validate_python(metadata.get(prop.type)) + subtype = calculation.get("type") + if (subtype == "unsupported") != (status == "unsupported"): + raise ProjectionMappingError("Notion property calculation status mismatch") + if prop.type == "rollup" and subtype == "array": + # show_original returns target values in results, not the metadata array. + return "\n".join( + _render(_Item.model_validate(v).type, v[_Item.model_validate(v).type]) + for v in values + ) + return _render(prop.type, calculation) + if status != "available": + raise ProjectionMappingError("Notion non-calculation property has invalid status") + rendered = [_render(prop.type, value[prop.type]) for value in values] + return ("" if prop.type in {"title", "rich_text"} else "\n").join(rendered) + + +async def map_notion_property(record: SourceRecord, storage: StorageBackend) -> ProjectionInputs: + """Validate retained authority and render values without fetching referenced objects.""" + manifest = _Manifest.model_validate(record.payload) + prop = _Value.model_validate(manifest.property) + parent = record.parent + if ( + record.payload_schema_version != 2 + or parent is None + or parent.record_type != "page" + or parent.container_id is not None + or parent.native_id != str(manifest.page_id) + or record.identity.container_id != str(manifest.page_id) + or record.identity.native_id != manifest.property_id + or prop.id != manifest.property_id + or len(record.blobs) != 1 + ): + raise ProjectionMappingError("Notion property archive authority mismatch") + blob = record.blobs[0] + if ( + blob.media_type != "application/json" + or blob.role is not None + or blob.source_path is not None + ): + raise ProjectionMappingError("Notion property archive descriptor mismatch") + if blob.size_bytes > MAX_PROPERTY_PROJECTION_BYTES: + raise ProjectionMappingError("Notion property exceeds the 16 MiB projection budget") + archive = _Archive.model_validate_json(await read_blob(record, blob, storage)) + header = tuple(_Header.model_fields) + if ( + archive.model_dump(include=set(header)) != manifest.model_dump(include=set(header)) + or len(archive.responses) != manifest.response_count + ): + raise ProjectionMappingError("Notion property archive provenance mismatch") + part = ExtractionPart(part_index=0, key=prop.id, kind="record", media_type="application/json") + try: + text = _text(archive, prop, manifest.value_status) + except _Unsupported: + return ProjectionInputs( + parts=(ProjectionInput(part=part, entity=None, omission="unsupported_format"),) + ) + entity = NotionOriginalEntity( + native_id=prop.id, + original_kind="page_property", + title=manifest.name, + text=text, + breadcrumbs=[], + native_url=None, + created_at=None, + updated_at=None, + content_coverage="This property only. References and file names exclude linked bodies and " + "file bytes. Related values are not a transactional snapshot.", + ) + return ProjectionInputs(parts=(ProjectionInput(part=part, entity=entity),)) diff --git a/backend/airweave/domains/entities/canonical/outlook_calendar_projection.py b/backend/airweave/domains/entities/canonical/outlook_calendar_projection.py new file mode 100644 index 0000000000..30b84b5597 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/outlook_calendar_projection.py @@ -0,0 +1,225 @@ +"""Offline projection of retained Graph calendar metadata and event bodies.""" + +from pathlib import Path + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.domains.entities.canonical.blob_materializer import write_blob +from airweave.domains.entities.canonical.extraction_models import ExtractionPart +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_inputs import ProjectionInput, ProjectionInputs +from airweave.domains.sync_pipeline.pipeline.text_models import NativeTextBody +from airweave.platform.entities.outlook_calendar import ( + OutlookCalendarHtmlEventEntity, + OutlookCalendarMetadataEntity, + OutlookCalendarSearchEventEntity, +) +from airweave.platform.sources.outlook_calendar_models import OutlookCalendar, OutlookCalendarEvent +from airweave.platform.sources.outlook_mail_models import OutlookAddress, OutlookRecipient + + +class _Response(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + response: str | None = None + + +class _Attendee(OutlookRecipient): + type: str | None = None + status: _Response | None = None + + +class _Pattern(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + type: str + interval: int | None = None + month: int | None = None + dayOfMonth: int | None = None + daysOfWeek: list[str] = Field(default_factory=list) + firstDayOfWeek: str | None = None + index: str | None = None + + +class _Range(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + type: str + startDate: str + endDate: str | None = None + recurrenceTimeZone: str | None = None + numberOfOccurrences: int | None = None + + +class _Recurrence(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + pattern: _Pattern + range: _Range + + +class _Location(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + displayName: str | None = None + + +class _Details(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + location: _Location | None = None + locations: list[_Location] = Field(default_factory=list) + categories: list[str] = Field(default_factory=list) + showAs: str | None = None + + +def _address(value: OutlookAddress) -> str | None: + return ( + " <".join([value.name, value.address]) + ">" + if value.name and value.address + else (value.name or value.address) + ) + + +def _attendee(value: _Attendee) -> str: + fields = [_address(value.emailAddress)] + if value.type is not None: + fields.append(f"type: {value.type}") + if value.status is not None and value.status.response is not None: + fields.append(f"response: {value.status.response}") + return "; ".join(field for field in fields if field) + + +def _recurrence(value: _Recurrence) -> str: + # Only declared native recurrence fields enter search, never extension JSON. + sections = [] + for label, fields in (("Pattern", value.pattern), ("Range", value.range)): + pairs = [] + for key, item in fields.model_dump(exclude_none=True).items(): + if item == []: + continue + rendered = ", ".join(item) if isinstance(item, list) else str(item) + pairs.append(f"{key}: {rendered}") + sections.append(f"{label}: " + "; ".join(pairs)) + return "\n".join(sections) + + +async def map_outlook_calendar(record: SourceRecord, directory: Path) -> ProjectionInputs: + """Use canonical originals only; conversion and retained/indexed text stay shared.""" + if ( + record.deleted_at is not None + or record.content_access != "available" + or record.payload_schema_version != 1 + or record.blobs + ): + raise ValueError("Outlook calendar projection requires active schema-one JSON originals") + if record.identity.record_type == "calendar": + if record.parent is not None or record.identity.container_id is not None: + raise ValueError("Outlook calendar must be a root identity") + native = OutlookCalendar.model_validate(record.payload) + if native.id != record.identity.native_id: + raise ValueError("Outlook calendar native identity mismatch") + entity = OutlookCalendarMetadataEntity( + native_id=native.id, + title=native.name or "Calendar", + breadcrumbs=[], + owner=_address(OutlookAddress.model_validate(native.owner)) if native.owner else None, + can_edit=native.canEdit, + can_view_private_items=native.canViewPrivateItems, + is_default_calendar=native.isDefaultCalendar, + ) + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart(part_index=0, key="/calendar", kind="record"), + entity=entity, + ), + ) + ) + parent = record.parent + if ( + record.identity.record_type != "event" + or parent is None + or parent.record_type != "calendar" + or parent.container_id is not None + or record.identity.container_id != parent.native_id + or record.completeness != "partial" + ): + raise ValueError("Outlook event requires its exact calendar parent and partial capture") + event = OutlookCalendarEvent.model_validate(record.payload) + if event.id != record.identity.native_id: + raise ValueError("Outlook event native identity mismatch") + details = _Details.model_validate(record.payload) + locations = ([details.location] if details.location else []) + details.locations + metadata = OutlookCalendarSearchEventEntity( + native_id=event.id, + title=event.subject or "Calendar event", + calendar_id=parent.native_id, + breadcrumbs=[], + created_at=record.source_created_at, + updated_at=record.source_updated_at, + event_type=event.type, + start=f"{event.start.dateTime} ({event.start.timeZone})", + end=f"{event.end.dateTime} ({event.end.timeZone})", + is_all_day=event.isAllDay, + is_cancelled=event.isCancelled, + organizer=_address(event.organizer.emailAddress) if event.organizer else None, + attendees=[_attendee(_Attendee.model_validate(value)) for value in event.attendees], + recurrence=_recurrence(_Recurrence.model_validate(event.recurrence)) + if event.recurrence is not None + else None, + series_master_id=event.seriesMasterId, + original_start=event.originalStart, + original_start_timezone=event.originalStartTimeZone, + original_end_timezone=event.originalEndTimeZone, + locations=list( + dict.fromkeys(location.displayName for location in locations if location.displayName) + ), + categories=details.categories, + show_as=details.showAs, + web_link=event.webLink, + ) + body = event.body + if body is not None and body.contentType == "html": + content = body.content.encode("utf-8") + local = await write_blob(content, directory, suffix=".html") + html_entity = OutlookCalendarHtmlEventEntity( + **metadata.model_dump(round_trip=True), + url=event.webLink or "https://outlook.office.com/calendar/", + size=len(content), + file_type="html", + mime_type="text/html", + local_path=str(local), + ) + parts = [ + ProjectionInput( + part=ExtractionPart( + part_index=0, + key="/body/content", + kind="body", + media_type="text/html", + extension=".html", + ), + entity=html_entity, + ) + ] + elif body is not None: + parts = [ + ProjectionInput( + part=ExtractionPart( + part_index=0, key="/body/content", kind="body", media_type="text/plain" + ), + entity=metadata, + native_body=NativeTextBody(text=body.content), + ) + ] + else: + parts = [ + ProjectionInput( + part=ExtractionPart(part_index=0, key="/event", kind="record"), entity=metadata + ), + ProjectionInput( + part=ExtractionPart(part_index=1, key="/body/content", kind="body"), entity=None + ), + ] + for key in ("/attachment_inventory", "/cancelledOccurrences_inventory"): + parts.append( + ProjectionInput( + part=ExtractionPart(part_index=len(parts), key=key, kind="record"), entity=None + ) + ) + return ProjectionInputs(parts=tuple(parts)) diff --git a/backend/airweave/domains/entities/canonical/outlook_projection.py b/backend/airweave/domains/entities/canonical/outlook_projection.py new file mode 100644 index 0000000000..ee2b3920a8 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/outlook_projection.py @@ -0,0 +1,318 @@ +"""Offline Outlook projection from native metadata and retained, verified RFC822 bytes.""" + +import hashlib +import html +import mimetypes +import re +from datetime import datetime +from email import policy +from email.errors import MessageError +from email.message import EmailMessage +from email.parser import BytesParser +from email.utils import quote +from pathlib import Path + +from airweave.domains.entities.canonical.blob_materializer import read_blob, write_blob +from airweave.domains.entities.canonical.extraction_models import CharsetRecovery, ExtractionPart +from airweave.domains.entities.canonical.mime_text import decode_mime_text +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_inputs import ProjectionInput, ProjectionInputs +from airweave.domains.storage.protocols import StorageBackend +from airweave.platform.entities._base import Breadcrumb +from airweave.platform.entities.outlook_mail import OutlookAttachmentEntity, OutlookMessageEntity +from airweave.platform.sources.outlook_mail_models import OutlookMessage, OutlookRecipient + +# Fail rather than publish a truncated body for pathological MIME trees. +MAX_MIME_DEPTH = 32 +MAX_MIME_PARTS = 1024 +_UNSUPPORTED = {"multipart/encrypted", "application/pkcs7-mime", "application/x-pkcs7-mime"} + + +def _address(recipient: OutlookRecipient | None) -> str | None: + if recipient is None or not recipient.emailAddress.address: + return None + address = recipient.emailAddress + # Search metadata is readable text, not a wire header requiring RFC2047 encoding. + return f'"{quote(address.name)}" <{address.address}>' if address.name else address.address + + +def _addresses(recipients: list[OutlookRecipient]) -> list[str]: + return [address for recipient in recipients if (address := _address(recipient))] + + +def _date(value: str | None) -> datetime | None: + if value is None: + return None + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + raise ValueError("Outlook metadata timestamp requires a timezone") + return parsed + + +def _decoded(part: EmailMessage) -> bytes: + encoding = str(part.get("Content-Transfer-Encoding", "7bit")).strip().lower() + if encoding not in {"7bit", "8bit", "binary", "base64", "quoted-printable"}: + raise ValueError("Unsupported MIME transfer encoding") + if encoding == "quoted-printable": + # The stdlib deliberately accepts invalid escapes. Preserve no silently repaired text. + raw = part.get_payload(decode=False) + if not isinstance(raw, str) or re.search(r"=(?![0-9A-Fa-f]{2}|\r?\n)", raw): + raise ValueError("Malformed quoted-printable MIME payload") + content = part.get_payload(decode=True) + if not isinstance(content, bytes): + raise ValueError("MIME leaf must contain bytes") + if part.defects: + raise ValueError("Malformed MIME transfer payload") + return content + + +def _parse(content: bytes) -> EmailMessage: + try: + message = BytesParser(policy=policy.default.clone(raise_on_defect=True)).parsebytes(content) + pending = [(message, 0)] + count = 0 + while pending: + part, depth = pending.pop() + count += 1 + if depth > MAX_MIME_DEPTH or count > MAX_MIME_PARTS: + raise ValueError("Outlook MIME tree exceeds projection depth or part limit") + if part.defects or any(header.defects for header in part.values()): + raise ValueError("Malformed Outlook MIME headers") + if part.is_multipart(): + pending.extend((child, depth + 1) for child in part.iter_parts()) + else: + _decoded(part) + return message + except (MessageError, RecursionError) as error: + raise ValueError("Malformed Outlook MIME content") from error + + +def _file(part: EmailMessage) -> bool: + return ( + bool(part.get_filename()) + or part.get_content_disposition() == "attachment" + or (not part.is_multipart() and part.get_content_type() not in {"text/plain", "text/html"}) + ) + + +def _suffix(part: EmailMessage) -> str: + # MIME type is authoritative over misleading names (e.g. a PDF called image.png). + mime = part.get_content_type() + mime_suffix = mimetypes.guess_extension(mime) if mime != "application/octet-stream" else None + if mime_suffix and re.fullmatch(r"\.[a-z0-9]{1,12}", mime_suffix): + return mime_suffix + filename_suffix = Path(part.get_filename() or "").suffix.lower() + return filename_suffix if re.fullmatch(r"\.[a-z0-9]{1,12}", filename_suffix) else ".bin" + + +class OutlookProjection: + """Disposable extraction inputs; canonical originals remain unchanged.""" + + def __init__(self, record: SourceRecord, directory: Path, native: OutlookMessage): + """Bind validated native metadata and caller-owned scratch directory.""" + self.record = record + self.directory = directory + self.native = native + self.parts: list[ProjectionInput] = [] + self.fragments: list[str] = [] + self.charset_recoveries: list[CharsetRecovery] = [] + + def unsupported(self, part: EmailMessage, path: str) -> None: + """Account for retained content whose semantics are not yet extractable.""" + self.parts.append( + ProjectionInput( + part=ExtractionPart( + part_index=len(self.parts) + 1, + key=path, + kind="file", + media_type=part.get_content_type(), + ), + entity=None, + omission="unsupported_format", + ) + ) + + async def attachment(self, part: EmailMessage, path: str) -> None: + """Materialize safe content-hash paths, independent of provider filenames.""" + content = _decoded(part) + suffix = _suffix(part) + local = await write_blob(content, self.directory, suffix=suffix) + message_id = self.native.id + self.parts.append( + ProjectionInput( + part=ExtractionPart( + part_index=len(self.parts) + 1, + key=path, + kind="file", + media_type=part.get_content_type(), + extension=suffix, + ), + entity=OutlookAttachmentEntity( + composite_id=f"{message_id}:{hashlib.sha256(path.encode()).hexdigest()}", + name=part.get_filename() or "Inline content", + message_id=message_id, + attachment_id=path, + content_type=part.get_content_type(), + is_inline=part.get_content_disposition() == "inline" + or bool(part.get("Content-ID")), + content_id=str(part["Content-ID"]) if part["Content-ID"] else None, + metadata={"identity_kind": "mime_path"}, + message_web_url=self.native.webLink, + url=self.native.webLink or "https://outlook.office.com/mail/", + size=len(content), + file_type=suffix[1:], + mime_type=part.get_content_type(), + local_path=str(local), + breadcrumbs=[ + Breadcrumb( + entity_id=message_id, + name=self.native.subject or "Email message", + entity_type="OutlookMessageEntity", + ) + ], + ), + ) + ) + + async def render(self, part: EmailMessage, path: str, *, body: bool = True) -> None: + """Select one alternative while still retaining independent inline/file parts.""" + mime = part.get_content_type() + if mime.startswith("message/") or mime in _UNSUPPORTED: + self.unsupported(part, path) + return + if part.is_multipart(): + if part.get_content_disposition() == "attachment": + self.unsupported(part, path) + return + children = list(part.iter_parts()) + selected = set(range(len(children))) + if mime == "multipart/alternative": + # get_body understands nested related parts and their explicit start CID. + preferred = part.get_body(preferencelist=("html", "plain")) + selected = { + i + for i, child in enumerate(children) + if preferred is not None and any(node is preferred for node in child.walk()) + } + elif mime == "multipart/related": + start = part.get_param("start") + selected = { + next( + ( + i + for i, child in enumerate(children) + if start and child.get("Content-ID") == start + ), + 0, + ) + } + for index, child in enumerate(children): + await self.render(child, f"{path}/parts/{index}", body=body and index in selected) + return + if _file(part) or (not body and part.get("Content-ID")): + await self.attachment(part, path) + elif body: + decoded = decode_mime_text( + _decoded(part), part.get_content_charset() or "us-ascii", path + ) + if decoded.recovery is not None: + self.charset_recoveries.append(decoded.recovery) + text = decoded.text + self.fragments.append( + text if mime == "text/html" else "
" + html.escape(text) + "" + ) + + async def body(self) -> ProjectionInput: + """Expose full decoded text and independently retained native message metadata.""" + content = "\n".join(self.fragments).encode("utf-8") + local = await write_blob(content, self.directory, suffix=".html") + native = self.native + return ProjectionInput( + part=ExtractionPart( + part_index=0, + key="/mime/body", + kind="body", + media_type="text/html", + extension=".html", + charset_recoveries=tuple(self.charset_recoveries), + ), + entity=OutlookMessageEntity( + id=native.id, + breadcrumbs=[], + folder_id=native.parentFolderId, + subject=native.subject.strip() + if native.subject and native.subject.strip() + else "Email message", + sender=_address(native.sender), + from_address=_address(native.from_), + conversation_id=native.conversationId, + to_recipients=_addresses(native.toRecipients), + cc_recipients=_addresses(native.ccRecipients), + bcc_recipients=_addresses(native.bccRecipients), + reply_to=_addresses(native.replyTo), + sent_date=_date(native.sentDateTime), + received_date=_date(native.receivedDateTime), + internet_message_id=native.internetMessageId, + is_draft=native.isDraft, + has_attachments=native.hasAttachments, + web_url_override=native.webLink, + url=native.webLink or "https://outlook.office.com/mail/", + size=len(content), + file_type="html", + mime_type="text/html", + local_path=str(local), + ), + ) + + +async def map_outlook( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> ProjectionInputs: + """Project schema-one mailbox messages without provider calls or body-preview fallback.""" + if ( + record.deleted_at is not None + or record.content_access != "available" + or record.payload_schema_version != 1 + or record.identity.record_type != "message" + or record.identity.container_id is not None + or record.parent is not None + or record.completeness not in {"partial", "metadata_only"} + ): + raise ValueError("Outlook projection requires an active schema-one mailbox message") + native = OutlookMessage.model_validate(record.payload) + if native.id != record.identity.native_id: + raise ValueError("Outlook projection identity differs from retained native metadata") + if len(record.blobs) > 1 or any( + ref.role is not None or ref.source_path is not None or ref.media_type != "message/rfc822" + for ref in record.blobs + ): + raise ValueError("Outlook projection requires one top-level RFC822 original") + if not record.blobs: + if record.completeness != "metadata_only": + raise ValueError("Missing Outlook MIME requires an explicit metadata-only capture") + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, key="/mime", kind="body", media_type="message/rfc822" + ), + entity=None, + ), + ) + ) + if record.completeness != "partial": + raise ValueError("Outlook MIME capture must report pending attachment inventory") + message = _parse(await read_blob(record, record.blobs[0], storage)) + projection = OutlookProjection(record, directory, native) + await projection.render(message, "/mime") + body = await projection.body() + if not projection.fragments and any(part.omission for part in projection.parts): + body = ProjectionInput(part=body.part, entity=None, omission="unsupported_format") + # MIME cannot prove native reference/item attachments were represented in the response. + inventory = ProjectionInput( + part=ExtractionPart( + part_index=len(projection.parts) + 1, key="/attachment_inventory", kind="record" + ), + entity=None, + ) + return ProjectionInputs(parts=(body, *projection.parts, inventory)) diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py new file mode 100644 index 0000000000..efe1d3cd68 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -0,0 +1,162 @@ +"""Opted-in scope pages; the existing capture pipeline owns durable progress.""" + +from typing import Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from airweave.domains.entities.canonical.cycle_models import ( + CaptureCycle, + CaptureMode, + CycleConfiguration, + ProviderCheckpoint, + SourcePlan, +) +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + ScopeRemovalReason, +) +from airweave.domains.entities.canonical.scan_models import ( + ChildScopeObservation, + ScanContinuation, + ScanState, +) +from airweave.domains.entities.canonical.scope_execution import ScopePlan +from airweave.domains.storage.file_service import FileService + + +class CapturePage(BaseModel): + """One provider response and the continuation committed with its original records.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + records: tuple[CaptureRecord, ...] = Field(max_length=500) + discovered_records: tuple[CaptureRecord, ...] = Field( + default=(), + max_length=500, + description="Verified independent originals from the bound source; never placeholders", + ) + child_scope_observations: tuple[ChildScopeObservation, ...] = Field( + default=(), + max_length=500, + description="Complete child inventories evidenced by active originals on this page", + ) + continuation: ScanContinuation + final: bool = False + provider_checkpoint: ProviderCheckpoint | None = None + + +class CapturePlan(BaseModel): + """Immutable source plan selected once before engine cycle creation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + mode: CaptureMode = "full" + starting_checkpoint: ProviderCheckpoint | None = None + source_plan: SourcePlan = Field(default_factory=dict) + + +class InvalidCaptureCheckpoint(Exception): + """Native evidence invalidates a whole cycle boundary, not merely a page token.""" + + +@runtime_checkable +class CheckpointedPageSource(Protocol): + """Sources with a native changes API reuse the same cycle and page authority.""" + + async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: + """Provider I/O occurs before the fenced cycle CAS; None requests a fresh full pass.""" + ... + + def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation: + """Derive resumable page state only from the persisted immutable plan.""" + ... + + +class InvalidScanContinuation(Exception): + """The provider rejected a saved cursor; restart the whole scope, never reconcile it.""" + + +class ScopeAccessLost(Exception): + """The source's audited contract says this child container is no longer available.""" + + def __init__( + self, + message: str, + *, + removal_reason: ScopeRemovalReason = "access_revoked", + ): + """Slack keeps its confirmed-access default; other sources must choose explicitly.""" + super().__init__(message) + self.removal_reason = TypeAdapter(ScopeRemovalReason).validate_python(removal_reason) + + +class RequiredScopeAccessLost(ScopeAccessLost): + """Withdraw known unavailable content, then fail an explicitly required source selection.""" + + +@runtime_checkable +class CanonicalPageSource(Protocol): + """A source fetches pages; SQL, attempts, scope discovery and checkpointing stay outside it.""" + + canonical_record_types: tuple[str, ...] + canonical_container_parents: dict[str, str | tuple[str | None, ...]] + capture_cycle_configuration: CycleConfiguration + + async def capture_page( + self, + scope: CompletedScope, + continuation: ScanContinuation, + *, + files: FileService, + parent: SourceRecord | None = None, + ) -> CapturePage: + """Fetch one page without advancing durable state or performing writes.""" + ... + + def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: + """Source-owned stable container policy; always retain the full parent identity.""" + ... + + async def confirm_absent(self, record: SourceRecord) -> None: + """Raise unless an omitted inventory record is confirmed outside accessible scope.""" + ... + + +@runtime_checkable +class KnownObjectSource(Protocol): + """Sources declaring exact known-object validation implement current reads.""" + + async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord: + """Return exact fresh state or explicit unavailability; errors never mean absence.""" + ... + + +class InvalidScopeCheckpoint(Exception): + """Native evidence requires a full restart of one scope, not the whole forest.""" + + +@runtime_checkable +class ScopedPageSource(Protocol): + """Mixed sources select immutable per-scope requests; the store attests authority.""" + + async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: + """Freeze source parameters once before cycle creation.""" + ... + + async def prepare_scope( + self, + scope: CompletedScope, + cycle: CaptureCycle, + previous: ScanState | None, + *, + parent: SourceRecord | None, + force_full: bool, + ) -> ScopePlan: + """Called outside SQL only when starting or explicitly restarting a sweep.""" + ... + + def initial_scope_continuation( + self, scope: CompletedScope, cycle: CaptureCycle, plan: ScopePlan + ) -> ScanContinuation: + """Derive progress from immutable persisted scope and cycle parameters.""" + ... diff --git a/backend/airweave/domains/entities/canonical/projection_gc.py b/backend/airweave/domains/entities/canonical/projection_gc.py new file mode 100644 index 0000000000..11edcb2c14 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_gc.py @@ -0,0 +1,150 @@ +"""Bounded recurring reclamation; retirement fences publication before any remote delete.""" + +from datetime import datetime, timedelta +from uuid import UUID, uuid4 + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.db.unit_of_work import UnitOfWork +from airweave.domains.entities.canonical.projection_models import ( + ProjectionCleanupPage, + ProjectionDocument, +) +from airweave.domains.entities.canonical.store import content_is_available +from airweave.domains.entities.canonical.text_models import TextArtifact +from airweave.models.entity import Entity +from airweave.models.projection_generation import ProjectionGeneration +from airweave.models.sync import Sync + + +class ProjectionGCStore: + """Retain manifests indefinitely: timed-out feed threads have no bounded lifetime.""" + + async def due(self, db: AsyncSession, *, now: datetime, limit: int = 20) -> tuple[UUID, ...]: + """Bound maintenance work independently from source sync frequency.""" + return tuple( + await db.scalars( + select(ProjectionGeneration.id) + .where( + ProjectionGeneration.next_gc_at <= now, + ) + .order_by(ProjectionGeneration.next_gc_at, ProjectionGeneration.id) + .limit(limit) + ) + ) + + async def claim( + self, + db: AsyncSession, + generation: UUID, + *, + now: datetime, + limit: int = 100, + ) -> ProjectionCleanupPage | None: + """Irrevocably retire under the same Sync lock used by publication.""" + if not 1 <= limit <= 100: + raise ValueError("Cleanup document budget must be between 1 and 100") + async with UnitOfWork(db): + initial = await db.get(ProjectionGeneration, generation) + if initial is None: + return None + await db.scalar(select(Sync.id).where(Sync.id == initial.sync_id).with_for_update()) + row = await db.scalar( + select(ProjectionGeneration) + .where( + ProjectionGeneration.id == generation, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + if row.next_gc_at > now: + return None + current = await db.scalar( + select(Entity).where( + Entity.id == row.record_id, + Entity.sync_id == row.sync_id, + Entity.organization_id == row.organization_id, + ) + ) + if current is not None and current.indexed_generation == generation: + available = await db.scalar( + select(content_is_available()).where(Entity.id == current.id) + ) + if not row.documents or (current.deleted_at is None and available): + row.next_gc_at = now + timedelta(days=1) + return None + # Capture and parent transitions share the Sync lock. Clear publication + # before retiring: restoration then requires a fresh generation, never + # resurrects documents that an already-issued delete may still remove. + current.indexed_generation = None + current.indexed_revision = None + current.indexed_pipeline_version = None + current.indexed_chunk_count = None + if row.retired_at is None: + row.retired_at = now + row.next_gc_at = now + timedelta(minutes=5) + row.gc_attempt = uuid4() + await db.flush() + artifacts = tuple( + TextArtifact.model_validate(item) for item in (row.text_representations or []) + ) + start, end = row.delete_cursor, row.delete_cursor + limit + artifact_start = max(0, start - len(row.documents)) + artifact_end = max(0, end - len(row.documents)) + return ProjectionCleanupPage( + generation=generation, + attempt=row.gc_attempt, + cursor=row.delete_cursor, + artifact_keys=tuple( + item.storage_key(row.sync_id, row.id) + for item in artifacts[artifact_start:artifact_end] + ), + documents=tuple( + ProjectionDocument.model_validate(item) + for item in row.documents[row.delete_cursor : row.delete_cursor + limit] + ), + ) + + async def acknowledge( + self, + db: AsyncSession, + page: ProjectionCleanupPage, + *, + now: datetime, + error: str | None = None, + ) -> None: + """Retry failed pages; successful full passes remain scheduled for late writes.""" + async with UnitOfWork(db): + row = await db.scalar( + select(ProjectionGeneration) + .where( + ProjectionGeneration.id == page.generation, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + if ( + row is None + or row.retired_at is None + or row.delete_cursor != page.cursor + or row.gc_attempt != page.attempt + ): + return + row.gc_attempt = None + if error: + row.gc_error = error[:200] + row.next_gc_at = now + timedelta(minutes=5) + return + row.gc_error = None + row.delete_cursor += len(page.documents) + len(page.artifact_keys) + if row.delete_cursor >= len(row.documents) + len(row.text_representations or []): + row.delete_cursor = 0 + row.gc_passes += 1 + row.last_gc_at = now + row.next_gc_at = now + timedelta( + minutes=min(1440, 15 * 2 ** min(row.gc_passes - 1, 7)) + ) + else: + row.next_gc_at = now + await db.flush() diff --git a/backend/airweave/domains/entities/canonical/projection_inputs.py b/backend/airweave/domains/entities/canonical/projection_inputs.py new file mode 100644 index 0000000000..791ce7fd14 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_inputs.py @@ -0,0 +1,38 @@ +"""Typed mapper outputs; source descriptors survive filtering and entity-ID stamping.""" + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, model_validator + +from airweave.domains.entities.canonical.extraction_models import ExtractionPart +from airweave.domains.sync_pipeline.pipeline.text_models import NativeTextBody +from airweave.platform.entities._base import BaseEntity + + +class ProjectionInput(BaseModel): + """An expected part, with None only for explicitly uncaptured original bytes.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + part: ExtractionPart + entity: BaseEntity | None + native_body: NativeTextBody | None = None + omission: Literal["unsupported_format"] | None = None + + @model_validator(mode="after") + def valid_omission(self) -> "ProjectionInput": + """Explicit unsupported originals cannot also supply indexable content.""" + if self.omission is not None and (self.entity is not None or self.native_body is not None): + raise ValueError("Unsupported projection parts cannot supply an entity or text") + return self + + +class ProjectionInputs(BaseModel): + """All expected parts, including deliberate original-capture omissions.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + parts: tuple[ProjectionInput, ...] + + @property + def entities(self) -> tuple[BaseEntity, ...]: + """Materialized entities for mapper callers; coverage lives in parts.""" + return tuple(p.entity for p in self.parts if p.entity is not None) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py new file mode 100644 index 0000000000..ce774293ae --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -0,0 +1,534 @@ +"""Typed search views of captured records. This module never contacts a provider.""" + +from __future__ import annotations + +import mimetypes +import re +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from pathlib import Path +from tempfile import TemporaryDirectory + +from pydantic import JsonValue, TypeAdapter + +from airweave.domains.entities.canonical.extraction_models import ExtractionPart +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_inputs import ProjectionInput, ProjectionInputs +from airweave.domains.entities.canonical.projection_policy import excluded_from_search +from airweave.domains.storage.protocols import StorageBackend +from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields +from airweave.platform.entities._base import BaseEntity, FileEntity +from airweave.platform.entities.google_calendar import ( + GoogleCalendarEventEntity, + GoogleCalendarListEntity, +) +from airweave.platform.entities.google_drive import GoogleDriveFileEntity, GoogleDriveFolderEntity +from airweave.platform.entities.slack import SlackChannelEntity, SlackMessageEntity +from airweave.platform.entities.wispr import WisprMeetingEntity, WisprNoteEntity +from airweave.platform.sources.records.sheets_manifest import GridGap + + +class ProjectionMappingError(ValueError): + """Captured representation cannot be faithfully mapped; keep projection pending.""" + + +def _string(value: JsonValue, *, default: str = "") -> str: + if value is None: + return default + if not isinstance(value, str): + raise ProjectionMappingError("Expected a captured string field") + return value + + +def _datetime(value: JsonValue) -> datetime | None: + if value is None or value == "": + return None + parsed = datetime.fromisoformat(_string(value).replace("Z", "+00:00")) + if parsed.tzinfo is None: + raise ProjectionMappingError("Captured timestamp lacks timezone") + return parsed + + +def _drive_omissions( + paths: tuple[str, ...], *, start_index: int = 1 +) -> tuple[ProjectionInput, ...]: + """Missing native representations remain explicit beside any usable text/export.""" + return tuple( + ProjectionInput( + part=ExtractionPart(part_index=index + start_index, key=path, kind="file"), + entity=None, + ) + for index, path in enumerate(dict.fromkeys(paths)) + ) + + +def _sheet_gap_key(gap: GridGap) -> str: + """Use native sheet IDs and exclusive bounds, never mutable sheet titles.""" + key = f"/native/sheets/{gap.sheet_id}" + bounds = gap.bounds + if bounds is not None: + key += ( + f"/rows/{bounds.start_row}:{bounds.end_row}" + f"/columns/{bounds.start_column}:{bounds.end_column}" + ) + return key + + +async def _drive_document( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> ProjectionInputs: + """Verify each retained Docs representation once and preserve missing-native evidence.""" + from airweave.domains.entities.canonical.blob_materializer import read_blob + from airweave.domains.entities.canonical.workspace_docs import CapturedDocument + from airweave.platform.sources.records.workspace_manifest import ( + parse_manifest, + validate_document, + ) + + if record.payload.get("mimeType") != "application/vnd.google-apps.document": + raise ProjectionMappingError("Captured file is not a Google document") + marked = [blob for blob in record.blobs if blob.role == "representation_manifest"] + if len(marked) != 1: + raise ProjectionMappingError("Drive document requires one representation manifest") + manifest = parse_manifest( + await read_blob(record, marked[0], storage), + file_id=record.identity.native_id, + drive_version=_string(record.payload.get("version")), + blobs=record.blobs, + ) + by_digest = {blob.sha256: blob for blob in record.blobs} + missing = tuple(gap.source_path or "/native" for gap in manifest.native.missing) + if manifest.native.status == "unavailable": + if manifest.export.status != "retained": + return ProjectionInputs(parts=_drive_omissions(missing, start_index=0)) + content = await read_blob(record, by_digest[manifest.export.blob], storage) + media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + suffix = ".docx" + else: + raw = await read_blob(record, by_digest[manifest.native.document_blob], storage) + document = TypeAdapter(dict[str, JsonValue]).validate_json(raw) + validate_document(document, file_id=record.identity.native_id) + content = CapturedDocument(manifest=manifest, document=document).text().encode("utf-8") + media_type, suffix = "text/plain", ".txt" + body = await _drive_file(record, content, media_type, suffix, directory) + return ProjectionInputs(parts=(body, *_drive_omissions(missing))) + + +async def _drive( + record: SourceRecord, + storage: StorageBackend, + directory: Path, +) -> ProjectionInputs: + from airweave.domains.entities.canonical.blob_materializer import read_blob + + if record.identity.record_type != "file": + raise ProjectionMappingError("Unsupported Drive record type") + data = record.payload + if data.get("mimeType") == "application/vnd.google-apps.folder": + return ProjectionInputs( + parts=( + _projection_input( + 0, + GoogleDriveFolderEntity( + folder_id=record.identity.native_id, + title=_string(data.get("name"), default="Untitled"), + description=_string(data.get("description")), + breadcrumbs=[], + ), + ), + ) + ) + if record.completeness == "metadata_only" and not record.blobs: + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, + key=record.identity.native_id, + kind="file", + media_type=_string(data.get("mimeType")) or None, + ), + entity=None, + ), + ) + ) + omissions: tuple[ProjectionInput, ...] = () + if any(blob.role == "representation_manifest" for blob in record.blobs): + if data.get("mimeType") == "application/vnd.google-apps.spreadsheet": + from airweave.domains.entities.canonical.workspace_sheets import read_spreadsheet + + spreadsheet = await read_spreadsheet(record, storage, for_projection=True) + omissions = _drive_omissions( + tuple(_sheet_gap_key(gap) for gap in spreadsheet.manifest.native.missing) + ) + if spreadsheet.manifest.native.status == "complete": + content = spreadsheet.text().encode("utf-8") + media_type, suffix = "text/plain", ".txt" + else: + digest = spreadsheet.manifest.export.blob + if digest is None: + raise ProjectionMappingError("Partial native grid has no complete export") + blob = next(blob for blob in record.blobs if blob.sha256 == digest) + content = await read_blob(record, blob, storage) + media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + suffix = ".xlsx" + else: + return await _drive_document(record, storage, directory) + else: + # Historical Drive records retain their export-only contract. + if len(record.blobs) != 1: + raise ProjectionMappingError("Drive content requires exactly one captured blob") + blob = record.blobs[0] + content = await read_blob(record, blob, storage) + media_type = blob.media_type or _string(data.get("mimeType")) + suffix = mimetypes.guess_extension(media_type) or Path(_string(data.get("name"))).suffix + if not suffix: + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, + key=record.identity.native_id, + kind="file", + media_type=media_type or None, + ), + entity=None, + omission="unsupported_format", + ), + ) + ) + suffix = suffix.lower() + body = await _drive_file(record, content, media_type, suffix, directory) + return ProjectionInputs(parts=(body, *omissions)) + + +async def _drive_file( + record: SourceRecord, content: bytes, media_type: str, suffix: str, directory: Path +) -> ProjectionInput: + """Materialize one verified representation without changing its captured original.""" + from airweave.domains.entities.canonical.blob_materializer import write_blob + + data = record.payload + local_path = await write_blob(content, directory, suffix=suffix) + metadata = { + **data, + "createdTime": data.get("createdTime") or record.observed_at.isoformat(), + "modifiedTime": data.get("modifiedTime") or record.observed_at.isoformat(), + } + entity = GoogleDriveFileEntity.from_api(metadata, breadcrumbs=[]) + entity.local_path = str(local_path) + entity.size = len(content) + entity.mime_type = media_type + entity.file_type = suffix.lstrip(".") + entity.url = entity.web_url + return _projection_input(0, entity) + + +def _calendar(record: SourceRecord) -> tuple[BaseEntity, ...]: + metadata = { + **record.payload, + "created": record.payload.get("created") or record.observed_at.isoformat(), + "updated": record.payload.get("updated") or record.observed_at.isoformat(), + } + if record.identity.record_type == "calendar": + return (GoogleCalendarListEntity.from_api(metadata),) + if record.identity.record_type != "event" or not record.identity.container_id: + raise ProjectionMappingError("Calendar event requires its native calendar identity") + return ( + GoogleCalendarEventEntity.from_api( + metadata, + calendar_key=record.identity.container_id, + breadcrumbs=[], + ), + ) + + +def _slack(record: SourceRecord) -> tuple[BaseEntity, ...]: + payload = record.payload + if record.identity.record_type == "channel": + purpose = payload.get("purpose", {}) + topic = payload.get("topic", {}) + if not isinstance(purpose, dict) or not isinstance(topic, dict): + raise ProjectionMappingError("Invalid Slack conversation description") + return ( + SlackChannelEntity( + channel_id=record.identity.native_id, + title=_string(payload.get("name"), default=record.identity.native_id), + purpose=_string(purpose.get("value")), + topic=_string(topic.get("value")), + breadcrumbs=[], + ), + ) + if record.identity.record_type != "message" or not record.identity.container_id: + raise ProjectionMappingError("Slack message requires its conversation identity") + timestamp = datetime.fromtimestamp(float(record.identity.native_id), timezone.utc) + metadata = {**payload, "channel": {"id": record.identity.container_id}} + text = payload.get("text", "") + if not isinstance(text, str): + raise ProjectionMappingError("Slack message text must be a string when present") + entity = SlackMessageEntity.from_api(metadata, breadcrumbs=[]) + entity.text = text + entity.created_at = entity.message_time = timestamp + edited = payload.get("edited") + if isinstance(edited, dict) and edited.get("ts"): + entity.updated_at = datetime.fromtimestamp(float(edited["ts"]), timezone.utc) + return (entity,) + + +def _wispr_range_text(response: dict[str, JsonValue], field: str, offset: int) -> str: + """Explicit provider absence has no transcript text; ambiguous missing data fails.""" + native = response.get(field) + if ( + field == "transcript" + and offset == 0 + and response.get("has_transcript") is False + and "transcript" in response + and native is None + ): + return "" + if not isinstance(native, str): + raise ProjectionMappingError("Wispr requested text is unavailable or malformed") + return native + + +_WISPR_TRANSCRIPT_HEADER = ( + "<<
नमस्ते — café
" + source = record( + part( + text.encode(), + "text/html", + headers=[{"name": "Content-Type", "value": "text/html; charset=gb2312"}], + ) + ) + before = source.model_dump() + result = await map_gmail(source, AsyncMock(), tmp_path) + assert next(tmp_path.glob("*.html")).read_text() == text + assert result.parts[0].part.charset_recoveries[0].source_path == "/payload" + assert result.parts[0].part.charset_recoveries[0].from_charset == "gb2312" + assert source.model_dump() == before + + +@pytest.mark.asyncio +async def test_alternative_prefers_full_html_not_plain_duplicate_or_snippet(tmp_path): + source = record( + { + "mimeType": "multipart/alternative", + "parts": [ + part(b"unselected plain"), + part(b"Full HTML body
", "text/html"), + ], + } + ) + entities = await map_gmail(source, AsyncMock(), tmp_path) + entities = entities.entities + text = next(tmp_path.glob("*.html")).read_text() + assert text == "Full HTML body
" + assert entities[0].local_path is not None + + +@pytest.mark.parametrize("html_body", [b"", b" \n\t "]) +@pytest.mark.asyncio +async def test_alternative_falls_back_to_plain_when_html_is_empty(tmp_path, html_body): + source = record( + { + "mimeType": "multipart/alternative", + "parts": [part(b"usable plain body", "text/plain"), part(html_body, "text/html")], + } + ) + + await map_gmail(source, AsyncMock(), tmp_path) + + assert next(tmp_path.glob("*.html")).read_text() == "usable plain body" + + +@pytest.mark.asyncio +async def test_related_empty_html_alternative_falls_back_without_materializing_siblings(tmp_path): + source = record( + { + "mimeType": "multipart/alternative", + "parts": [ + part(b"usable plain body", "text/plain"), + { + "mimeType": "multipart/related", + "parts": [ + part(b" \n ", "text/html"), + part(b"image bytes", "image/png", filename="image.png"), + ], + }, + ], + } + ) + + entities = await map_gmail(source, AsyncMock(), tmp_path) + + assert len(entities.entities) == 1 + assert next(tmp_path.glob("*.html")).read_text() == "
usable plain body" + + +@pytest.mark.asyncio +async def test_unreadable_html_alternative_does_not_fall_back_to_plain(tmp_path): + source = record( + { + "mimeType": "multipart/alternative", + "parts": [ + part(b"usable plain body", "text/plain"), + {"mimeType": "text/html", "body": {"attachmentId": "missing", "size": 0}}, + ], + } + ) + + with pytest.raises(ValueError, match="exactly one"): + await map_gmail(source, AsyncMock(), tmp_path) + + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("nested", [False, True]) +@pytest.mark.asyncio +async def test_last_readable_html_skips_earlier_missing_blob_and_reads_selected_once( + tmp_path, nested +): + payload = { + "mimeType": "multipart/alternative", + "parts": [ + {"mimeType": "text/html", "body": {"attachmentId": "missing", "size": 0}}, + {"mimeType": "text/html", "body": {"attachmentId": "selected", "size": 0}}, + ], + } + body_path = "/payload/parts/1/body" + if nested: + payload = { + "mimeType": "multipart/alternative", + "parts": [ + part(b"plain", "text/plain"), + {"mimeType": "multipart/related", "parts": [payload]}, + ], + } + body_path = "/payload/parts/1/parts/0/parts/1/body" + source = record(payload) + content = b"
Selected external HTML
" + source = source.model_copy(update={"blobs": (blob(source, content, body_path),)}) + storage = AsyncMock() + storage.read_file.return_value = content + + await map_gmail(source, storage, tmp_path) + + assert next(tmp_path.glob("*.html")).read_text() == content.decode() + storage.read_file.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_external_blob_and_duplicate_filename_attachments_get_distinct_identity(tmp_path): + source = record( + { + "mimeType": "multipart/mixed", + "parts": [ + {"mimeType": "text/html", "body": {"attachmentId": "body", "size": 10}}, + part(b"one", "text/plain", filename="../../same.txt", partId="1"), + part(b"two", "text/plain", filename="../../same.txt", partId="2"), + ], + } + ) + data = b"body" + source = source.model_copy(update={"blobs": (blob(source, data, "/payload/parts/0/body"),)}) + # Canonical SHA and actual size are authoritative even when native size differs. + assert source.payload["payload"]["parts"][0]["body"]["size"] != len(data) + storage = AsyncMock() + storage.read_file.return_value = data + entities = await map_gmail(source, storage, tmp_path) + entities = entities.entities + assert len(entities) == 3 + assert entities[1].attachment_key != entities[2].attachment_key + from pathlib import Path + + assert all(Path(e.local_path).parent == tmp_path for e in entities) + assert {Path(e.local_path).read_bytes() for e in entities[1:]} == {b"one", b"two"} + + +@pytest.mark.asyncio +async def test_missing_blob_ref_fails_instead_of_indexing_snippet(tmp_path): + source = record({"mimeType": "text/html", "body": {"attachmentId": "body", "size": 3}}) + with pytest.raises(ValueError, match="exactly one"): + await map_gmail(source, AsyncMock(), tmp_path) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_cross_source_blob_is_rejected_before_storage_access(): + source = record(part()) + reference = blob(source, b"abc").model_copy(update={"key": "canonical/other/blobs/sha256/hash"}) + storage = AsyncMock() + with pytest.raises(ValueError, match="source scope"): + await read_blob(source, reference, storage) + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_corrupt_blob_bytes_fail_before_projection(tmp_path): + source = record({"mimeType": "text/plain", "body": {"attachmentId": "body", "size": 3}}) + source = source.model_copy(update={"blobs": (blob(source, b"abc"),)}) + storage = AsyncMock() + storage.read_file.return_value = b"bad" + with pytest.raises(ValueError, match="digest"): + await map_gmail(source, storage, tmp_path) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_missing_attachment_keeps_available_body_searchable_and_original_partial(tmp_path): + source = record( + { + "mimeType": "multipart/mixed", + "parts": [ + part(b"Available message text"), + { + "mimeType": "application/pdf", + "filename": "uncaptured.pdf", + "body": {"attachmentId": "large", "size": 500000000}, + }, + ], + }, + completeness="partial", + ) + before = source.model_dump() + storage = AsyncMock() + entities = await map_gmail(source, storage, tmp_path) + entities = entities.entities + assert len(entities) == 1 + assert "Available message text" in next(tmp_path.glob("*.html")).read_text() + assert source.model_dump() == before + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_partial_capture_still_requires_actual_message_body(tmp_path): + source = record( + {"mimeType": "text/plain", "body": {"attachmentId": "body", "size": 3}}, + completeness="partial", + ) + with pytest.raises(ValueError, match="exactly one"): + await map_gmail(source, AsyncMock(), tmp_path) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_related_html_branch_does_not_get_replaced_by_plain_alternative(tmp_path): + source = record( + { + "mimeType": "multipart/alternative", + "parts": [ + part(b"plain duplicate"), + { + "mimeType": "multipart/related", + "parts": [ + part(b"Complete rich content
", "text/html"), + part(b"image bytes", "image/png", filename="image.png"), + ], + }, + ], + } + ) + entities = await map_gmail(source, AsyncMock(), tmp_path) + entities = entities.entities + assert len(entities) == 2 + from pathlib import Path + + assert Path(entities[0].local_path).read_text().strip() == "Complete rich content
" + assert entities[1].file_type == "png" + + +@pytest.mark.asyncio +async def test_inline_native_size_discrepancy_preserves_data_and_metadata(tmp_path): + source = record(part(b"Retained full body
", "text/html")) + source.payload["payload"]["body"]["size"] -= 1 + before = source.model_dump() + mapped = await map_gmail(source, AsyncMock(), tmp_path) + from pathlib import Path + + assert Path(mapped.entities[0].local_path).read_bytes() == b"Retained full body
" + assert source.model_dump() == before + source.payload["payload"]["body"]["data"] = "invalid!!" + with pytest.raises(ValueError): + await map_gmail(source, AsyncMock(), tmp_path) diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py new file mode 100644 index 0000000000..899fe92df4 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py @@ -0,0 +1,340 @@ +"""Actual Gmail transport/page source, lifecycle and PostgreSQL; native HTTP is synthetic.""" + +from contextlib import aclosing +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.requests import CaptureBatch +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components +from airweave.domains.entities.canonical.tests.test_changes_cycles import record +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline +from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt +from airweave.models.entity import Entity +from airweave.models.sync_cursor import SyncCursor +from airweave.models.sync_job import SyncJob +from airweave.platform.configs.config import GmailConfig +from airweave.platform.sources.gmail import GmailSource +from airweave.platform.sources.tests.test_gmail_capture import BASE, message + + +class NativeHTTP: + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + + async def handle(self, request): + if request.url.params.get("fields") == "emailAddress": + return httpx.Response(200, json={"emailAddress": "synthetic@example.test"}) + self.calls.append((str(request.url).split("?")[0], dict(request.url.params))) + expected, response = self.responses.pop(0) + assert str(request.url).split("?")[0] == BASE + expected + if isinstance(response, Exception): + raise response + status, payload = response if isinstance(response, tuple) else (200, response) + return httpx.Response(status, json=payload) + + +async def setup(database, source, native, *, attempt=1, query=None): + service, fence = source + client = httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) + connector = await GmailSource.create( + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=client, + config=GmailConfig( + expected_mailbox="synthetic@example.test", + gmail_query=query, + included_labels=[], + excluded_labels=[], + excluded_categories=[], + ), + ) + ctx, _, runtime, bus = components(database, source) + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + connector.canonical_record_types, + CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), + connector.canonical_container_parents, + page_source=connector, + files=MagicMock(), + ) + await pipeline.start(ctx) + return pipeline, ctx, runtime, client + + +async def no_limits(): + pass + + +async def run(pipeline, ctx, runtime): + await pipeline.run_scans(ctx, runtime, no_limits) + await pipeline.cleanup_orphaned_entities(ctx, runtime) + await pipeline.save_checkpoint(ctx, runtime) + + +async def test_full_resume_interleaved_history_known_omission_and_next_delta(database, source): + service, fence = source + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(record("known"),))) + native = NativeHTTP( + [ + ("/profile", {"historyId": "100"}), + ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "second"}), + ("/messages/a", message("a")), + ( + "/history", + {"historyId": "110", "history": [{"messagesAdded": [{"message": {"id": "c"}}]}]}, + ), + ("/messages/c", message("c")), + ("/messages", ConnectionError("synthetic interruption")), + ("/messages", {"messages": [{"id": "b"}]}), + ("/messages/b", message("b")), + ( + "/history", + {"historyId": "120", "history": [{"messagesDeleted": [{"message": {"id": "a"}}]}]}, + ), + ("/messages/a", (404, {"error": {"message": "gone"}})), + ("/messages/known", message("known")), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with aclosing(client): + with pytest.raises(ConnectionError, match="interruption"): + await run(pipeline, ctx, runtime) + async with database() as db: + cursor = await db.scalar(select(SyncCursor)) + assert "canonical_checkpoint" not in cursor.cursor_data + assert cursor.cursor_data["canonical_cycle"]["promoted_checkpoint"] is None + pipeline, ctx, runtime, client = await setup(database, source, native, attempt=2) + async with aclosing(client): + await run(pipeline, ctx, runtime) + assert len([url for url, _ in native.calls if url.endswith("/profile")]) == 1 + assert [ + params["startHistoryId"] for url, params in native.calls if url.endswith("/history") + ] == ["100", "110"] + assert native.calls[6][1]["pageToken"] == "second" + assert not native.responses + async with database() as db: + full = await service.read_cycle(db, pipeline._writer()) + assert full.promoted_checkpoint.checkpoint.value == {"history_id": "120"} + assert { + r.native_id for r in (await db.scalars(select(Entity))).all() if r.deleted_at is None + } == {"b", "c", "known"} + old = await db.get(SyncJob, fence.job_id) + old.status = "completed" + job_id = uuid4() + db.add( + SyncJob( + id=job_id, + sync_id=fence.sync_id, + organization_id=fence.organization_id, + status="running", + ) + ) + await db.commit() + next_fence = await service.activate_writer( + db, fence.organization_id, fence.sync_id, job_id, attempt_id=uuid4(), attempt_number=1 + ) + delta = NativeHTTP( + [ + ( + "/history", + {"historyId": "130", "history": [{"labelsAdded": [{"message": {"id": "b"}}]}]}, + ), + ("/messages/b", message("b", labelIds=["TRASH"])), + ] + ) + pipeline, ctx, runtime, client = await setup(database, (service, next_fence), delta) + async with aclosing(client): + await run(pipeline, ctx, runtime) + assert delta.calls[0][1]["startHistoryId"] == "120" + assert not delta.responses + async with database() as db: + changed = await service.read_cycle(db, pipeline._writer()) + assert changed.mode == "changes" and changed.last_full_capture == full.last_full_capture + assert changed.promoted_checkpoint.checkpoint.value == {"history_id": "130"} + assert ( + await db.scalar(select(Entity).where(Entity.native_id == "known")) + ).deleted_at is None + assert (await db.scalar(select(Entity).where(Entity.native_id == "b"))).source_payload[ + "labelIds" + ] == ["TRASH"] + + +async def test_filtered_mutating_view_omission_is_not_provider_deletion(database, source): + service, fence = source + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(record("no-longer-listed"),))) + native = NativeHTTP( + [("/messages", {"messages": [{"id": "a"}]}), ("/messages/a", message("a", labelIds=[]))] + ) + pipeline, ctx, runtime, client = await setup(database, source, native, query="in:inbox") + async with aclosing(client): + await run(pipeline, ctx, runtime) + assert native.calls[0][1]["q"] == "in:inbox" + async with database() as db: + cycle = await service.read_cycle(db, pipeline._writer()) + removed = await db.scalar(select(Entity).where(Entity.native_id == "no-longer-listed")) + assert cycle.promoted_checkpoint is None + assert removed.removal_reason == "absent" and removed.deleted_at is not None + + +async def test_history_expiry_restarts_full_without_using_partial_sightings(database, source): + native = NativeHTTP( + [ + ("/profile", {"historyId": "old"}), + ("/messages", {"messages": [{"id": "partial"}]}), + ("/messages/partial", message("partial")), + ("/history", (404, {"error": {"message": "expired"}})), + ("/profile", {"historyId": "fresh"}), + ("/messages", {}), + ("/history", {"historyId": "terminal"}), + ("/messages/partial", (404, {"error": {"message": "gone"}})), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with aclosing(client): + await run(pipeline, ctx, runtime) + async with database() as db: + cycle = await source[0].read_cycle(db, pipeline._writer()) + partial = await db.scalar(select(Entity).where(Entity.native_id == "partial")) + assert cycle.starting_checkpoint.value == {"history_id": "fresh"} + assert cycle.promoted_checkpoint.checkpoint.value == {"history_id": "terminal"} + assert partial.deleted_at is not None and partial.removal_reason == "provider_deleted" + + +async def test_history_offset_resume_and_lost_final_ack_do_not_rehydrate_committed_messages( + database, source +): + raw = { + "historyId": "20", + "history": [{"id": "entry", "messages": [{"id": str(i)} for i in range(51)]}], + } + responses = [("/profile", {"historyId": "10"}), ("/messages", {}), ("/history", raw)] + responses.extend((f"/messages/{i}", message(str(i))) for i in range(50)) + responses.extend( + [ + ("/history", ConnectionError("synthetic partial history stop")), + ("/history", raw), + ("/messages/50", message("50")), + ] + ) + native = NativeHTTP(responses) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with aclosing(client): + with pytest.raises(ConnectionError, match="partial history"): + await run(pipeline, ctx, runtime) + pipeline, ctx, runtime, client = await setup(database, source, native, attempt=2) + async with aclosing(client): + await pipeline.run_scans(ctx, runtime, no_limits) + # The process disappears after its durable final page, before final checkpoint publication. + assert not native.responses + calls = len(native.calls) + pipeline, ctx, runtime, client = await setup(database, source, native, attempt=3) + async with aclosing(client): + await run(pipeline, ctx, runtime) + assert len(native.calls) == calls + assert len([url for url, _ in native.calls if "/messages/" in url]) == 51 + async with database() as db: + state = await source[0].read_cycle(db, pipeline._writer()) + assert state.promoted_checkpoint.checkpoint.value == {"history_id": "20"} + assert state.phase == "complete" + + +async def test_new_job_resumes_active_cycle_after_exhausted_server_error( + database, source, monkeypatch +): + from tenacity import wait_none + + from airweave.domains.sources.exceptions import SourceServerError + from airweave.domains.syncs.jobs.repository import SyncJobRepository + + monkeypatch.setattr( + GmailSource, + "_get_capture_json", + GmailSource._get_capture_json.retry_with(wait=wait_none()), + ) + native = NativeHTTP( + [ + ("/profile", {"historyId": "100"}), + ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "second"}), + ("/messages/a", message("a")), + ("/history", {"historyId": "110"}), + *[("/messages", (503, {"error": {"message": "synthetic transient"}}))] * 5, + ("/messages", {"messages": [{"id": "b"}]}), + ("/messages/b", message("b")), + ("/history", {"historyId": "120"}), + ] + ) + service, fence = source + pipeline, ctx, runtime, client = await setup(database, source, native) + async with aclosing(client): + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + incomplete = await service.read_cycle(db, pipeline._writer()) + assert incomplete.phase == "active" + assert incomplete.promoted_checkpoint is None + # Model the production orchestrator's terminal failure, then the next + # eligible scheduled job. This is not a Temporal workflow execution. + old = await db.get(SyncJob, fence.job_id) + old.status = "failed" + await db.commit() + assert not await SyncJobRepository().get_active_for_sync(db, fence.sync_id, ctx) + job_id = uuid4() + db.add( + SyncJob( + id=job_id, + sync_id=fence.sync_id, + organization_id=fence.organization_id, + status="running", + ) + ) + await db.commit() + next_fence = await service.activate_writer( + db, + fence.organization_id, + fence.sync_id, + job_id, + attempt_id=uuid4(), + attempt_number=1, + ) + pipeline, ctx, runtime, client = await setup(database, (service, next_fence), native) + async with aclosing(client): + await run(pipeline, ctx, runtime) + assert not native.responses + assert sum(url.endswith("/messages/a") for url, _ in native.calls) == 1 + assert all(params["pageToken"] == "second" for url, params in native.calls[4:10]) + async with database() as db: + completed = await service.read_cycle(db, pipeline._writer()) + assert completed.version.cycle_id == incomplete.version.cycle_id + assert completed.phase == "complete" + assert completed.promoted_checkpoint.checkpoint.value == {"history_id": "120"} + assert {row.native_id for row in (await db.scalars(select(Entity))).all()} == {"a", "b"} + + +async def test_wrong_native_mailbox_never_starts_capture(database, source): + native = NativeHTTP([]) + async with database() as db: + before = (await db.execute(select(SyncCursor))).scalars().all() + before = [(item.id, item.cursor_data) for item in before] + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) as client: + with pytest.raises(ValueError, match="does not match"): + await GmailSource.create( + auth=StaticTokenProvider("same-broker-account"), + logger=MagicMock(), + http_client=client, + config=GmailConfig(expected_mailbox="different@example.test"), + ) + assert native.calls == [] + async with database() as db: + assert (await db.execute(select(Entity))).scalars().all() == [] + after = (await db.execute(select(SyncCursor))).scalars().all() + assert [(item.id, item.cursor_data) for item in after] == before diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py new file mode 100644 index 0000000000..57c88dd83f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py @@ -0,0 +1,261 @@ +"""Gmail harness proof uses fresh processes and isolated PostgreSQL, no live provider.""" + +import json + +from sqlalchemy import text + +from airweave.domains.entities.canonical.tests.test_resume_probe import subprocess_script + +SCRIPT = r""" +import asyncio, os, sys +from pathlib import Path +from uuid import uuid4 +from unittest.mock import MagicMock +sys.path.insert(0, "tests/live") +import conftest +import httpx +from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker +from airweave.domains.entities.canonical.requests import WriterFence +from airweave.domains.entities.canonical.service import CanonicalCaptureService +from airweave.domains.entities.canonical.store import CanonicalRecordStore +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components, orchestrator +from airweave.domains.entities.canonical.tests.test_slack_recovery import run +from airweave.domains.entities.canonical.tests.test_gmail_recovery import NativeHTTP +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline +from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt +from airweave.platform.configs.config import GmailConfig +from airweave.platform.sources.gmail import GmailSource +from airweave.platform.sources.tests.test_gmail_capture import message +from gmail_resume import GmailResumeProbe +from provider_lifecycle import bounded_page_source + +async def main(): + engine = create_async_engine(os.environ["CANONICAL_TEST_DATABASE_URL"], + connect_args={"server_settings": {"search_path": os.environ["TEST_SCHEMA"]}}) + sessions = async_sessionmaker(engine, expire_on_commit=False) + fence = WriterFence.model_validate_json(os.environ["TEST_FENCE"]) + service = CanonicalCaptureService(CanonicalRecordStore()) + mode = os.environ["TEST_MODE"] + responses = [('/profile', {'historyId': 'initial'}), + ('/messages', {'messages': [{'id': 'a'}], 'nextPageToken': 'next'}), + ('/messages/a', message('a'))] if mode == 'interrupt' else [ + ('/history', {'historyId': 'later', 'history': [{'messages': [{'id': 'b'}]}]}), + ('/messages/b', message('b'))] + native = NativeHTTP(responses) + counters = {'provider_requests': 0, 'records_observed': 0, 'capture_profile_requests': 0, + 'blob_bytes_written': 0} + async def handle(request): + counters['provider_requests'] += 1 + if request.url.params.get('fields') == 'emailAddress': + return await native.handle(request) + if request.url.path.endswith('/profile'): + counters['capture_profile_requests'] += 1 + return await native.handle(request) + client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + connector = await GmailSource.create(auth=StaticTokenProvider('synthetic'), + logger=MagicMock(), http_client=client, config=GmailConfig( + expected_mailbox="synthetic@example.test", + included_labels=[], excluded_labels=[], excluded_categories=[])) + ctx, _, runtime, bus = components(sessions, (service, fence)) + attempt = CaptureAttempt(id=fence.attempt_id if mode == 'interrupt' else uuid4(), + number=1 if mode == 'interrupt' else 2) + probe = GmailResumeProbe(sessions, fence.organization_id, fence.sync_id, fence.job_id, + attempt.id, Path(os.environ['TEST_TARGET']), mode, counters) + wrapper = bounded_page_source(connector, counters, 250, probe) + pipeline = CanonicalCapturePipeline(service, sessions, bus, connector.canonical_record_types, + attempt, connector.canonical_container_parents, page_source=wrapper, files=MagicMock()) + runtime.source, runtime.canonical_capture = connector, pipeline + instance = orchestrator(ctx, pipeline, runtime, None, bus) + instance.stream = None + await run(instance) +asyncio.run(main()) +""" + + +async def test_gmail_process_loss_resumes_saved_plan_without_profile(database, source, tmp_path): + async with database() as db: + schema = await db.scalar(text("select current_schema()")) + env = { + "TEST_SCHEMA": schema, + "TEST_FENCE": source[1].model_dump_json(), + "TEST_TARGET": str(tmp_path / "target.json"), + } + code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "interrupt"}) + assert code == 75, error + first = json.loads(output.splitlines()[-1]) + assert first["intentional_interruption"] and first["records_observed"] == 1 + assert first["capture_profile_requests"] == 1 + code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "resume"}) + assert code == 76, error + second = json.loads(output.splitlines()[-1]) + assert second["gmail_recovery_verified"] and second["resumed_saved_page"] + assert second["capture_profile_requests"] == 0 and second["records_observed"] == 1 + assert not second["full_scope_completed"] and not second["checkpoint_saved"] + assert (tmp_path / "target.json").stat().st_mode & 0o077 == 0 + + +async def test_gmail_harness_budgets_capabilities_and_early_completion(tmp_path): + script = r""" +import asyncio, json, os, sys +from types import SimpleNamespace +sys.path.insert(0, 'tests/live') +import conftest +import httpx +import provider_lifecycle as lifecycle +from airweave.domains.entities.canonical.page_source import ( + CheckpointedPageSource, KnownObjectSource) + +class Plain: + canonical_record_types = ('message',) + canonical_container_parents = {} + capture_cycle_configuration = None + +class Known(Plain): + async def refresh_known(self, record, *, files): + return record + +class Planned(Known): + async def prepare_cycle(self, previous): + return previous + def initial_continuation(self, cycle): + return cycle + +async def main(): + counters = {'records_observed': 0} + for source, known, planned in [(Plain(), False, False), (Known(), True, False), + (Planned(), True, True)]: + wrapper = lifecycle.bounded_page_source(source, counters, 1) + assert isinstance(wrapper, CheckpointedPageSource) == planned + assert isinstance(wrapper, KnownObjectSource) == known + assert await wrapper.prepare_cycle('plan') == 'plan' + assert wrapper.initial_continuation('saved') == 'saved' + assert await wrapper.refresh_known('record', files=None) == 'record' + try: + await wrapper.refresh_known('record', files=None) + except lifecycle.BudgetExceeded: + pass + else: + raise AssertionError('Exact reads bypassed record budget') + counters = {'identity_profile_requests': 0, 'capture_profile_requests': 0, + 'resumed_changes_requests': 0} + request = httpx.Request('GET', 'https://gmail.googleapis.com/gmail/v1/users/me/profile') + lifecycle.count_gmail_request('gmail', request, {}, counters, False) + lifecycle.count_gmail_request('gmail', request, {}, counters, True) + assert counters['identity_profile_requests'] == counters['capture_profile_requests'] == 1 + for early in (False, True): + snapshots = [] + async def execute(path, timeout): + snapshots.append(json.loads(path.read_text())) + if early: + return 0, {'failed': False, 'full_scope_completed': True} + second = len(snapshots) == 2 + return (76 if second else 75), {'failed': False, 'intentional_interruption': not second, + 'gmail_recovery_verified': second, 'provider_requests': 7, + 'records_observed': 50, 'blob_bytes_written': 20} + lifecycle.execute_trial = execute + results = await lifecycle.run_trials({'provider': 'gmail', 'gmail_resume': True, + 'root': os.environ['TEST_ROOT'], 'request_limit': 600, 'record_limit': 250, + 'blob_byte_limit': 100, 'timeout': 600}) + if early: + assert len(snapshots) == 1 and results[0]['recovery_not_exercised'] + else: + assert len(snapshots) == 2 + assert [s['request_limit'] for s in snapshots] == [600, 593] + assert [s['record_limit'] for s in snapshots] == [250, 200] + assert [s['blob_byte_limit'] for s in snapshots] == [100, 80] + assert snapshots[0]['job_id'] == snapshots[1]['job_id'] + assert [s['attempt_number'] for s in snapshots] == [1, 2] + print('verified') +asyncio.run(main()) +""" + code, output, error = await subprocess_script(script, {"TEST_ROOT": str(tmp_path)}) + assert code == 0, error + assert output.splitlines()[-1] == "verified" + + +async def test_actual_lifecycle_child_context_reaches_committed_gmail_page( + database, source, tmp_path +): + """Exercise the real child composition, not the separate pipeline-test context.""" + script = r""" +import asyncio, json, os, sys +from contextlib import asynccontextmanager +from unittest.mock import MagicMock +sys.path.insert(0, 'tests/live') +import conftest +import httpx +import provider_lifecycle as lifecycle +from airweave.domains.entities.canonical.tests.test_gmail_recovery import NativeHTTP +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.platform.configs.config import GmailConfig +from airweave.platform.sources.gmail import GmailSource +from airweave.platform.sources.tests.test_gmail_capture import message + +@asynccontextmanager +async def synthetic_source(name, account, expected_email, key, fence, **options): + assert name == 'gmail' and options['gmail_unfiltered'] is True + native = NativeHTTP([ + ('/profile', {'historyId': 'capture-boundary'}), + ('/messages', {'messages': [{'id': 'synthetic-a'}], 'nextPageToken': 'next'}), + ('/messages/synthetic-a', message('synthetic-a')), + ]) + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle), + event_hooks={'request': [options['request_hook']]}) as client: + connector = await GmailSource.create(auth=StaticTokenProvider('synthetic'), + logger=MagicMock(), http_client=client, config=GmailConfig( + expected_mailbox=expected_email, + included_labels=[], excluded_labels=[], excluded_categories=[])) + yield connector, 'provider_email' + +lifecycle.rest_source = synthetic_source +# This offline test uses the disposable database fixture (TCP in CI, socket locally). +# Keep the live harness's private-socket guard unchanged for real account trials. +lifecycle.harness.test_database_url = lambda: os.environ['CANONICAL_TEST_DATABASE_URL'] +# All provider traffic uses the scripted MockTransport; synthetic credentials only. +os.environ['COMPOSIO_API_KEY'] = 'synthetic-test-only' +os.environ['LIVE_GMAIL_ACCOUNT_ID'] = 'synthetic-test-only' +os.environ['LIVE_EXPECTED_EMAIL'] = 'synthetic@example.test' +raise SystemExit(asyncio.run(lifecycle.child(json.loads(os.environ['TEST_MANIFEST'])))) +""" + from sqlalchemy import select + + from airweave.models.capture_scan import CaptureScan + from airweave.models.entity import Entity + from airweave.models.sync_cursor import SyncCursor + + fence = source[1] + async with database() as db: + schema = await db.scalar(text("select current_schema()")) + manifest = { + "provider": "gmail", + "schema": schema, + "root": str(tmp_path), + "organization_id": str(fence.organization_id), + "sync_id": str(fence.sync_id), + "job_id": str(fence.job_id), + "attempt_number": 2, + "gmail_resume": True, + "gmail_unfiltered": True, + "resume_stage": "interrupt", + "record_limit": 250, + "request_limit": 600, + "timeout": 30, + "file_byte_limit": 10 * 1024 * 1024, + "blob_byte_limit": 256 * 1024 * 1024, + } + code, output, error = await subprocess_script(script, {"TEST_MANIFEST": json.dumps(manifest)}) + assert code == 75, (output, error) + result = json.loads(output.splitlines()[-1]) + assert result["intentional_interruption"] and not result["failed"] + assert result["identity_profile_requests"] == result["capture_profile_requests"] == 1 + assert result["provider_requests"] == 4 and result["records_observed"] == 1 + assert not result["checkpoint_saved"] and not result["full_scope_completed"] + async with database() as db: + record = await db.scalar(select(Entity).where(Entity.native_id == "synthetic-a")) + scan = await db.scalar(select(CaptureScan).where(CaptureScan.sync_id == fence.sync_id)) + cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + assert record is not None and record.record_revision == 1 + assert scan.phase == "collecting" and scan.continuation["phase"] == "history" + assert cursor.cursor_data["canonical_cycle"]["phase"] == "active" + assert "canonical_checkpoint" not in cursor.cursor_data diff --git a/backend/airweave/domains/entities/canonical/tests/test_http.py b/backend/airweave/domains/entities/canonical/tests/test_http.py new file mode 100644 index 0000000000..3d767f3f67 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_http.py @@ -0,0 +1,149 @@ +"""HTTP contract verification using the real PostgreSQL-backed query service.""" + +from types import SimpleNamespace +from uuid import uuid4 + +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient + +from airweave.api import deps +from airweave.api.v1.endpoints.records import record_error_response, router +from airweave.db.session import get_db +from airweave.domains.entities.canonical.query import CanonicalQueryService +from airweave.domains.entities.canonical.query_store import CanonicalQueryStore +from airweave.domains.entities.canonical.requests import BlobReference, CaptureBatch, RecordIdentity +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.tests.helpers import observation + + +def query_app(database, context): + """Use the real query service with synthetic HTTP authentication and signing.""" + app = FastAPI() + app.include_router(router, prefix="/sync") + app.add_exception_handler(CanonicalStoreError, record_error_response) + + async def session(): + async with database() as db: + yield db + + app.dependency_overrides[deps.get_context] = context + app.dependency_overrides[get_db] = session + app.dependency_overrides[deps.get_canonical_query_service] = lambda: CanonicalQueryService( + CanonicalRecordStore(), CanonicalQueryStore(), "synthetic-cursor-secret" + ) + return app + + +async def test_http_read_list_changes_and_cross_tenant_denial(database, source): + capture, fence = source + async with database() as db: + result = await capture.capture( + db, CaptureBatch(fence=fence, records=(observation("one"), observation("two"))) + ) + owner = fence.organization_id + + async def context(): + return SimpleNamespace(organization=SimpleNamespace(id=owner)) + + app = query_app(database, context) + base = f"/sync/{fence.sync_id}/records" + record_path = f"{base}/{result.changes[0].record.id}" + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.get(record_path) + assert response.status_code == 200 + assert response.json()["identity"]["native_id"] == "one" + listed = await client.get(base, params={"limit": 1}) + assert listed.status_code == 200 + assert listed.json()["has_more"] + continued = await client.get( + base, params={"cursor": listed.json()["next_cursor"], "limit": 1} + ) + assert continued.status_code == 200 and not continued.json()["has_more"] + changes = await client.get(base + "/changes") + assert changes.status_code == 200 + assert len(changes.json()["changes"]) == 2 + invalid = await client.get(base, params={"cursor": "not-a-valid-cursor"}) + assert invalid.status_code == 400 + assert invalid.json()["error"]["code"] == "invalid_cursor" + oversized = await client.get(base, params={"limit": 501}) + assert oversized.status_code == 422 + owner = uuid4() + assert (await client.get(record_path)).status_code == 404 + assert (await client.get(base)).status_code == 404 + assert (await client.get(base + "/changes")).status_code == 404 + + +async def test_http_change_continuation_redacts_withdrawn_snapshots_and_rejects_foreign_cursors( + database, source +): + capture, fence = source + parent_identity = RecordIdentity(record_type="calendar", native_id="cal") + parent = observation(identity=parent_identity) + child = observation( + "event", + "cal", + parent=parent_identity, + blobs=(BlobReference(key="immutable", sha256="a" * 64, size_bytes=7),), + ) + async with database() as db: + initial = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, child))) + original_child = initial.changes[1].record + owner = fence.organization_id + + async def context(): + return SimpleNamespace(organization=SimpleNamespace(id=owner)) + + app = query_app(database, context) + base = f"/sync/{fence.sync_id}/records" + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + first_response = await client.get(base + "/changes", params={"limit": 1}) + assert first_response.status_code == 200 + first = first_response.json() + assert first["has_more"] and first["high_watermark"] == 2 + assert [change["sequence"] for change in first["changes"]] == [1] + + # Commit withdrawal between pages, before descendant reconciliation. + tombstone = parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}) + async with database() as db: + withdrawn = await capture.capture(db, CaptureBatch(fence=fence, records=(tombstone,))) + assert withdrawn.sequence == 3 + + params = {"cursor": first["next_cursor"], "limit": 1} + continued_response = await client.get(base + "/changes", params=params) + assert continued_response.status_code == 200 + continued = continued_response.json() + assert not continued["has_more"] and continued["high_watermark"] == 2 + assert [change["sequence"] for change in continued["changes"]] == [2] + historical = continued["changes"][0] + assert historical["kind"] == "upsert" + assert historical["record"]["id"] == str(original_child.id) + assert historical["record"]["revision"] == original_child.revision + assert historical["record"]["deleted_at"] is None + assert historical["record"]["content_access"] == "unavailable" + assert historical["record"]["payload"] == {} and historical["record"]["blobs"] == [] + replay = await client.get(base + "/changes", params=params) + assert replay.status_code == 200 and replay.json() == continued + + polled_response = await client.get( + base + "/changes", params={"cursor": continued["next_cursor"], "limit": 1} + ) + assert polled_response.status_code == 200 + polled = polled_response.json() + assert not polled["has_more"] and polled["high_watermark"] == 3 + assert [change["sequence"] for change in polled["changes"]] == [3] + deletion = polled["changes"][0] + assert deletion["kind"] == "delete" + assert deletion["record"]["removal_reason"] == "access_revoked" + assert deletion["record"]["content_access"] == "unavailable" + assert deletion["record"]["payload"] == {} and deletion["record"]["blobs"] == [] + empty = await client.get(base + "/changes", params={"cursor": polled["next_cursor"]}) + assert empty.status_code == 200 + assert empty.json()["changes"] == [] and empty.json()["high_watermark"] == 3 + + foreign_source = await client.get(f"/sync/{uuid4()}/records/changes", params=params) + foreign_operation = await client.get(base, params=params) + owner = uuid4() + foreign_owner = await client.get(base + "/changes", params=params) + for rejected in (foreign_source, foreign_operation, foreign_owner): + assert rejected.status_code == 400 + assert rejected.json()["error"]["code"] == "invalid_cursor" diff --git a/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py b/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py new file mode 100644 index 0000000000..de75b08699 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py @@ -0,0 +1,150 @@ +"""Original-only Linear projections with no provider calls or fabricated file content.""" + +import hashlib +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import AsyncMock +from uuid import UUID, uuid4 + +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity +from airweave.platform.entities.linear import LinearLinkedAttachmentEntity + +ISSUE = UUID(int=3) +STAMP = "2026-09-30T00:00:00Z" +CONTEXT = { + "id": str(ISSUE), + "identifier": "ENG-3", + "title": "Native title", + "url": "https://linear.app/example/issue/ENG-3", + "team": {"id": str(UUID(int=2)), "name": "Engineering"}, + "project": None, +} + + +def original(kind, **extra): + native_id = ISSUE if kind == "issue" else UUID(int=4) + payload = { + "id": str(native_id), + "createdAt": STAMP, + "updatedAt": STAMP, + "url": "https://linear.app/native-link", + **extra, + } + if kind == "issue": + payload.update(CONTEXT) + else: + payload["issue"] = CONTEXT + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity( + record_type=kind, + native_id=str(native_id), + container_id=None if kind == "issue" else str(ISSUE), + ), + revision=1, + payload=payload, + payload_schema_version=2, + capture_hash="hash", + content_hash=None, + completeness="complete", + observed_at=datetime.now(timezone.utc), + source_created_at=None, + source_updated_at=None, + deleted_at=None, + removal_reason=None, + blobs=(), + indexed_revision=None, + indexed_pipeline_version=None, + ) + + +@pytest.mark.asyncio +async def test_native_issue_and_comment_render_without_provider_or_blob_access(): + for kind, fields in [ + ("issue", {"description": "**native**"}), + ("comment", {"body": "reply", "user": None}), + ]: + record = original(kind, **fields) + storage = AsyncMock() + async with map_record(record, "linear", storage) as entities: + entities = entities.entities + assert entities[0].web_url == record.payload["url"] + assert entities[0].created_at.tzinfo is not None + storage.read_file.assert_not_called() + assert record.payload.get("description", record.payload.get("body")) == next( + iter(fields.values()) + ) + + +@pytest.mark.asyncio +async def test_link_is_explicit_metadata_not_invented_file(): + record = original("attachment", title="Document").model_copy(update={"completeness": "partial"}) + async with map_record(record, "linear", AsyncMock()) as entities: + entities = entities.entities + assert isinstance(entities[0], LinearLinkedAttachmentEntity) + assert "not retained" in entities[0].content_coverage + + +@pytest.mark.asyncio +async def test_missing_binary_cannot_claim_complete(): + with pytest.raises(ValueError, match="without retained bytes"): + async with map_record(original("attachment", title="Document"), "linear", AsyncMock()): + pass + + +@pytest.mark.asyncio +async def test_retained_file_verified_and_disposable(): + record = original("attachment", title="Document") + data = b"%PDF fixture" + digest = hashlib.sha256(data).hexdigest() + ref = BlobReference( + key=f"canonical/{record.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(data), + media_type="application/pdf", + source_path="/url", + ) + record = record.model_copy(update={"blobs": (ref,)}) + storage = AsyncMock() + storage.read_file.return_value = data + async with map_record(record, "linear", storage) as entities: + entities = entities.entities + path = Path(entities[0].local_path) + assert path.read_bytes() == data + assert entities[0].issue_identifier == "ENG-3" + assert not path.exists() + storage.read_file.assert_awaited_once_with(ref.key, max_bytes=len(data)) + storage.read_file.return_value = b"corrupt" + with pytest.raises(ValueError, match="digest"): + async with map_record(record, "linear", storage): + pass + + +@pytest.mark.asyncio +async def test_wrong_parent_or_missing_context_never_invents_metadata(): + record = original("comment", body="reply") + record = record.model_copy( + update={"identity": record.identity.model_copy(update={"container_id": str(UUID(int=99))})} + ) + with pytest.raises(ValueError, match="parent"): + async with map_record(record, "linear", AsyncMock()): + pass + payload = {**record.payload, "issue": {"id": str(ISSUE)}} + with pytest.raises(ValueError): + async with map_record( + record.model_copy(update={"payload": payload}), "linear", AsyncMock() + ): + pass + + +def test_link_projection_is_registered(): + from airweave.domains.entities.registry import EntityDefinitionRegistry + + registry = EntityDefinitionRegistry() + registry.build() + assert registry.get_short_name_by_class(LinearLinkedAttachmentEntity) diff --git a/backend/airweave/domains/entities/canonical/tests/test_linear_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_linear_recovery.py new file mode 100644 index 0000000000..2cd7793155 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_linear_recovery.py @@ -0,0 +1,202 @@ +"""Actual Linear query adapter and SQL capture authority with synthetic provider responses.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID, uuid4 + +import httpx +import pytest + +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components, orchestrator +from airweave.domains.entities.canonical.tests.test_slack_recovery import run, saved +from airweave.domains.sources.exceptions import SourceError +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline +from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt +from airweave.platform.configs.config import LinearConfig +from airweave.platform.sources.linear import LinearSource + +WORKSPACE, TEAM, ISSUE, COMMENT, ATTACHMENT = (UUID(int=n) for n in range(1, 6)) +STAMP = "2026-09-30T00:00:00Z" +ROOT = { + "id": str(ISSUE), + "team": {"id": str(TEAM)}, + "createdAt": STAMP, + "updatedAt": STAMP, + "archivedAt": STAMP, + "trashed": True, +} +CHILD = { + "id": str(COMMENT), + "issue": {"id": str(ISSUE)}, + "createdAt": STAMP, + "updatedAt": STAMP, + "body": "retained original", +} + + +def connection(nodes, *, following=False, cursor=None): + return {"nodes": nodes, "pageInfo": {"hasNextPage": following, "endCursor": cursor}} + + +def result(nodes, *, errors=None, following=False): + return { + "data": {"issues": connection(nodes, following=following)}, + **({"errors": errors} if errors else {}), + } + + +def children(kind, nodes, *, following=False): + return result( + [ + { + **ROOT, + kind: connection( + nodes, following=following, cursor="child-next" if following else None + ), + } + ] + ) + + +async def runner(database, source, responses, *, attempt=1): + service, fence = source + client = AsyncMock() + envelopes = [ + {"data": {"organization": {"id": str(WORKSPACE)}, "viewer": {"id": str(UUID(int=9))}}}, + {"data": {"teams": connection([{"id": str(TEAM)}])}}, + *responses, + ] + client.post.side_effect = [ + httpx.Response( + 200, json=value, request=httpx.Request("POST", "https://api.linear.app/graphql") + ) + for value in envelopes + ] + connector = await LinearSource.create( + auth=StaticTokenProvider("fixture"), + logger=MagicMock(), + http_client=client, + config=LinearConfig(workspace_id=WORKSPACE, team_ids=(TEAM,)), + ) + ctx, _, runtime, bus = components(database, source) + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + connector.canonical_record_types, + CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), + connector.canonical_container_parents, + page_source=connector, + files=MagicMock(), + ) + runtime.source, runtime.canonical_capture = connector, pipeline + instance = orchestrator(ctx, pipeline, runtime, None, bus) + instance.stream = None + return instance, client + + +async def assert_hidden(database, source, rows): + assert all(row.removal_reason == "scope_removed" and row.deleted_at is not None for row in rows) + async with database() as db: + for row in rows: + record = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, row.id + ) + assert record.content_access == "unavailable" and record.payload == {} + + +async def test_zero_outer_child_connection_withdraws_root_and_prior_child(database, source): + instance, client = await runner( + database, source, [result([ROOT]), children("comments", [CHILD]), result([])] + ) + await run(instance) + rows, cursor, _ = await saved(database) + assert {row.native_id for row in rows} == {str(ISSUE), str(COMMENT)} + await assert_hidden(database, source, rows) + assert cursor["canonical_cycle"]["phase"] == "complete" + child_queries = [call.kwargs["json"]["query"] for call in client.post.call_args_list[3:]] + assert all("filter: {id: {eq: $issueId}}" in query for query in child_queries) + + +async def test_zero_exact_membership_reconciles_old_root_and_both_children(database, source): + attachment = {**CHILD, "id": str(ATTACHMENT), "url": "https://external.example/link"} + first, _ = await runner( + database, + source, + [result([ROOT]), children("comments", [CHILD]), children("attachments", [attachment])], + ) + await run(first, checkpoint=False) + rows, _, _ = await saved(database) + assert len(rows) == 3 and all(row.deleted_at is None for row in rows) + # Archived and trashed fields are retained states, never deletion instructions. + async with database() as db: + original = await source[0].store.read( + db, + source[1].organization_id, + source[1].sync_id, + next(row.id for row in rows if row.native_id == str(ISSUE)), + ) + assert original.payload["archivedAt"] == STAMP and original.payload["trashed"] is True + second, _ = await runner(database, source, [result([]), result([])], attempt=2) + await run(second) + rows, cursor, _ = await saved(database) + await assert_hidden(database, source, rows) + assert cursor["canonical_cycle"]["phase"] == "complete" + + +@pytest.mark.parametrize( + "invalid", + [ + result([], errors=[{"extensions": {"type": "INTERNAL_SERVER_ERROR"}}]), + {"data": {"issues": None}}, + result([], following=True), + result([{**ROOT, "id": str(UUID(int=99))}]), + result([ROOT, ROOT]), + {"data": {"issues": {"nodes": []}}}, + ], +) +async def test_invalid_membership_keeps_visible_originals_and_saved_continuation( + database, source, invalid +): + instance, _ = await runner( + database, source, [result([ROOT]), children("comments", [CHILD], following=True), invalid] + ) + with pytest.raises((ValueError, SourceError)): + await run(instance) + rows, cursor, scans = await saved(database) + assert {row.native_id for row in rows} == {str(ISSUE), str(COMMENT)} + assert all(row.deleted_at is None and row.removal_reason is None for row in rows) + assert "canonical_checkpoint" not in cursor + child_scan = next(scan for scan in scans if scan.record_type == "comment") + assert child_scan.phase == "collecting" and child_scan.continuation["after"] == "child-next" + async with database() as db: + for row in rows: + original = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, row.id + ) + assert original.content_access == "available" and original.payload + + +@pytest.mark.parametrize( + "confirmation", + [ + result([ROOT]), + result([], following=True), + result([], errors=[{"extensions": {"type": "INTERNAL_SERVER_ERROR"}}]), + ], +) +async def test_root_omission_cannot_hide_readable_or_unconfirmed_records( + database, source, confirmation +): + first, _ = await runner( + database, + source, + [result([ROOT]), children("comments", [CHILD]), children("attachments", [])], + ) + await run(first, checkpoint=False) + second, _ = await runner(database, source, [result([]), confirmation], attempt=2) + with pytest.raises((ValueError, SourceError)): + await run(second) + rows, cursor, _ = await saved(database) + assert len(rows) == 2 and all(row.deleted_at is None for row in rows) + assert "canonical_checkpoint" not in cursor diff --git a/backend/airweave/domains/entities/canonical/tests/test_mail_http.py b/backend/airweave/domains/entities/canonical/tests/test_mail_http.py new file mode 100644 index 0000000000..01b8c8ee7f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_mail_http.py @@ -0,0 +1,192 @@ +"""Real PostgreSQL mail traversal and HTTP blob authorization over synthetic bytes.""" + +import hashlib +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from sqlalchemy import text + +from airweave.api import deps +from airweave.api.v1.endpoints.records import record_error_response, router +from airweave.db.session import get_db +from airweave.domains.entities.canonical.query import CanonicalQueryService, InvalidRecordCursor +from airweave.domains.entities.canonical.query_store import CanonicalQueryStore +from airweave.domains.entities.canonical.requests import BlobReference, CaptureBatch, RecordIdentity +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.tests.helpers import observation +from airweave.domains.storage.exceptions import StorageConnectionError, StorageNotFoundError +from airweave.models import Sync, SyncJob + + +def message(native_id, created=None, **kwargs): + return observation( + native_id, + identity=RecordIdentity(record_type="message", native_id=native_id), + payload={"id": native_id, "threadId": "same-thread"}, + source_created_at=created, + **kwargs, + ) + + +def service(): + return CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-secret") + + +async def test_thread_chronology_ties_unknown_dates_and_scope(database, source): + capture, fence = source + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + other_sync, other_job = uuid4(), uuid4() + async with database() as db: + db.add(Sync(id=other_sync, organization_id=fence.organization_id, name="Other Gmail")) + await db.flush() + db.add( + SyncJob( + id=other_job, + organization_id=fence.organization_id, + sync_id=other_sync, + status="running", + ) + ) + await db.commit() + other_fence = await capture.activate_writer( + db, fence.organization_id, other_sync, other_job, attempt_id=uuid4(), attempt_number=1 + ) + await capture.capture(db, CaptureBatch(fence=other_fence, records=(message("other", now),))) + await capture.capture( + db, + CaptureBatch( + fence=fence, + records=( + message("unknown"), + message("unknown2"), + message("b", now, completeness="partial"), + message("a", now), + ), + ), + ) + query = service() + rows, cursor = [], None + async with database() as db: + while True: + page = await query.mail_thread( + db, fence.organization_id, fence.sync_id, "same-thread", limit=1, cursor=cursor + ) + rows.extend(page.messages) + assert page.coverage == "stored_messages_only" + if not page.has_more: + break + cursor = page.next_cursor + assert {r.identity.native_id for r in rows[-2:]} == {"unknown", "unknown2"} + assert rows[-2].id < rows[-1].id + assert rows[0].id < rows[1].id + assert {r.identity.native_id for r in rows} == {"a", "b", "unknown", "unknown2"} + assert all(r.identity.container_id is None for r in rows) + assert any(r.completeness == "partial" for r in rows) + for org, sync, thread in [ + (uuid4(), fence.sync_id, "same-thread"), + (fence.organization_id, other_sync, "same-thread"), + (fence.organization_id, fence.sync_id, "different"), + ]: + with pytest.raises(InvalidRecordCursor): + await query.mail_thread(db, org, sync, thread, cursor=cursor) + other = await query.mail_thread(db, fence.organization_id, other_sync, "same-thread") + assert [r.identity.native_id for r in other.messages] == ["other"] + index = await db.scalar( + text( + "SELECT indexdef FROM pg_indexes WHERE schemaname=current_schema() " + "AND indexname='ix_entity_mail_thread'" + ) + ) + assert "threadId" in index and "source_created_at" in index + + +async def test_http_blob_scope_revision_missing_corrupt_and_revoke_during_io(database, source): + capture, fence = source + content = b"synthetic MIME body" + digest = hashlib.sha256(content).hexdigest() + ref = BlobReference( + key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type="text/plain", + ) + item = message("one", blobs=(ref,), completeness="partial") + async with database() as db: + result = await capture.capture(db, CaptureBatch(fence=fence, records=(item,))) + record_id = result.changes[0].record.id + storage = AsyncMock() + storage.read_file.return_value = content + owner = fence.organization_id + app = FastAPI() + app.include_router(router, prefix="/sync") + app.add_exception_handler(CanonicalStoreError, record_error_response) + + async def context(): + return SimpleNamespace(organization=SimpleNamespace(id=owner)) + + async def session(): + async with database() as db: + yield db + + app.dependency_overrides[deps.get_context] = context + app.dependency_overrides[get_db] = session + app.dependency_overrides[deps.get_canonical_query_service] = service + app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(storage_backend=storage) + path = f"/sync/{fence.sync_id}/records/{record_id}/blobs/{digest}" + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.get(path, params={"revision": 1}) + assert response.status_code == 200 and response.content == content + assert response.headers["cache-control"] == "private, no-store" + thread = await client.get(f"/sync/{fence.sync_id}/mail/threads/same-thread") + assert ( + thread.status_code == 200 and thread.json()["messages"][0]["completeness"] == "partial" + ) + assert (await client.get(path, params={"revision": 2})).status_code == 409 + assert ( + await client.get(path.replace(digest, "f" * 64), params={"revision": 1}) + ).status_code == 404 + assert ( + await client.get(path.replace(str(fence.sync_id), str(uuid4())), params={"revision": 1}) + ).status_code == 404 + owner = uuid4() + assert (await client.get(path, params={"revision": 1})).status_code == 404 + assert ( + await client.get(f"/sync/{fence.sync_id}/mail/threads/same-thread") + ).status_code == 404 + owner = fence.organization_id + for failure in (StorageNotFoundError("missing"), StorageConnectionError("backend down")): + storage.read_file.side_effect = failure + unavailable = await client.get(path, params={"revision": 1}) + assert unavailable.status_code == 503 + assert unavailable.json()["error"]["code"] == "blob_unavailable" + storage.read_file.side_effect = ValueError("programming error") + with pytest.raises(ValueError, match="programming error"): + await client.get(path, params={"revision": 1}) + storage.read_file.side_effect = None + storage.read_file.return_value = b"corrupt" + assert (await client.get(path, params={"revision": 1})).status_code == 503 + + async def revoke(_, *, max_bytes): + async with database() as db: + await capture.capture( + db, + CaptureBatch( + fence=fence, + records=( + item.model_copy( + update={"kind": "delete", "removal_reason": "access_revoked"} + ), + ), + ), + ) + return content + + storage.read_file.side_effect = revoke + assert (await client.get(path, params={"revision": 1})).status_code == 404 + hidden = await client.get(f"/sync/{fence.sync_id}/mail/threads/same-thread") + assert hidden.json()["messages"] == [] diff --git a/backend/airweave/domains/entities/canonical/tests/test_metadata_counts.py b/backend/airweave/domains/entities/canonical/tests/test_metadata_counts.py new file mode 100644 index 0000000000..a4c18518f2 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_metadata_counts.py @@ -0,0 +1,74 @@ +"""Metadata distinguishes owned records from independently indexed entity classes.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +from airweave.domains.entities.canonical.query_store import CanonicalQueryStore +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.tests.helpers import capture, observation +from airweave.domains.search.builders.collection_metadata import CollectionMetadataBuilder + + +async def test_captured_counts_respect_visibility_and_do_not_claim_indexed_counts(database, source): + service, fence = source + parent_id = RecordIdentity(record_type="calendar", native_id="cal") + parent = observation(identity=parent_id) + child = observation("event", "cal", parent=parent_id) + await capture(database, service, fence, parent, child) + query = CanonicalQueryStore() + async with database() as db: + assert await query.captured_counts(db, fence.organization_id, fence.sync_id) == { + "calendar": 1, + "event": 1, + } + assert await query.captured_counts(db, uuid4(), fence.sync_id) == {} + collection_repo = AsyncMock() + collection_repo.get_by_readable_id.return_value = SimpleNamespace(id=uuid4()) + sc_repo = AsyncMock() + sc_repo.get_by_collection_ids.return_value = [ + SimpleNamespace(short_name="google_calendar", sync_id=fence.sync_id, is_authenticated=True) + ] + registry = MagicMock() + registry.get.return_value = SimpleNamespace( + source_class_ref=SimpleNamespace(canonical_record_types=("calendar", "event")), + federated_search=False, + ) + definitions = MagicMock() + definitions.list_for_source.return_value = [ + SimpleNamespace( + name="GoogleCalendarEventEntity", + short_name="GoogleCalendarEventEntity", + entity_schema={"properties": {"title": {"description": "Event title"}}}, + ) + ] + legacy_counts = AsyncMock() + builder = CollectionMetadataBuilder( + collection_repo, + sc_repo, + registry, + definitions, + legacy_counts, + ) + context = SimpleNamespace(organization=SimpleNamespace(id=fence.organization_id)) + async with database() as db: + metadata = await builder.build(db, context, "calendar") + assert metadata.sources[0].captured_record_counts == {"calendar": 1, "event": 1} + assert metadata.sources[0].entity_types[0].count is None + assert "indexing may lag" in metadata.to_md() + assert "indexed count unknown" in metadata.to_md() + assert "(no entities)" not in metadata.to_md() + legacy_counts.get_counts_per_sync_and_type.assert_not_called() + await capture( + database, + service, + fence, + parent.model_copy( + update={ + "kind": "delete", + "removal_reason": "access_revoked", + } + ), + ) + async with database() as db: + assert await query.captured_counts(db, fence.organization_id, fence.sync_id) == {} diff --git a/backend/airweave/domains/entities/canonical/tests/test_mime_text.py b/backend/airweave/domains/entities/canonical/tests/test_mime_text.py new file mode 100644 index 0000000000..b651483819 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_mime_text.py @@ -0,0 +1,33 @@ +"""Charset recovery never guesses, replaces invalid bytes or overrides valid declarations.""" + +import pytest + +from airweave.domains.entities.canonical.mime_text import decode_mime_text + + +def test_valid_declared_charset_wins_even_when_bytes_also_form_valid_utf8(): + # UTF-8 for café is also valid Latin-1: the explicit declaration remains authoritative. + raw = "café".encode("utf-8") + result = decode_mime_text(raw, "iso-8859-1", "/body") + assert result.text.encode("iso-8859-1") == raw + assert result.recovery is None + + +def test_invalid_declared_charset_recovers_exact_utf8_and_records_evidence(): + raw = "नमस्ते — café".encode("utf-8") + result = decode_mime_text(raw, "us-ascii", "/parts/1") + assert result.text.encode("utf-8") == raw + assert result.recovery.model_dump() == { + "source_path": "/parts/1", + "from_charset": "us-ascii", + "to_charset": "utf-8", + } + + +@pytest.mark.parametrize( + "raw,charset,error", + [(b"\xff", "utf-8", UnicodeDecodeError), (b"valid", "unknown-charset", LookupError)], +) +def test_invalid_or_unknown_encoding_still_fails(raw, charset, error): + with pytest.raises(error): + decode_mime_text(raw, charset, "/body") diff --git a/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py new file mode 100644 index 0000000000..f048701256 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py @@ -0,0 +1,589 @@ +"""Real SQL and production page driver with a synthetic mixed-scope source.""" + +from datetime import datetime, timezone +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.coverage import capture_coverage, mixed_scope_summary +from airweave.domains.entities.canonical.cycle_models import CycleConfiguration, ProviderCheckpoint +from airweave.domains.entities.canonical.page_source import ( + CapturePage, + CapturePlan, + InvalidScopeCheckpoint, +) +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, +) +from airweave.domains.entities.canonical.scan_models import ReconcileScan, ScanContinuation +from airweave.domains.entities.canonical.scan_store import ScanConflict +from airweave.domains.entities.canonical.scope_execution import ScopePlan +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components +from airweave.domains.sources.exceptions import SourceServerError +from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline +from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt +from airweave.models.capture_scan import CaptureScan +from airweave.models.entity import Entity +from airweave.models.sync_job import SyncJob + + +class MixedSource: + canonical_record_types = ("calendar", "event", "event_occurrence") + canonical_container_parents = {"event": "calendar", "event_occurrence": "calendar"} + capture_cycle_configuration = CycleConfiguration( + fingerprint="a" * 64, + parents={"calendar": (None,), "event": ("calendar",), "event_occurrence": ("calendar",)}, + scope_changes=("event",), + ) + + def __init__(self, *, fail_window=False, expire=False, window="fixed"): + self.fail_window, self.expire, self.window = fail_window, expire, window + self.calls = [] + self.prepared = 0 + + async def prepare_cycle(self, previous): + self.prepared += 1 + return CapturePlan(mode="mixed", source_plan={"window": self.window}) + + async def prepare_scope(self, scope, cycle, previous, *, parent, force_full): + context = ( + {"parameters": "masters"} + if scope.record_type == "event" + else {"window": cycle.source_plan["window"]} + ) + published = previous.execution.published if previous and previous.execution else None + changes = ( + scope.record_type == "event" and published and published.checkpoint and not force_full + ) + return ScopePlan( + mode="changes" if changes else "full", + request_context=context, + starting_checkpoint=published.checkpoint if changes else None, + ) + + def initial_scope_continuation(self, scope, cycle, plan): + return ScanContinuation( + value={ + "mode": plan.mode, + "token": plan.starting_checkpoint.value if plan.starting_checkpoint else None, + "window": cycle.source_plan["window"], + } + ) + + def child_scope(self, parent, record_type): + return CompletedScope( + record_type=record_type, container_id=parent.identity.native_id, parent=parent.identity + ) + + async def confirm_absent(self, record): + raise AssertionError("Synthetic membership is unchanged") + + async def capture_page(self, scope, continuation, *, files, parent=None): + self.calls.append((scope.record_type, continuation.value)) + if scope.record_type == "event_occurrence" and self.fail_window: + raise SourceServerError("Synthetic interruption", status_code=503) + if scope.record_type == "event" and continuation.value["mode"] == "changes" and self.expire: + self.expire = False + raise InvalidScopeCheckpoint("Synthetic410") + identities = ( + ("cal",) + if scope.record_type == "calendar" + else ("a",) + if continuation.value["mode"] == "changes" + else ("a", "b") + ) + records = tuple( + CaptureRecord( + identity=RecordIdentity( + record_type=scope.record_type, native_id=native, container_id=scope.container_id + ), + payload={"id": native, "mode": continuation.value["mode"]}, + observed_at=datetime.now(timezone.utc), + ) + for native in identities + ) + return CapturePage( + records=records, + continuation=continuation, + final=True, + provider_checkpoint=ProviderCheckpoint(value={"sync_token": str(len(self.calls))}) + if scope.record_type == "event" + else None, + ) + + +async def setup(database, source, connector, attempt=1): + service, fence = source + ctx, _, runtime, bus = components(database, source) + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + connector.canonical_record_types, + CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), + connector.canonical_container_parents, + page_source=connector, + files=MagicMock(), + ) + await pipeline.start(ctx) + return pipeline, ctx, runtime + + +async def no_limits(): + pass + + +async def run(pipeline, ctx, runtime): + await pipeline.run_scans(ctx, runtime, no_limits) + await pipeline.cleanup_orphaned_entities(ctx, runtime) + await pipeline.save_checkpoint(ctx, runtime) + + +async def next_job(database, source): + service, fence = source + async with database() as db: + (await db.get(SyncJob, fence.job_id)).status = "completed" + job = SyncJob( + id=uuid4(), + sync_id=fence.sync_id, + organization_id=fence.organization_id, + status="running", + ) + db.add(job) + await db.commit() + return service, await service.activate_writer( + db, fence.organization_id, fence.sync_id, job.id, attempt_id=uuid4(), attempt_number=1 + ) + + +async def test_mixed_scope_resume_preserves_baseline_and_fixed_window(database, source): + connector = MixedSource(fail_window=True) + pipeline, ctx, runtime = await setup(database, source, connector) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + event = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + old_full = event.execution_state["last_full"] + assert event.phase == "complete" and old_full is not None + coverage = (await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)))[ + source[1].sync_id + ] + assert coverage.mode == "mixed" and coverage.scope_summary.unfinished == 1 + resumed = MixedSource(window="must-not-replace-active-window") + pipeline, ctx, runtime = await setup(database, source, resumed, 2) + await run(pipeline, ctx, runtime) + assert resumed.prepared == 0 + assert [kind for kind, _ in resumed.calls] == ["calendar", "event_occurrence"] + assert all(progress["window"] == "fixed" for _, progress in resumed.calls) + source = await next_job(database, source) + delta = MixedSource(window="next-window") + pipeline, ctx, runtime = await setup(database, source, delta) + await run(pipeline, ctx, runtime) + assert [progress["mode"] for kind, progress in delta.calls if kind == "event"] == ["changes"] + async with database() as db: + event = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + assert event.execution_state["last_full"] == old_full + assert ( + await db.scalar( + select(Entity).where( + Entity.entity_definition_short_name == "event", Entity.native_id == "b" + ) + ) + ).deleted_at is None + coverage = (await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)))[ + source[1].sync_id + ] + assert coverage.scope_summary.model_dump() == { + "eligible": 3, + "completed_full": 2, + "completed_changes": 1, + "unfinished": 0, + } + assert ( + coverage.last_full_capture is None and coverage.provider_checkpoint_promoted_at is None + ) + assert coverage.discovery == "scope_enumeration_complete" + + +async def test_scope_expiry_restarts_only_event_inventory(database, source): + first = MixedSource() + pipeline, ctx, runtime = await setup(database, source, first) + await run(pipeline, ctx, runtime) + source = await next_job(database, source) + connector = MixedSource(expire=True) + pipeline, ctx, runtime = await setup(database, source, connector) + await run(pipeline, ctx, runtime) + assert [progress["mode"] for kind, progress in connector.calls if kind == "event"] == [ + "changes", + "full", + ] + assert connector.prepared == 1 + assert sum(kind == "calendar" for kind, _ in connector.calls) == 1 + + +async def test_changes_scope_rejects_absence_even_before_cycle_completion(database, source): + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + source = await next_job(database, source) + pipeline, ctx, runtime = await setup(database, source, MixedSource(fail_window=True)) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + state = await source[0].read_scan( + db, pipeline._writer(), CompletedScope(record_type="event", container_id="cal") + ) + assert state.mode == "changes" + with pytest.raises(ScanConflict, match="cannot reconcile"): + await source[0].reconcile_scan( + db, + ReconcileScan( + fence=pipeline._writer(), + scope=state.scope, + cycle_id=state.cycle_id, + expected=state.version, + observed_at=datetime.now(timezone.utc), + ), + ) + + +async def test_changed_capture_configuration_cannot_reuse_same_context_baseline(database, source): + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + source = await next_job(database, source) + connector = MixedSource() + connector.capture_cycle_configuration = connector.capture_cycle_configuration.model_copy( + update={"fingerprint": "b" * 64} + ) + pipeline, ctx, runtime = await setup(database, source, connector) + # Synthetic adapter deliberately proposes an old token despite configuration change. + with pytest.raises(ScanConflict, match="compatible full evidence"): + await run(pipeline, ctx, runtime) + async with database() as db: + row = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + assert row.fingerprint == "a" * 64 + + +async def test_forced_mixed_cycle_resumes_intent_across_attempts_and_jobs(database, source): + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + source = await next_job(database, source) + pipeline, ctx, runtime = await setup(database, source, MixedSource(fail_window=True)) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + forced = MixedSource(fail_window=True, window="forced-window") + pipeline, ctx, runtime = await setup(database, source, forced, 2) + ctx.force_full_sync = True + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + assert [p["mode"] for k, p in forced.calls if k == "event"] == ["full"] + async with database() as db: + initial = await source[0].read_cycle(db, pipeline._writer()) + assert initial.force_full_scopes + repeated = MixedSource(fail_window=True, window="wrong") + pipeline, ctx, runtime = await setup(database, source, repeated, 3) + ctx.force_full_sync = True + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + assert repeated.prepared == 0 and not any(k == "event" for k, _ in repeated.calls) + source = await next_job(database, source) + resumed = MixedSource(window="also-wrong") + pipeline, ctx, runtime = await setup(database, source, resumed) + await run(pipeline, ctx, runtime) + async with database() as db: + completed = await source[0].read_cycle(db, pipeline._writer()) + assert completed.version.cycle_id == initial.version.cycle_id + assert completed.source_plan == {"window": "forced-window"} + assert resumed.prepared == 0 + + +async def test_mixed_coverage_wire_from_sql(database, source, tmp_path): + import json + + from airweave.domains.entities.canonical.cycle_models import BeginCycle + + connector = MixedSource(fail_window=True) + pipeline, ctx, runtime = await setup(database, source, connector) + async with database() as db: + await source[0].begin_cycle( + db, + BeginCycle( + fence=pipeline._writer(), + configuration=connector.capture_cycle_configuration, + mode="mixed", + source_plan={"window": "fixed"}, + ), + ) + unknown = (await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)))[ + source[1].sync_id + ] + assert unknown.scope_summary is None and unknown.discovery == "pending" + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + partial = (await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)))[ + source[1].sync_id + ] + assert partial.scope_summary.unfinished == 1 + pipeline, ctx, runtime = await setup(database, source, MixedSource(), 2) + await run(pipeline, ctx, runtime) + async with database() as db: + complete = (await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)))[ + source[1].sync_id + ] + assert complete.scope_summary.unfinished == 0 and complete.phase == "complete" + assert complete.scope_summary.completed_changes == 0 + source = await next_job(database, source) + delta = MixedSource(fail_window=True, window="next-window") + pipeline, ctx, runtime = await setup(database, source, delta) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + assert [value["mode"] for kind, value in delta.calls if kind == "event"] == ["changes"] + async with database() as db: + partial_delta = ( + await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)) + )[source[1].sync_id] + assert partial_delta.phase == "active" and partial_delta.discovery == "pending" + assert partial_delta.scope_summary.model_dump() == { + "eligible": 3, + "completed_full": 1, + "completed_changes": 1, + "unfinished": 1, + } + resumed = MixedSource(window="must-preserve-next-window") + pipeline, ctx, runtime = await setup(database, source, resumed, 2) + await run(pipeline, ctx, runtime) + assert [kind for kind, _ in resumed.calls] == ["calendar", "event_occurrence"] + async with database() as db: + complete_delta = ( + await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)) + )[source[1].sync_id] + assert complete_delta.phase == "complete" + assert complete_delta.discovery == "scope_enumeration_complete" + assert complete_delta.scope_summary.model_dump() == { + "eligible": 3, + "completed_full": 2, + "completed_changes": 1, + "unfinished": 0, + } + assert complete_delta.last_full_capture is None + assert complete_delta.provider_checkpoint_promoted_at is None + unchanged = await db.scalar( + select(Entity).where( + Entity.entity_definition_short_name == "event", Entity.native_id == "b" + ) + ) + assert unchanged.deleted_at is None + (tmp_path / "mixed-coverage.json").write_text( + json.dumps( + [ + item.model_dump(mode="json") + for item in (unknown, partial, complete, partial_delta, complete_delta) + ], + indent=2, + ) + ) + + +async def test_parent_revival_invalidates_published_scope_history(database, source): + from airweave.domains.entities.canonical.requests import CaptureBatch + + pipeline, ctx, runtime = await setup(database, source, MixedSource(fail_window=True)) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + old = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + epoch = old.execution_state["last_full"]["parent_visibility_epoch"] + await source[0].capture( + db, + CaptureBatch( + fence=pipeline._writer(), + records=( + CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id="cal"), + payload={"id": "cal"}, + kind="delete", + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ), + ), + ), + ) + await source[0].capture( + db, + CaptureBatch( + fence=pipeline._writer(), + records=( + CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id="cal"), + payload={"id": "cal"}, + observed_at=datetime.now(timezone.utc), + ), + ), + ), + ) + restored = MixedSource() + pipeline, ctx, runtime = await setup(database, source, restored, 2) + await run(pipeline, ctx, runtime) + assert [p["mode"] for k, p in restored.calls if k == "event"] == ["full"] + async with database() as db: + row = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + assert row.execution_state["last_full"]["parent_visibility_epoch"] > epoch + + +@pytest.mark.parametrize("race", ["owner", "writer"]) +async def test_awaited_scope_plan_cannot_commit_after_owner_or_writer_change( + database, source, race +): + from airweave.domains.entities.canonical.requests import CaptureBatch + from airweave.domains.entities.canonical.store import StaleWriter + + service, fence = source + + class RacingSource(MixedSource): + async def prepare_scope(self, scope, cycle, previous, *, parent, force_full): + plan = await super().prepare_scope( + scope, cycle, previous, parent=parent, force_full=force_full + ) + if scope.record_type == "event": + async with database() as db: + if race == "writer": + await service.activate_writer( + db, + fence.organization_id, + fence.sync_id, + fence.job_id, + attempt_id=uuid4(), + attempt_number=2, + ) + else: + await service.capture( + db, + CaptureBatch( + fence=fence, + records=( + CaptureRecord( + identity=parent.identity, + payload={"id": "cal", "route": "changed"}, + observed_at=datetime.now(timezone.utc), + ), + ), + ), + ) + return plan + + pipeline, ctx, runtime = await setup(database, source, RacingSource()) + with pytest.raises(StaleWriter if race == "writer" else ScanConflict): + await run(pipeline, ctx, runtime) + async with database() as db: + assert ( + await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) is None + ) + assert ( + await db.scalar(select(Entity).where(Entity.entity_definition_short_name == "event")) + is None + ) + + +async def test_nonmixed_replacement_clears_prior_scope_evidence(database, source): + from airweave.domains.entities.canonical.cycle_models import BeginCycle + from airweave.domains.entities.canonical.scan_models import BeginScan + + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + async with database() as db: + service = source[0] + fence = pipeline._writer() + previous = await service.read_cycle(db, fence) + config = previous.configuration.model_copy( + update={"fingerprint": "b" * 64, "scope_changes": ()} + ) + current = await service.begin_cycle( + db, BeginCycle(fence=fence, configuration=config, expected=previous.version) + ) + scope = CompletedScope(record_type="calendar") + old = await service.read_scan(db, fence, scope) + assert old.execution is not None + new = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=scope, + cycle_id=current.version.cycle_id, + fingerprint=config.fingerprint, + expected=old.version, + ), + ) + assert new.execution is None + + +async def test_completed_coverage_survives_job_completion_but_not_new_writer(database, source): + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + service, fence = source + async with database() as db: + (await db.get(SyncJob, fence.job_id)).status = "completed" + await db.commit() + coverage = await capture_coverage(db, fence.organization_id, (fence.sync_id,)) + assert coverage[fence.sync_id].phase == "complete" + assert coverage[fence.sync_id].scope_summary.unfinished == 0 + source = await next_job(database, source) + async with database() as db: + coverage = await capture_coverage(db, fence.organization_id, (fence.sync_id,)) + assert fence.sync_id not in coverage + pipeline, ctx, runtime = await setup(database, source, MixedSource(fail_window=True)) + with pytest.raises(SourceServerError): + await run(pipeline, ctx, runtime) + async with database() as db: + coverage = await capture_coverage(db, fence.organization_id, (fence.sync_id,)) + assert coverage[fence.sync_id].phase == "active" + assert coverage[fence.sync_id].scope_summary.unfinished == 1 + + +@pytest.mark.parametrize( + "change", ["stale_root_inventory", "unknown_root_mode", "parent_epoch", "leaf_attempt"] +) +async def test_coverage_root_batch_preserves_scope_evidence(database, source, change): + pipeline, ctx, runtime = await setup(database, source, MixedSource()) + await run(pipeline, ctx, runtime) + service, fence = source + async with database() as db: + cycle = await service.read_cycle(db, fence) + root = await db.scalar( + select(CaptureScan).where( + CaptureScan.sync_id == fence.sync_id, CaptureScan.record_type == "calendar" + ) + ) + if change == "stale_root_inventory": + root.membership_attempt_id = uuid4() + elif change == "unknown_root_mode": + root.execution_state = None + elif change == "parent_epoch": + parent = await db.scalar( + select(Entity).where( + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == "calendar", + ) + ) + parent.visibility_epoch += 1 + else: + leaf = await db.scalar( + select(CaptureScan).where( + CaptureScan.sync_id == fence.sync_id, CaptureScan.record_type == "event" + ) + ) + leaf.membership_attempt_id = uuid4() + await db.flush() + + summary = await mixed_scope_summary(db, fence.organization_id, fence.sync_id, cycle) + if change == "stale_root_inventory": + assert summary is None + else: + expected_unfinished = {"unknown_root_mode": 1, "parent_epoch": 2, "leaf_attempt": 0} + assert summary.eligible == 3 + assert summary.unfinished == expected_unfinished[change] + assert summary.completed_full == 3 - summary.unfinished + assert await mixed_scope_summary(db, uuid4(), fence.sync_id, cycle) is None diff --git a/backend/airweave/domains/entities/canonical/tests/test_native_text.py b/backend/airweave/domains/entities/canonical/tests/test_native_text.py new file mode 100644 index 0000000000..f35b9e239f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_native_text.py @@ -0,0 +1,163 @@ +"""Source-selected native bodies share the retained and indexed construction.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from pydantic import ValidationError + +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.tests.test_slack_projection import message +from airweave.domains.entities.canonical.text_models import TextArtifact +from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder + + +@pytest.mark.parametrize("schema", [1, 2]) +@pytest.mark.parametrize( + "text", ["", " \n", "Hello <@U1> 👋\n*bold*नमस्ते café & body
"), + ], +) +async def test_body_provenance_unicode_schedule_and_partial_inventory( + tmp_path, monkeypatch, body_type, text +): + # Even an accidental future native HTTP request would fail this offline test. + monkeypatch.setattr( + "httpx.AsyncClient.request", AsyncMock(side_effect=AssertionError("offline")) + ) + record = event({"contentType": body_type, "content": text}) + before = record.model_dump() + mapped, coverage, batch = await build(record, tmp_path) + built = batch.representations[0] + assert built.kind == ("native_text" if body_type == "text" else "extracted_text") + assert built.text == batch.entities[0].textual_representation + body = built.text[built.content_start :] + if body_type == "text": + assert body == text + else: + assert "नमस्ते" in body and "café" in body and "&" in body + assert "" not in body and "" not in body
+ assert mapped.parts[0].native_body is None
+ assert "2026-11-01T01:30:00.0000000 (Pacific Standard Time)" in built.text
+ assert "01:30:00+00:00" not in built.text
+ assert "tzone://Microsoft/Custom" in built.text
+ assert "**Is Cancelled**: True" in built.text
+ assert all(value in built.text for value in ("Éva", "accepted", "weekly", "monday", "Office"))
+ assert "DO NOT INDEX" not in built.text and "LOSSY PREVIEW" not in built.text
+ assert coverage.status == "partial"
+ assert [(part.key, part.outcome) for part in coverage.parts[1:]] == [
+ ("/attachment_inventory", "unavailable_original"),
+ ("/cancelledOccurrences_inventory", "unavailable_original"),
+ ]
+ assert record.model_dump() == before
+
+
+async def test_missing_body_is_not_preview_and_calendar_unknown_permissions_stay_unknown(tmp_path):
+ record = event(None)
+ _, coverage, batch = await build(record, tmp_path)
+ assert batch.representations[0].kind == "generated_text"
+ assert any(
+ part.key == "/body/content" and part.outcome == "unavailable_original"
+ for part in coverage.parts
+ )
+ calendar = record.model_copy(
+ update={
+ "identity": record.parent,
+ "parent": None,
+ "completeness": "complete",
+ "payload": {
+ "id": "calendar-1",
+ "name": "Équipe",
+ "owner": {"name": "Owner", "address": "o@x.test"},
+ "unknown": "DO NOT INDEX",
+ },
+ }
+ )
+ mapped = await map_outlook_calendar(calendar, tmp_path)
+ assert mapped.entities[0].owner == "Owner नमस्ते café, पूरा पाठनमस्ते 原文 नमस्ते 原文", "… नमस्ते 原文"),
+ ("semantic", "… नमस्ते 原文", None),
+ ("hybrid", "An unmatched lead without term markers", None),
+ ("keyword", None, None),
+ ],
+)
+async def test_owned_lexical_fragment_keeps_original_chunk(
+ database, indexed, http_search, mode, snippet, expected_fragment
+):
+ fence, locator, _ = indexed
+ client, vector, _, _, _ = http_search
+ http_search[3]._dense_embedder = FakeDenseEmbedder(dimensions=3)
+ original = "Leading context. " * 200 + "नमस्ते 原文"
+ candidate = hit(fence, locator.encode()).model_copy(
+ update={
+ "textual_representation": original,
+ "query_snippet": snippet,
+ }
+ )
+ vector.seed_results(SearchResults(results=[candidate]))
+ response = await client.post(
+ "/sync/search",
+ json={
+ "query": "नमस्ते",
+ "sync_ids": [str(fence.sync_id)],
+ "mode": mode,
+ },
+ )
+ assert response.status_code == 200, response.text
+ expected = expected_fragment if expected_fragment is not None else original[:2000]
+ assert response.json()["items"][0]["excerpts"] == [expected]
+ assert response.json()["items"][0]["record_id"] == str(locator.record_id)
+ assert candidate.textual_representation == original
diff --git a/backend/airweave/domains/entities/canonical/tests/test_page_child_receipts.py b/backend/airweave/domains/entities/canonical/tests/test_page_child_receipts.py
new file mode 100644
index 0000000000..bc779b28f8
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_page_child_receipts.py
@@ -0,0 +1,320 @@
+"""Native page observations atomically attest existing child scans without another queue."""
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import func, select
+
+from airweave.domains.entities.canonical.cycle_models import BeginCycle
+from airweave.domains.entities.canonical.cycle_store import CycleConflict
+from airweave.domains.entities.canonical.requests import CompletedScope, parent_container_key
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ ChildScopeObservation,
+ CommitScanPage,
+ ScanContinuation,
+)
+from airweave.domains.entities.canonical.scan_store import ScanConflict
+from airweave.domains.entities.canonical.store import StaleWriter
+from airweave.domains.entities.canonical.tests.helpers import capture
+from airweave.domains.entities.canonical.tests.test_cycles import finish, page, scan
+from airweave.domains.entities.canonical.tests.test_exact_parent_validation import CONFIG
+from airweave.domains.entities.canonical.tests.test_forest_scans import collect, item
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+
+
+async def history(database, source, config=CONFIG):
+ service, fence = source
+ async with database() as db:
+ cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config))
+ channel = item("channel", "C")
+ await collect(database, service, fence, cycle, CompletedScope(record_type="channel"), channel)
+ messages = tuple(item("message", name, channel.identity) for name in ("one", "two"))
+ state = await scan(
+ database,
+ service,
+ fence,
+ cycle,
+ CompletedScope(
+ record_type="message",
+ container_id=messages[0].identity.container_id,
+ parent=channel.identity,
+ ),
+ )
+ return cycle, channel, messages, state
+
+
+def declaration(message):
+ return ChildScopeObservation(
+ scope=CompletedScope(
+ record_type="file",
+ container_id=parent_container_key(message.identity),
+ parent=message.identity,
+ )
+ )
+
+
+def request(fence, state, messages, declarations=None):
+ return CommitScanPage(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ records=messages,
+ continuation=ScanContinuation(value={"next": str(state.version.revision)}),
+ child_scope_observations=(
+ tuple(declaration(m) for m in messages) if declarations is None else declarations
+ ),
+ )
+
+
+async def commit(database, service, request):
+ async with database() as db:
+ return await service.commit_scan_page(db, request)
+
+
+@pytest.mark.parametrize("completed", [False, True])
+async def test_page_receipts_include_unchanged_owners_and_preserve_child_progress(
+ database, source, completed
+):
+ service, fence = source
+ _, _, messages, state = await history(database, source)
+ # Complete child inventory does not imply complete unrelated body coverage.
+ messages = (messages[0].model_copy(update={"completeness": "partial"}), messages[1])
+ async with database() as db:
+ before_cycle = await service.read_cycle(db, fence)
+ result = await commit(database, service, request(fence, state, messages))
+ async with database() as db:
+ cycle = await service.read_cycle(db, fence)
+ children = [await service.read_scan(db, fence, declaration(m).scope) for m in messages]
+ assert cycle.version.revision == before_cycle.version.revision + len(messages)
+ assert all(c.parent_verified_attempt_id == fence.attempt_id for c in children)
+ child = await page(
+ database,
+ service,
+ fence,
+ children[0],
+ item("file", "F", messages[0].identity),
+ final=completed,
+ )
+ if completed:
+ child = await finish(database, service, fence, child)
+ repeated = await commit(database, service, request(fence, result.state, messages))
+ assert repeated.capture.unchanged == 2 and repeated.capture.changes == ()
+ async with database() as db:
+ current = await service.read_scan(db, fence, child.scope)
+ owner = await db.scalar(select(Entity).where(Entity.native_id == "one"))
+ assert current.version.sweep_id == child.version.sweep_id
+ assert current.continuation == child.continuation and current.phase == child.phase
+ assert current.parent_verified_revision == owner.record_revision == 1
+ assert current.parent_visibility_epoch == owner.visibility_epoch
+
+
+async def test_changed_inventory_restarts_child_and_hides_removed_attachment(database, source):
+ service, fence = source
+ _, _, messages, state = await history(database, source)
+ message = messages[0].model_copy(
+ update={"payload": {"files": ["F"]}, "descendant_visibility_fields": ("files",)}
+ )
+ result = await commit(database, service, request(fence, state, (message,)))
+ async with database() as db:
+ child = await service.read_scan(db, fence, declaration(message).scope)
+ child = await page(
+ database, service, fence, child, item("file", "F", message.identity), final=True
+ )
+ child = await finish(database, service, fence, child)
+ changed = message.model_copy(update={"payload": {"files": []}})
+ await commit(database, service, request(fence, result.state, (changed,)))
+ async with database() as db:
+ current = await service.read_scan(db, fence, child.scope)
+ file = await db.scalar(select(Entity).where(Entity.native_id == "F"))
+ retained = await service.store.read(db, fence.organization_id, fence.sync_id, file.id)
+ assert retained.content_access == "unavailable"
+ assert current.phase == "collecting" and current.continuation == ScanContinuation()
+ assert current.version.sweep_id != child.version.sweep_id
+ assert current.parent_verified_revision == 2
+ empty = await page(database, service, fence, current, final=True)
+ await finish(database, service, fence, empty)
+ async with database() as db:
+ assert (await db.get(Entity, file.id)).deleted_at is not None
+
+
+@pytest.mark.parametrize(
+ "invalid",
+ [
+ "outside_page",
+ "duplicate",
+ "duplicate_namespace",
+ "delete",
+ "reparent",
+ "relationship",
+ "not_opted_in",
+ ],
+)
+async def test_page_rejects_unqualified_child_receipts_without_advancing(database, source, invalid):
+ service, fence = source
+ config = (
+ CONFIG.model_copy(update={"exact_parent_validation": ()})
+ if invalid == "not_opted_in"
+ else CONFIG
+ )
+ _, _, messages, state = await history(database, source, config)
+ message = messages[0]
+ declared = declaration(message)
+ if invalid == "outside_page":
+ await capture(database, service, fence, messages[1])
+ declared = declaration(messages[1])
+ elif invalid == "delete":
+ message = message.model_copy(
+ update={"kind": "delete", "removal_reason": "provider_deleted"}
+ )
+ elif invalid == "reparent":
+ message = message.model_copy(update={"allow_reparent": True})
+ elif invalid == "relationship":
+ declared = declared.model_copy(
+ update={"scope": declared.scope.model_copy(update={"record_type": "channel"})}
+ )
+ declarations = (declared, declared) if invalid == "duplicate" else (declared,)
+ if invalid == "duplicate_namespace":
+ declarations = (
+ declared,
+ declared.model_copy(
+ update={"scope": declared.scope.model_copy(update={"container_id": "other"})}
+ ),
+ )
+ with pytest.raises(ScanConflict):
+ await commit(database, service, request(fence, state, (message,), declarations))
+ async with database() as db:
+ assert (await service.read_scan(db, fence, state.scope)).version == state.version
+ assert await db.scalar(
+ select(func.count())
+ .select_from(Entity)
+ .where(Entity.entity_definition_short_name == "message")
+ ) == (1 if invalid == "outside_page" else 0)
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(CaptureScan)
+ .where(CaptureScan.record_type == "file")
+ )
+ == 0
+ )
+
+
+async def test_receipt_failure_rolls_back_originals_scans_and_page_cursor(
+ database, source, monkeypatch
+):
+ service, fence = source
+ _, _, messages, state = await history(database, source)
+ original = service.scans._begin_verified
+ calls = 0
+
+ async def fail_second(*args, **kwargs):
+ nonlocal calls
+ calls += 1
+ value = await original(*args, **kwargs)
+ if calls == 2:
+ raise RuntimeError("synthetic crash after second receipt")
+ return value
+
+ monkeypatch.setattr(service.scans, "_begin_verified", fail_second)
+ with pytest.raises(RuntimeError, match="synthetic crash"):
+ await commit(database, service, request(fence, state, messages))
+ async with database() as db:
+ assert (await service.read_scan(db, fence, state.scope)).version == state.version
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(Entity)
+ .where(Entity.entity_definition_short_name == "message")
+ )
+ == 0
+ )
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(CaptureScan)
+ .where(CaptureScan.record_type == "file")
+ )
+ == 0
+ )
+
+
+async def test_new_writer_needs_new_observation_and_rejects_old_writer(database, source):
+ service, fence = source
+ cycle, channel, messages, state = await history(database, source)
+ result = await commit(database, service, request(fence, state, messages))
+ async with database() as db:
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ root = await service.read_scan(db, newer, CompletedScope(record_type="channel"))
+ with pytest.raises(StaleWriter):
+ await commit(database, service, request(fence, result.state, messages))
+ root = await scan(
+ database, service, newer, cycle, root.scope, expected=root.version, restart=True
+ )
+ root = await page(database, service, newer, root, channel, final=True)
+ await finish(database, service, newer, root)
+ async with database() as db:
+ old_child = await service.read_scan(db, newer, declaration(messages[1]).scope)
+ old_admission = BeginScan(
+ fence=newer,
+ scope=old_child.scope,
+ cycle_id=cycle.version.cycle_id,
+ fingerprint=CONFIG.fingerprint,
+ expected=old_child.version,
+ expected_parent_epoch=old_child.parent_visibility_epoch,
+ )
+ with pytest.raises(CycleConflict, match="fresh exact"):
+ async with database() as db:
+ await service.begin_scan(db, old_admission)
+ await commit(database, service, request(newer, result.state, (messages[0],)))
+ async with database() as db:
+ fresh = await service.read_scan(db, newer, declaration(messages[0]).scope)
+ old = await service.read_scan(db, newer, old_child.scope)
+ assert fresh.parent_verified_attempt_id == newer.attempt_id
+ assert old.parent_verified_attempt_id == fence.attempt_id
+
+
+async def test_withdrawn_ancestor_rejects_page_receipts(database, source):
+ service, fence = source
+ _, channel, messages, state = await history(database, source)
+ await capture(
+ database,
+ service,
+ fence,
+ channel.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ with pytest.raises(CycleConflict, match="visible parent"):
+ await commit(database, service, request(fence, state, messages))
+ async with database() as db:
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(CaptureScan)
+ .where(CaptureScan.record_type == "file")
+ )
+ == 0
+ )
+
+
+async def test_established_child_namespace_change_rolls_back_page(database, source):
+ service, fence = source
+ _, _, messages, state = await history(database, source)
+ result = await commit(database, service, request(fence, state, messages))
+ declared = declaration(messages[0])
+ wrong = declared.model_copy(
+ update={"scope": declared.scope.model_copy(update={"container_id": "other"})}
+ )
+ with pytest.raises(CycleConflict, match="identity mapping changed"):
+ await commit(database, service, request(fence, result.state, (messages[0],), (wrong,)))
+ async with database() as db:
+ assert (await service.read_scan(db, fence, state.scope)).version == result.state.version
+ assert (await service.read_scan(db, fence, declared.scope)).scope == declared.scope
diff --git a/backend/airweave/domains/entities/canonical/tests/test_parent_visibility.py b/backend/airweave/domains/entities/canonical/tests/test_parent_visibility.py
new file mode 100644
index 0000000000..c8de7ee6de
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_parent_visibility.py
@@ -0,0 +1,84 @@
+"""Real database proof of immediate visibility and retryable descendant removals."""
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.query_models import RecordFilters
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.models import Entity, EntityChange
+
+pytestmark = pytest.mark.integration
+
+
+async def test_parent_deletion_hides_content_before_sweep_and_replay_recovers(database, source):
+ service, fence = source
+ parent_id = RecordIdentity(record_type="calendar", native_id="cal")
+ parent = observation(identity=parent_id)
+ child = observation(
+ "event",
+ "cal",
+ parent=parent_id,
+ blobs=(BlobReference(key="immutable", sha256="a" * 64, size_bytes=7),),
+ )
+ initial = await capture(database, service, fence, parent, child)
+ child_id = initial.changes[1].record.id
+ tombstone = parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"})
+ await capture(database, service, fence, tombstone)
+
+ # Simulate a crash after the parent transaction, before any child cleanup.
+ async with database() as db:
+ persisted = await db.get(Entity, child_id)
+ assert persisted.deleted_at is None
+ exact = await service.store.read(db, fence.organization_id, fence.sync_id, child_id)
+ assert exact.content_access == "unavailable"
+ assert exact.payload == {} and exact.blobs == ()
+ visible = await CanonicalQueryStore().list_records(
+ db, fence.organization_id, fence.sync_id, RecordFilters(), after_id=None, limit=100
+ )
+ assert not visible
+ history = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert all(change.record.content_access == "unavailable" for change in history.changes)
+ assert all(change.record.payload == {} for change in history.changes)
+ snapshot = await db.scalar(select(EntityChange.snapshot).where(EntityChange.sequence == 2))
+ assert snapshot["payload"]["summary"] == "Synthetic event"
+ assert snapshot["blobs"][0]["key"] == "immutable"
+
+ replay = await capture(database, service, fence, tombstone)
+ assert replay.unchanged == 1
+ async with database() as db:
+ swept = await service.reconcile_parents(db, fence)
+ assert len(swept.capture.changes) == 1
+ assert swept.capture.changes[0].record.id == child_id
+ assert swept.capture.changes[0].record.removal_reason == "access_revoked"
+ async with database() as db:
+ repeated = await service.reconcile_parents(db, fence)
+ assert not repeated.capture.changes
+ async with database() as db:
+ history = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert [item.sequence for item in history.changes] == [1, 2, 3, 4]
+ assert history.changes[-1].kind == "delete"
+ assert history.changes[-1].record.content_access == "unavailable"
+
+
+@pytest.mark.parametrize("reason", ["access_revoked", "scope_removed"])
+async def test_revoked_root_tombstones_redact_retained_payload(database, source, reason):
+ service, fence = source
+ record = observation(kind="delete", removal_reason=reason)
+ result = await capture(database, service, fence, record)
+ async with database() as db:
+ exact = await service.store.read(
+ db, fence.organization_id, fence.sync_id, result.changes[0].record.id
+ )
+ assert exact.content_access == "unavailable" and exact.payload == {}
+ listed = await CanonicalQueryStore().list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordFilters(state="deleted"),
+ after_id=None,
+ limit=100,
+ )
+ assert len(listed) == 1
+ assert listed[0].content_access == "unavailable" and listed[0].payload == {}
diff --git a/backend/airweave/domains/entities/canonical/tests/test_postgres.py b/backend/airweave/domains/entities/canonical/tests/test_postgres.py
new file mode 100644
index 0000000000..72eb3062cc
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_postgres.py
@@ -0,0 +1,344 @@
+"""Real PostgreSQL proofs; never connect without an explicit synthetic-test URL."""
+
+import asyncio
+from datetime import datetime, timedelta, timezone
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import delete, func, select, text, update
+
+from airweave.domains.entities.canonical.requests import (
+ CaptureBatch,
+ CompletedScope,
+ ReconcileScope,
+)
+from airweave.domains.entities.canonical.store import (
+ SourceNotFound,
+ StaleWriter,
+ WriterBusy,
+)
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.models import Entity, EntityChange, Sync, SyncCursor, SyncJob
+
+pytestmark = pytest.mark.integration
+
+
+async def test_identity_replay_metadata_tombstone_and_historical_snapshots(database, source):
+ service, fence = source
+ first = observation(container_id="cal-a")
+ second = observation(container_id="cal-b")
+ result = await capture(database, service, fence, first, second)
+ assert [c.sequence for c in result.changes] == [1, 2]
+ assert result.changes[0].record.id != result.changes[1].record.id
+ replay = await capture(
+ database,
+ service,
+ fence,
+ first.model_copy(update={"observed_at": first.observed_at + timedelta(hours=1)}),
+ )
+ assert replay.unchanged == 1 and replay.sequence == 2 and not replay.changes
+ changed = first.model_copy(update={"payload": {"id": "one", "labels": ["IMPORTANT"]}})
+ changed_result = await capture(database, service, fence, changed)
+ assert changed_result.changes[0].record.revision == 2
+ assert changed_result.changes[0].record.content_hash == "same-body"
+ cancelled = observation(
+ "never-seen",
+ "cal-a",
+ payload={
+ "id": "never-seen",
+ "recurringEventId": "series",
+ "originalStartTime": {"date": "2026-09-30"},
+ },
+ kind="delete",
+ removal_reason="provider_deleted",
+ )
+ deleted = await capture(database, service, fence, cancelled)
+ assert deleted.changes[0].record.payload["recurringEventId"] == "series"
+ async with database() as db:
+ page = await service.store.changes(db, fence.organization_id, fence.sync_id, limit=2)
+ assert page.has_more and page.next_sequence == 2
+ assert page.changes[0].record.payload["summary"] == "Synthetic event"
+ tail = await service.store.changes(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ after=page.next_sequence,
+ high_watermark=page.high_watermark,
+ )
+ assert [c.sequence for c in tail.changes] == [3, 4]
+ assert not tail.has_more
+ assert (
+ await service.store.read(db, uuid4(), fence.sync_id, result.changes[0].record.id)
+ is None
+ )
+ with pytest.raises(SourceNotFound):
+ await service.store.changes(db, uuid4(), fence.sync_id)
+ pending = await db.scalar(
+ select(func.count())
+ .select_from(Entity)
+ .where(Entity.indexed_revision.is_distinct_from(Entity.record_revision))
+ )
+ assert pending == 3
+
+
+async def test_rollback_has_no_record_journal_or_sequence_hole(database, source):
+ service, fence = source
+ async with database() as db:
+ await service.store.capture(db, CaptureBatch(fence=fence, records=(observation(),)))
+ await db.rollback()
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 0
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 0
+ assert await db.scalar(select(Sync.observed_change_sequence)) == 0
+ committed = await capture(database, service, fence, observation())
+ assert committed.sequence == 1
+
+
+async def test_parallel_commits_cannot_skip_a_change_watermark(database, source):
+ service, fence = source
+ first_ready, release_first, second_started = asyncio.Event(), asyncio.Event(), asyncio.Event()
+
+ async def first_writer():
+ async with database() as db:
+ result = await service.store.capture(
+ db, CaptureBatch(fence=fence, records=(observation("first"),))
+ )
+ first_ready.set()
+ await release_first.wait()
+ await db.commit()
+ return result
+
+ async def second_writer():
+ await first_ready.wait()
+ second_started.set()
+ return await capture(database, service, fence, observation("second"))
+
+ first_task = asyncio.create_task(first_writer())
+ second_task = asyncio.create_task(second_writer())
+ await asyncio.wait_for(second_started.wait(), 5)
+ async with database() as db:
+ page = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert page.high_watermark == 0 and not page.changes
+ assert not second_task.done()
+ release_first.set()
+ a, b = await asyncio.wait_for(asyncio.gather(first_task, second_task), 5)
+ assert (a.sequence, b.sequence) == (1, 2)
+ async with database() as db:
+ page = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert [c.sequence for c in page.changes] == [1, 2]
+
+
+async def test_writer_fences_checkpoint_and_job_deletion_preserves_records(database, source):
+ service, fence = source
+ captured = await capture(database, service, fence, observation())
+ next_job = uuid4()
+ async with database() as db:
+ db.add(
+ SyncJob(
+ id=next_job,
+ organization_id=fence.organization_id,
+ sync_id=fence.sync_id,
+ status="running",
+ )
+ )
+ await db.commit()
+ async with database() as db:
+ with pytest.raises(WriterBusy):
+ await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ next_job,
+ attempt_id=uuid4(),
+ attempt_number=1,
+ )
+ async with database() as db:
+ await db.execute(
+ update(SyncJob).where(SyncJob.id == fence.job_id).values(status="completed")
+ )
+ await db.commit()
+ async with database() as db:
+ next_fence = await service.activate_writer(
+ db, fence.organization_id, fence.sync_id, next_job, attempt_id=uuid4(), attempt_number=1
+ )
+ assert next_fence.epoch == fence.epoch + 1
+ async with database() as db:
+ with pytest.raises(StaleWriter):
+ await service.save_checkpoint(db, fence, {"history_id": "old"})
+ async with database() as db:
+ await service.save_checkpoint(db, next_fence, {"history_id": "new"})
+ async with database() as db:
+ await db.execute(delete(SyncJob).where(SyncJob.id == fence.job_id))
+ await db.commit()
+ entity = await db.get(Entity, captured.changes[0].record.id)
+ assert entity is not None and entity.sync_job_id is None
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 1
+ assert (await db.scalar(select(SyncCursor.cursor_data)))["history_id"] == "new"
+
+
+async def test_reconcile_only_complete_exact_container_and_keep_marks_seen(database, source):
+ service, fence = source
+ keep, removed, outside = (
+ observation("keep", "a"),
+ observation("gone", "a"),
+ observation("other", "b"),
+ )
+ await capture(database, service, fence, keep, removed, outside)
+ async with database() as db:
+ await db.execute(
+ update(SyncJob).where(SyncJob.id == fence.job_id).values(status="completed")
+ )
+ new_job = SyncJob(
+ organization_id=fence.organization_id, sync_id=fence.sync_id, status="running"
+ )
+ db.add(new_job)
+ await db.commit()
+ new_fence = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ new_job.id,
+ attempt_id=uuid4(),
+ attempt_number=1,
+ )
+ kept = await capture(database, service, new_fence, keep)
+ assert kept.unchanged == 1
+ async with database() as db:
+ result = await service.reconcile_scope(
+ db,
+ ReconcileScope(
+ fence=new_fence,
+ scope=CompletedScope(record_type="event", container_id="a"),
+ observed_at=datetime.now(timezone.utc),
+ ),
+ )
+ assert [c.record.identity.native_id for c in result.capture.changes] == ["gone"]
+ assert result.capture.changes[0].record.removal_reason == "absent"
+ assert not result.has_more
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(Entity).where(Entity.deleted_at.is_(None))
+ )
+ == 2
+ )
+
+
+async def test_retry_attempt_fences_prior_activity_and_seen_scope(database, source):
+ service, old = source
+ await capture(database, service, old, observation("from-first-attempt", "a"))
+ attempt_id = uuid4()
+ async with database() as db:
+ newer = await service.activate_writer(
+ db,
+ old.organization_id,
+ old.sync_id,
+ old.job_id,
+ attempt_id=attempt_id,
+ attempt_number=2,
+ )
+ repeated = await service.activate_writer(
+ db,
+ old.organization_id,
+ old.sync_id,
+ old.job_id,
+ attempt_id=attempt_id,
+ attempt_number=2,
+ )
+ assert repeated == newer and newer.epoch > old.epoch
+ with pytest.raises(StaleWriter):
+ await capture(database, service, old, observation("late-old"))
+ async with database() as db:
+ with pytest.raises(StaleWriter):
+ await service.activate_writer(
+ db,
+ old.organization_id,
+ old.sync_id,
+ old.job_id,
+ attempt_id=old.attempt_id,
+ attempt_number=1,
+ )
+ # Seen in an interrupted attempt is not seen in the retry's complete scan.
+ async with database() as db:
+ result = await service.reconcile_scope(
+ db,
+ ReconcileScope(
+ fence=newer,
+ scope=CompletedScope(record_type="event", container_id="a"),
+ observed_at=datetime.now(timezone.utc),
+ ),
+ )
+ assert len(result.capture.changes) == 1
+ assert result.capture.changes[0].record.deleted_at is not None
+
+
+async def test_cancellation_and_capture_have_a_database_barrier(database, source):
+ service, fence = source
+ async with database() as writer:
+ await service.store.capture(writer, CaptureBatch(fence=fence, records=(observation(),)))
+ cancellation_started = asyncio.Event()
+
+ async def cancel():
+ async with database() as db:
+ cancellation_started.set()
+ await db.execute(
+ update(SyncJob).where(SyncJob.id == fence.job_id).values(status="cancelled")
+ )
+ await db.commit()
+
+ task = asyncio.create_task(cancel())
+ await cancellation_started.wait()
+ assert not task.done()
+ await writer.commit()
+ await asyncio.wait_for(task, 5)
+ # Cancellation has committed; old capture and checkpoint now both fail.
+ with pytest.raises(StaleWriter):
+ await capture(database, service, fence, observation("after-cancel"))
+ async with database() as db:
+ with pytest.raises(StaleWriter):
+ await service.save_checkpoint(db, fence, {"history_id": "late"})
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncCursor)) == 0
+
+
+async def test_service_rolls_back_partial_batch_on_real_database_failure(database, source):
+ service, fence = source
+ async with database() as db:
+ await db.execute(
+ text("""
+ CREATE FUNCTION reject_second_change() RETURNS trigger LANGUAGE plpgsql AS $$
+ BEGIN
+ IF NEW.sequence = 2 THEN RAISE EXCEPTION 'synthetic journal failure'; END IF;
+ RETURN NEW;
+ END $$
+ """)
+ )
+ await db.execute(
+ text("""
+ CREATE TRIGGER reject_second_change BEFORE INSERT ON entity_change
+ FOR EACH ROW EXECUTE FUNCTION reject_second_change()
+ """)
+ )
+ await db.commit()
+ from sqlalchemy.exc import DBAPIError
+
+ with pytest.raises(DBAPIError, match="synthetic journal failure"):
+ await capture(database, service, fence, observation("first"), observation("second"))
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 0
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 0
+ assert await db.scalar(select(Sync.observed_change_sequence)) == 0
+
+
+@pytest.mark.parametrize("database", ["legacy"], indirect=True)
+async def test_upgrade_preserves_existing_metadata_without_claiming_capture(database):
+ async with database() as db:
+ entity = (await db.scalars(select(Entity))).one()
+ assert entity.entity_id == "legacy-native-key" and entity.hash == "old-hash"
+ assert entity.record_revision == 0 and entity.source_payload is None
+ assert entity.native_id is None and entity.indexed_revision is None
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 0
+ await db.execute(delete(SyncJob).where(SyncJob.id == entity.sync_job_id))
+ await db.commit()
+ await db.refresh(entity)
+ assert entity.sync_job_id is None
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py
new file mode 100644
index 0000000000..8f54f6faeb
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py
@@ -0,0 +1,193 @@
+"""Real PostgreSQL publication races: late feeds never regain search eligibility."""
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select, update
+
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import (
+ CanonicalProjectionStore,
+ publication_matches,
+)
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.tests.helpers import (
+ bind_projection,
+ capture,
+ observation,
+ publish_prepared,
+)
+from airweave.models import Entity, Sync
+
+pytestmark = pytest.mark.integration
+
+
+async def pending(database, fence):
+ async with database() as db:
+ return await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id)
+
+
+async def test_publication_compare_and_swap_revision_generation_and_pipeline(database, source):
+ service, fence = source
+ await bind_projection(database, fence, "slack")
+ await capture(database, service, fence, observation())
+ work = (await pending(database, fence))[0]
+ store = CanonicalProjectionStore()
+ generation = uuid4()
+ async with database() as db:
+ assert await publish_prepared(store, db, work, generation, 3)
+ async with database() as db:
+ assert not await publish_prepared(store, db, work, uuid4(), 3) # timed-out prior feed
+ assert not await pending(database, fence)
+ locator = ProjectionLocator(
+ record_id=work.record.id,
+ revision=1,
+ pipeline_version=1,
+ generation=generation,
+ part_index=0,
+ )
+ assert ProjectionLocator.parse(locator.encode()) == locator
+ async with database() as db:
+ assert (
+ await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator)))
+ == work.record.id
+ )
+ await capture(database, service, fence, observation(payload={"summary": "Changed"}))
+ async with database() as db:
+ assert (
+ await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator)))
+ is None
+ )
+ assert not await publish_prepared(store, db, work, uuid4(), 3)
+ changed = (await pending(database, fence))[0]
+ async with database() as db:
+ await db.execute(update(Sync).values(index_pipeline_version=2))
+ await db.commit()
+ async with database() as db:
+ assert not await publish_prepared(store, db, changed, uuid4(), 3)
+ assert (await pending(database, fence))[0].pipeline_version == 2
+
+
+async def test_parent_access_loss_fences_publication_and_tombstone_is_empty(database, source):
+ service, fence = source
+ await bind_projection(database, fence, "slack")
+ parent_id = RecordIdentity(record_type="calendar", native_id="cal")
+ parent = observation(identity=parent_id)
+ child = observation("event", "cal", parent=parent_id)
+ await capture(database, service, fence, parent, child)
+ rows = await pending(database, fence)
+ child_work = next(row for row in rows if row.record.identity.record_type == "event")
+ await capture(
+ database,
+ service,
+ fence,
+ parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ assert not await publish_prepared(store, db, child_work, uuid4(), 2)
+ rows = await pending(database, fence)
+ assert len(rows) == 1 and rows[0].record.deleted_at is not None
+ async with database() as db:
+ assert await publish_prepared(store, db, rows[0], uuid4(), 0)
+ assert not await pending(database, fence)
+
+
+async def test_failed_work_retries_and_stale_failure_preserves_publication(database, source):
+ service, fence = source
+ await bind_projection(database, fence, "slack")
+ await capture(database, service, fence, observation())
+ work = (await pending(database, fence))[0]
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ await store.fail(db, work, "conversion failed")
+ assert len(await pending(database, fence)) == 1
+ async with database() as db:
+ assert await publish_prepared(store, db, work, uuid4(), 1)
+ async with database() as db:
+ await store.fail(db, work, "late failure")
+ async with database() as db:
+ assert (await db.get(Entity, work.record.id)).projection_error is None
+
+
+def test_locator_rejects_malformed_canonical_values():
+ assert ProjectionLocator.parse("legacy-native-id") is None
+ for value in ("canonical:v2:bad", "canonical:v1:bad", "canonical:"):
+ with pytest.raises(ValueError):
+ ProjectionLocator.parse(value)
+
+
+async def test_owned_payload_mapper_to_publication_pipeline(database, source):
+ from unittest.mock import AsyncMock, MagicMock
+
+ from airweave.domains.entities.canonical.projector import CanonicalProjector
+ from airweave.platform.destinations.vespa.transformer import EntityTransformer
+
+ service, fence = source
+ binding = await bind_projection(database, fence, "slack")
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=RecordIdentity(record_type="channel", native_id="C1"),
+ payload={"id": "C1", "name": "general", "purpose": {"value": "Project work"}},
+ ),
+ )
+ work = (await pending(database, fence))[0]
+ processor = MagicMock()
+
+ async def process(entities, context, runtime, *, strict, expected_ids):
+ assert strict and context.source_short_name == "slack"
+ assert entities[0].name == "general"
+ assert entities[0].airweave_system_metadata.sync_id == fence.sync_id
+ locator = ProjectionLocator.parse(entities[0].entity_id)
+ assert locator.record_id == work.record.id
+ chunk = entities[0].model_copy(deep=True)
+ chunk.airweave_system_metadata.original_entity_id = chunk.entity_id
+ chunk.entity_id += "__chunk_0"
+ return [chunk]
+
+ from airweave.domains.sync_pipeline.pipeline.text_models import BuiltText, BuiltTextBatch
+
+ async def build_text(entities, context, runtime, *, native_bodies=None):
+ return BuiltTextBatch(
+ entities=entities,
+ representations=tuple(
+ BuiltText(entity_id=e.entity_id, text="Synthetic complete text") for e in entities
+ ),
+ )
+
+ processor.build_text = build_text
+ processor.process_built_text = process
+ destination = MagicMock()
+ destination.collection_id = binding.collection_id
+ destination.prepare_documents = lambda chunks: {
+ "base_entity": [
+ EntityTransformer(collection_id=destination.collection_id).transform(chunk)
+ for chunk in chunks
+ ]
+ }
+
+ async def feed_prepared(documents):
+ from airweave.models.projection_generation import ProjectionGeneration
+
+ async with database() as db:
+ manifest = await db.scalar(select(ProjectionGeneration))
+ assert manifest is not None
+ assert manifest.documents[0]["document_id"] == documents["base_entity"][0].id
+ assert (await db.get(Entity, work.record.id)).indexed_generation is None
+ document = documents["base_entity"][0]
+ assert document.id.startswith(f"{fence.sync_id}_{destination.collection_id}_")
+ # Direct reads/navigation query entity fields, not the remote feed key.
+ original = document.fields["airweave_system_metadata_original_entity_id"]
+ assert ProjectionLocator.parse(original).record_id == work.record.id
+ assert document.fields["entity_id"] == original + "__chunk_0"
+ assert document.fields["airweave_system_metadata_sync_id"] == str(fence.sync_id)
+
+ destination.feed_prepared = AsyncMock(side_effect=feed_prepared)
+ storage = MagicMock(write_file=AsyncMock())
+ projector = CanonicalProjector(CanonicalProjectionStore(), database, processor, storage)
+ assert await projector.project_one(work, "slack", destination, MagicMock())
+ destination.feed_prepared.assert_awaited_once()
+ assert not await pending(database, fence)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_admission.py b/backend/airweave/domains/entities/canonical/tests/test_projection_admission.py
new file mode 100644
index 0000000000..a93ef8295d
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection_admission.py
@@ -0,0 +1,176 @@
+"""Authenticated binding and current-input admission before external projection work."""
+
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select, update
+
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.tests.helpers import (
+ bind_projection,
+ capture,
+ observation,
+ publish_prepared,
+)
+from airweave.models import Entity, Organization, Sync
+from airweave.models.collection import Collection
+from airweave.models.source_connection import SourceConnection
+
+
+async def test_pending_requires_exact_authenticated_tenant_binding(database, source):
+ service, fence = source
+ await capture(database, service, fence, observation())
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ assert not await store.pending(db, fence.organization_id, fence.sync_id)
+ binding = await bind_projection(database, fence)
+ async with database() as db:
+ assert len(await store.pending(db, fence.organization_id, fence.sync_id)) == 1
+ assert not await store.pending(db, uuid4(), fence.sync_id)
+ await db.execute(update(SourceConnection).values(is_authenticated=False))
+ await db.commit()
+ async with database() as db:
+ assert not await store.pending(db, fence.organization_id, fence.sync_id)
+ foreign = Organization(name="Other synthetic tenant")
+ db.add(foreign)
+ await db.flush()
+ await db.execute(update(SourceConnection).values(is_authenticated=True))
+ await db.execute(
+ update(Collection)
+ .where(Collection.id == binding.collection_id)
+ .values(organization_id=foreign.id)
+ )
+ await db.commit()
+ async with database() as db:
+ assert not await store.pending(db, fence.organization_id, fence.sync_id)
+
+
+@pytest.mark.parametrize("change", ["auth", "revision", "pipeline", "source", "collection"])
+async def test_stale_work_is_denied_before_mapper_storage_or_processor(
+ database, source, monkeypatch, change
+):
+ service, fence = source
+ binding = await bind_projection(database, fence)
+ await capture(database, service, fence, observation())
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ if change == "revision":
+ await capture(database, service, fence, observation(payload={"summary": "new"}))
+ else:
+ async with database() as db:
+ if change == "auth":
+ await db.execute(update(SourceConnection).values(is_authenticated=False))
+ elif change == "pipeline":
+ await db.execute(update(Sync).values(index_pipeline_version=2))
+ elif change == "source":
+ await db.execute(update(SourceConnection).values(short_name="slack"))
+ else:
+ await db.execute(update(SourceConnection).values(readable_collection_id=None))
+ await db.commit()
+ mapper = MagicMock(side_effect=AssertionError("mapper must not read stale input"))
+ monkeypatch.setattr("airweave.domains.entities.canonical.projection_mappers.map_record", mapper)
+ processor, storage = MagicMock(), AsyncMock()
+ destination = MagicMock(collection_id=binding.collection_id, feed_prepared=AsyncMock())
+ project = CanonicalProjector(store, database, processor, storage)
+ assert not await project.project_one(work, "gmail", destination, MagicMock())
+ mapper.assert_not_called()
+ assert not processor.mock_calls and not storage.mock_calls
+ destination.feed_prepared.assert_not_awaited()
+
+
+async def test_deauth_after_preparation_rejects_publication_and_failure_write(database, source):
+ service, fence = source
+ binding = await bind_projection(database, fence)
+ await capture(
+ database, service, fence, observation(kind="delete", removal_reason="provider_deleted")
+ )
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ assert await store.admit(db, work)
+ generation = uuid4()
+ async with database() as db:
+ assert await store.prepare(db, work, generation, binding.collection_id, ())
+ async with database() as db:
+ # Same Sync fence used by owned provisioning pause/withdraw.
+ await db.scalar(select(Sync).where(Sync.id == fence.sync_id).with_for_update())
+ await db.execute(update(SourceConnection).values(is_authenticated=False))
+ await db.commit()
+ async with database() as db:
+ assert not await store.publish(db, work, generation, 0)
+ await store.fail(db, work, "late failure")
+ row = await db.get(Entity, work.record.id)
+ assert row.indexed_generation is None and row.projection_error is None
+
+
+async def test_wrong_destination_rejected_before_processing_and_manifest(database, source):
+ service, fence = source
+ await bind_projection(database, fence)
+ await capture(database, service, fence, observation())
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ processor = MagicMock()
+ project = CanonicalProjector(store, database, processor, AsyncMock())
+ assert not await project.project_one(
+ work, "gmail", MagicMock(collection_id=uuid4()), MagicMock()
+ )
+ assert not processor.mock_calls
+ async with database() as db:
+ with pytest.raises(ValueError, match="destination"):
+ await publish_prepared(store, db, work, uuid4(), 1, collection_id=uuid4())
+
+
+async def test_publication_waits_for_owned_pause_fence_and_observes_revocation(database, source):
+ import asyncio
+
+ service, fence = source
+ binding = await bind_projection(database, fence)
+ await capture(
+ database, service, fence, observation(kind="delete", removal_reason="provider_deleted")
+ )
+ store = CanonicalProjectionStore()
+ generation = uuid4()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ assert await store.prepare(db, work, generation, binding.collection_id, ())
+
+ async def publish():
+ async with database() as db:
+ return await store.publish(db, work, generation, 0)
+
+ async with database() as paused:
+ await paused.scalar(select(Sync).where(Sync.id == fence.sync_id).with_for_update())
+ publishing = asyncio.create_task(publish())
+ try:
+ # Publication cannot pass the same Sync row held by owned pause.
+ with pytest.raises(TimeoutError):
+ await asyncio.wait_for(asyncio.shield(publishing), timeout=0.1)
+ await paused.execute(update(SourceConnection).values(is_authenticated=False))
+ await paused.commit()
+ assert not await asyncio.wait_for(publishing, timeout=5)
+ finally:
+ if not publishing.done():
+ publishing.cancel()
+ await asyncio.gather(publishing, return_exceptions=True)
+
+
+async def test_stale_activity_route_does_not_poison_new_binding_work(database, source):
+ service, fence = source
+ binding = await bind_projection(database, fence, "slack")
+ await capture(database, service, fence, observation())
+ project = CanonicalProjector(CanonicalProjectionStore(), database, MagicMock(), AsyncMock())
+ result = await project.batch(
+ fence.organization_id,
+ fence.sync_id,
+ "gmail",
+ MagicMock(collection_id=binding.collection_id),
+ MagicMock(),
+ )
+ assert result.superseded == 1 and result.failed == 0
+ async with database() as db:
+ row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id))
+ assert row.projection_error is None and row.indexed_generation is None
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_projection_exclusions.py
new file mode 100644
index 0000000000..f7bdc96a3b
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection_exclusions.py
@@ -0,0 +1,150 @@
+"""Deliberate search exclusions publish normally without losing retained originals."""
+
+from contextlib import asynccontextmanager
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage
+from airweave.domains.entities.canonical.projection_inputs import ProjectionInputs
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.query import CanonicalQueryService
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation
+from airweave.domains.native_ingestion.tests.test_ingestion import NOW, bind, ingest
+from airweave.domains.native_ingestion.tests.test_projection import knowledge, message, session
+from airweave.models import Entity
+
+
+@pytest.mark.parametrize(
+ "kind",
+ [
+ "meeting_listing",
+ "scratchpad_listing",
+ "knowledge",
+ "message",
+ "event_occurrence",
+ "event",
+ ],
+)
+async def test_excluded_originals_publish_zero_documents(database, source, kind):
+ service, fence = source
+ if kind in {"knowledge", "message"}:
+ provider = "almanac"
+ await bind(database, fence, dataset="knowledge" if kind == "knowledge" else "sessions")
+ snapshots = (
+ (knowledge(archived_at=NOW.isoformat()),)
+ if kind == "knowledge"
+ else (
+ session(),
+ message(
+ "Original retained",
+ import_provenance={
+ "legacy_row_id": 1,
+ "active": False,
+ "compacted": True,
+ },
+ ),
+ )
+ )
+ await ingest(database, fence, *snapshots)
+ else:
+ provider = "wispr" if kind.endswith("listing") else "google_calendar"
+ await bind_projection(database, fence, provider)
+ payload = {"id": "original", "title": "Original retained"}
+ if kind == "event":
+ payload.update(
+ status="cancelled",
+ recurringEventId="series",
+ originalStartTime={"date": "2026-10-01"},
+ )
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=RecordIdentity(record_type=kind, native_id="original"),
+ payload=payload,
+ completeness="metadata_only",
+ ),
+ )
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = next(
+ item
+ for item in await store.pending(db, fence.organization_id, fence.sync_id)
+ if item.record.identity.record_type == kind
+ )
+ processor, storage = MagicMock(), MagicMock()
+ destination = MagicMock(collection_id=work.binding.collection_id)
+ projector = CanonicalProjector(store, database, processor, storage)
+ assert await projector.project_one(work, provider, destination, MagicMock())
+ processor.build_text.assert_not_called()
+ destination.feed_prepared.assert_not_called()
+ reader = CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test")
+ async with database() as db:
+ retained = await reader.read(db, fence.organization_id, fence.sync_id, work.record.id)
+ assert retained.payload == work.record.payload
+ row = await db.scalar(select(Entity).where(Entity.id == work.record.id))
+ assert row.indexed_generation is not None and row.indexed_chunk_count == 0
+ assert row.projection_error is None
+ assert work.record.id not in {
+ item.record.id for item in await store.pending(db, fence.organization_id, fence.sync_id)
+ }
+
+
+@pytest.mark.parametrize("record_type", ["unknown", "event_occurrence", "event"])
+async def test_unknown_empty_projection_remains_an_error(
+ database, source, monkeypatch, record_type
+):
+ service, fence = source
+ binding = await bind_projection(database, fence, "wispr")
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=RecordIdentity(record_type=record_type, native_id="original"),
+ payload={
+ "status": "cancelled",
+ "recurringEventId": "series",
+ "originalStartTime": {"date": "2026-10-01"},
+ },
+ completeness="metadata_only",
+ ),
+ )
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ (work,) = await store.pending(db, fence.organization_id, fence.sync_id)
+
+ @asynccontextmanager
+ async def empty_mapper(*args):
+ yield ProjectionInputs(parts=())
+
+ monkeypatch.setattr(
+ "airweave.domains.entities.canonical.projection_mappers.map_record", empty_mapper
+ )
+ projector = CanonicalProjector(store, database, MagicMock(), MagicMock())
+ with pytest.raises(ValueError, match="no required content"):
+ await projector.project_one(
+ work, "wispr", MagicMock(collection_id=binding.collection_id), MagicMock()
+ )
+ generation = uuid4()
+ async with database() as db:
+ assert await store.prepare(
+ db,
+ work,
+ generation,
+ binding.collection_id,
+ (),
+ coverage=ExtractionCoverage(parts=()),
+ text_representations=(),
+ )
+ async with database() as db:
+ with pytest.raises(ValueError, match="indexed content or explicit unavailable extraction"):
+ await store.publish(db, work, generation, 0)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py
new file mode 100644
index 0000000000..b9f0b0343d
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py
@@ -0,0 +1,285 @@
+"""Real PostgreSQL retirement fences and repeatable cleanup after late remote writes."""
+
+from datetime import datetime, timedelta, timezone
+from uuid import UUID, uuid4
+
+from sqlalchemy import delete, select
+
+from airweave.domains.entities.canonical.projection_gc import ProjectionGCStore
+from airweave.domains.entities.canonical.projection_models import (
+ ProjectionDocument,
+ ProjectionLocator,
+ scope_projection_document_id,
+)
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation
+from airweave.models import Entity, Sync
+from airweave.models.projection_generation import ProjectionGeneration
+
+
+async def prepare(database, source, count=3):
+ service, fence = source
+ binding = await bind_projection(database, fence)
+ await capture(database, service, fence, observation())
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ generation = uuid4()
+ locator = ProjectionLocator(
+ record_id=work.record.id,
+ revision=work.record.revision,
+ pipeline_version=work.pipeline_version,
+ generation=generation,
+ part_index=0,
+ ).encode()
+ collection_id = binding.collection_id
+ docs = tuple(
+ ProjectionDocument(
+ schema_name="base_entity",
+ document_id=scope_projection_document_id(
+ work.record.sync_id, collection_id, f"Entity_{locator}__chunk_{n}"
+ ),
+ )
+ for n in range(count)
+ )
+ async with database() as db:
+ assert await store.prepare(db, work, generation, collection_id, docs)
+ return store, work, generation, docs
+
+
+async def test_retirement_prevents_late_publish_and_repeats_after_late_feed(database, source):
+ store, work, generation, docs = await prepare(database, source)
+ gc = ProjectionGCStore()
+ now = datetime.now(timezone.utc) + timedelta(hours=2)
+ async with database() as db:
+ page = await gc.claim(db, generation, now=now, limit=2)
+ assert page.documents == docs[:2]
+ async with database() as db:
+ assert not await store.publish(db, work, generation, 3)
+ async with database() as db:
+ await gc.acknowledge(db, page, now=now, error="HTTPError")
+ async with database() as db:
+ row = await db.get(ProjectionGeneration, generation)
+ assert row.delete_cursor == 0 and row.gc_error == "HTTPError"
+ now += timedelta(minutes=6)
+ async with database() as db:
+ retry = await gc.claim(db, generation, now=now, limit=2)
+ assert retry.documents == page.documents
+ async with database() as db:
+ await gc.acknowledge(db, retry, now=now)
+ async with database() as db:
+ tail = await gc.claim(db, generation, now=now, limit=2)
+ assert tail.documents == docs[2:]
+ async with database() as db:
+ await gc.acknowledge(db, tail, now=now)
+ # A timed-out feed can recreate documents after this full delete pass.
+ # The retained ledger schedules all exact IDs again; no query-index discovery required.
+ async with database() as db:
+ repeated = await gc.claim(db, generation, now=now + timedelta(hours=1), limit=100)
+ assert repeated.documents == docs
+ async with database() as db:
+ await gc.acknowledge(db, tail, now=now) # stale callback cannot ack newer attempt
+ row = await db.get(ProjectionGeneration, generation)
+ assert row.gc_attempt == repeated.attempt and row.gc_passes == 1
+
+
+async def test_current_generation_protected_then_atomically_retired_on_replacement(
+ database, source
+):
+ store, work, generation, docs = await prepare(database, source)
+ gc = ProjectionGCStore()
+ async with database() as db:
+ assert await store.publish(db, work, generation, len(docs))
+ async with database() as db:
+ assert (
+ await gc.claim(db, generation, now=datetime.now(timezone.utc) + timedelta(days=2))
+ is None
+ )
+ service, fence = source
+ await capture(database, service, fence, observation(payload={"changed": True}))
+ async with database() as db:
+ newer = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ new_generation = uuid4()
+ new_docs = tuple(
+ doc.model_copy(
+ update={
+ "document_id": doc.document_id.replace(
+ str(generation), str(new_generation)
+ ).replace(":1:1:", ":2:1:")
+ }
+ )
+ for doc in docs
+ )
+ async with database() as db:
+ assert await store.prepare(
+ db, newer, new_generation, UUID(docs[0].document_id.split("_")[1]), new_docs
+ )
+ async with database() as db:
+ assert await store.publish(db, newer, new_generation, len(docs))
+ async with database() as db:
+ old = await db.get(ProjectionGeneration, generation)
+ assert old.retired_at is not None
+ assert (await db.get(Entity, work.record.id)).indexed_generation == new_generation
+
+
+async def test_source_deletion_does_not_erase_cleanup_obligations(database, source):
+ store, work, generation, docs = await prepare(database, source)
+ async with database() as db:
+ await db.execute(delete(Sync).where(Sync.id == work.record.sync_id))
+ await db.commit()
+ gc = ProjectionGCStore()
+ async with database() as db:
+ page = await gc.claim(db, generation, now=datetime.now(timezone.utc) + timedelta(hours=2))
+ assert page.documents == docs
+ async with database() as db:
+ assert not await store.publish(db, work, generation, len(docs))
+
+
+async def test_gc_waits_on_publication_lock_and_cannot_retire_current(database, source):
+ import asyncio
+
+ store, work, generation, docs = await prepare(database, source)
+ gc = ProjectionGCStore()
+ entered = asyncio.Event()
+ async with database() as publishing:
+ # Hold the real shared lock and publish the pointer in the same transaction.
+ await publishing.scalar(
+ select(Sync.id).where(Sync.id == work.record.sync_id).with_for_update()
+ )
+ row = await publishing.get(Entity, work.record.id)
+ row.indexed_generation = generation
+ await publishing.flush()
+
+ async def collect():
+ entered.set()
+ async with database() as db:
+ return await gc.claim(
+ db, generation, now=datetime.now(timezone.utc) + timedelta(hours=2)
+ )
+
+ attempt = asyncio.create_task(collect())
+ await entered.wait()
+ await asyncio.sleep(0.05)
+ assert not attempt.done()
+ await publishing.commit()
+ assert await attempt is None
+ async with database() as db:
+ assert (await db.get(ProjectionGeneration, generation)).retired_at is None
+
+
+async def test_deleted_current_generation_retired_without_another_projection(database, source):
+ store, work, generation, docs = await prepare(database, source)
+ async with database() as db:
+ assert await store.publish(db, work, generation, len(docs))
+ service, fence = source
+ await capture(
+ database,
+ service,
+ fence,
+ observation(kind="delete", removal_reason="provider_deleted", payload={"id": "one"}),
+ )
+ async with database() as db:
+ page = await ProjectionGCStore().claim(
+ db, generation, now=datetime.now(timezone.utc) + timedelta(days=2)
+ )
+ assert page.documents == docs
+ record = await db.get(Entity, work.record.id)
+ assert record.indexed_generation is None and record.indexed_revision is None
+ assert record.deleted_at is not None
+ assert (await db.get(ProjectionGeneration, generation)).retired_at is not None
+
+
+async def test_hidden_parent_retirement_restoration_requires_fresh_generation(database, source):
+ from airweave.domains.entities.canonical.requests import RecordIdentity
+
+ store, work, generation, docs = await prepare(database, source)
+ async with database() as db:
+ assert await store.publish(db, work, generation, len(docs))
+ service, fence = source
+ parent = RecordIdentity(record_type="calendar", native_id="calendar-one")
+ # Add a real captured parent reference, then simulate the crash window between
+ # parent revocation and bounded child reconciliation: child remains active.
+ await capture(database, service, fence, observation(identity=parent))
+ await capture(database, service, fence, observation(parent=parent, allow_reparent=True))
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=parent,
+ kind="delete",
+ removal_reason="access_revoked",
+ payload={"id": "calendar-one"},
+ ),
+ )
+ async with database() as db:
+ page = await ProjectionGCStore().claim(
+ db, generation, now=datetime.now(timezone.utc) + timedelta(days=2)
+ )
+ assert page.documents == docs
+ assert (await db.get(Entity, work.record.id)).deleted_at is None
+ await capture(database, service, fence, observation(identity=parent))
+ async with database() as db:
+ assert all(
+ item.record.id != work.record.id
+ for item in await store.pending(db, fence.organization_id, fence.sync_id)
+ )
+ # Restoration of the parent alone cannot reauthorize retained descendant bytes.
+ await capture(database, service, fence, observation(parent=parent))
+ async with database() as db:
+ restored = next(
+ item
+ for item in await store.pending(db, fence.organization_id, fence.sync_id)
+ if item.record.id == work.record.id
+ )
+ assert restored.previous_generation is None
+ assert not await store.publish(db, restored, generation, len(docs))
+
+
+async def test_manifest_rejects_foreign_scope_and_unanchored_locator(database, source):
+ import pytest
+
+ store, work, generation, docs = await prepare(database, source)
+ collection_id = UUID(docs[0].document_id.split("_")[1])
+ for document_id in (
+ docs[0].document_id.replace(str(work.record.sync_id), str(uuid4())),
+ docs[0].document_id.replace(str(collection_id), str(uuid4())),
+ "arbitrary_" + docs[0].document_id,
+ docs[0].document_id + "_suffix",
+ ):
+ async with database() as db:
+ with pytest.raises(ValueError, match="Projection manifest"):
+ await store.prepare(
+ db,
+ work,
+ generation,
+ collection_id,
+ (docs[0].model_copy(update={"document_id": document_id}),),
+ )
+
+
+async def test_empty_tombstone_publication_does_not_become_pending_each_gc_pass(database, source):
+ service, fence = source
+ binding = await bind_projection(database, fence)
+ await capture(
+ database,
+ service,
+ fence,
+ observation(kind="delete", removal_reason="provider_deleted", payload={"id": "one"}),
+ )
+ store = CanonicalProjectionStore()
+ generation = uuid4()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ async with database() as db:
+ assert await store.prepare(db, work, generation, binding.collection_id, ())
+ async with database() as db:
+ assert await store.publish(db, work, generation, 0)
+ async with database() as db:
+ assert (
+ await ProjectionGCStore().claim(
+ db, generation, now=datetime.now(timezone.utc) + timedelta(days=2)
+ )
+ is None
+ )
+ assert not await store.pending(db, fence.organization_id, fence.sync_id)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py b/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py
new file mode 100644
index 0000000000..94b8897837
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py
@@ -0,0 +1,77 @@
+"""OCR-dependent bytes remain captured and pending instead of vanishing from search."""
+
+import hashlib
+from io import BytesIO
+from unittest.mock import AsyncMock, MagicMock
+
+import fitz
+from PIL import Image
+from sqlalchemy import select
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.core.logging import logger
+from airweave.domains.converters.registry import ConverterRegistry
+from airweave.domains.embedders.fakes.embedder import FakeDenseEmbedder, FakeSparseEmbedder
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity
+from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation
+from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor
+from airweave.models import Entity
+
+
+async def test_scanned_pdf_without_ocr_stays_pending_with_owned_bytes(database, source, tmp_path):
+ """Use real PDF conversion and SQL; no mocked conversion or tracker behavior."""
+ image = BytesIO()
+ Image.new("RGB", (100, 100), color="blue").save(image, format="PNG")
+ document = fitz.open()
+ page = document.new_page()
+ page.insert_image(page.rect, stream=image.getvalue())
+ content = document.tobytes()
+ document.close()
+ service, fence = source
+ binding = await bind_projection(database, fence, "google_drive")
+ digest = hashlib.sha256(content).hexdigest()
+ blob = BlobReference(
+ key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}",
+ sha256=digest,
+ size_bytes=len(content),
+ media_type="application/pdf",
+ )
+ storage = FilesystemBackend(tmp_path)
+ await storage.write_file(blob.key, content)
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=RecordIdentity(record_type="file", native_id="scanned"),
+ payload={"id": "scanned", "name": "scan.pdf", "mimeType": "application/pdf"},
+ blobs=(blob,),
+ ),
+ )
+ destination = MagicMock(collection_id=binding.collection_id, feed_prepared=AsyncMock())
+ store = CanonicalProjectionStore()
+ projector = CanonicalProjector(
+ store,
+ database,
+ ChunkEmbedProcessor(
+ ConverterRegistry(ocr_provider=None), FakeDenseEmbedder(), FakeSparseEmbedder()
+ ),
+ storage,
+ )
+ result = await projector.batch(
+ fence.organization_id, fence.sync_id, "google_drive", destination, logger
+ )
+ assert result.failed == 1 and result.published == 0
+ destination.prepare_documents.assert_not_called()
+ destination.feed_prepared.assert_not_awaited()
+ async with database() as db:
+ pending = await store.pending(db, fence.organization_id, fence.sync_id)
+ assert len(pending) == 1
+ row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id))
+ assert row.projection_error is not None and row.indexed_generation is None
+ record = await service.store.read(db, fence.organization_id, fence.sync_id, row.id)
+ assert record.completeness == "complete" and record.content_access == "available"
+ assert record.blobs == (blob,) and record.payload["id"] == "scanned"
+ assert await storage.read_file(blob.key, max_bytes=blob.size_bytes) == content
diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py
new file mode 100644
index 0000000000..0302c91b6b
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py
@@ -0,0 +1,164 @@
+"""Recovery discovers exact pending scope and never retries failed originals implicitly."""
+
+from unittest.mock import AsyncMock
+from uuid import uuid4
+
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.projection_models import ProjectionBatchResult
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation
+from airweave.models import Collection, Entity, Organization, SourceConnection, Sync, SyncJob
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+async def test_source_pages_skip_failed_indexed_unsupported_and_foreign_rows(database, source):
+ service, seed = source
+ foreign = uuid4()
+ async with database() as db:
+ db.add(Organization(id=foreign, name="Foreign"))
+ metadata = VectorDbDeploymentMetadata(
+ dense_embedder="test", sparse_embedder="test", embedding_dimensions=3
+ )
+ db.add(metadata)
+ await db.flush()
+ for owner, readable in ((seed.organization_id, "owned"), (foreign, "foreign")):
+ db.add(
+ Collection(
+ id=uuid4(),
+ name=readable,
+ readable_id=readable,
+ organization_id=owner,
+ vector_db_deployment_metadata_id=metadata.id,
+ )
+ )
+ await db.commit()
+ expected = set()
+ for case in (
+ "fresh",
+ "provider",
+ "unauthenticated",
+ "missing-collection",
+ "foreign-collection",
+ "tombstone",
+ "failed",
+ "indexed",
+ "unsupported",
+ "foreign-source",
+ "foreign-record",
+ "withdrawn",
+ ):
+ sync_id, job_id = uuid4(), uuid4()
+ async with database() as db:
+ db.add(Sync(id=sync_id, organization_id=seed.organization_id, name=case))
+ await db.flush()
+ db.add(
+ SyncJob(
+ id=job_id,
+ sync_id=sync_id,
+ organization_id=seed.organization_id,
+ status="running",
+ )
+ )
+ db.add(
+ SourceConnection(
+ id=uuid4(),
+ sync_id=sync_id,
+ organization_id=foreign if case == "foreign-source" else seed.organization_id,
+ name=case,
+ short_name=(
+ "legacy"
+ if case == "unsupported"
+ else "slack"
+ if case == "provider"
+ else "almanac"
+ ),
+ is_authenticated=case != "unauthenticated",
+ readable_collection_id=(
+ None
+ if case == "missing-collection"
+ else "foreign"
+ if case == "foreign-collection"
+ else "owned"
+ ),
+ )
+ )
+ await db.commit()
+ async with database() as db:
+ fence = await service.activate_writer(
+ db,
+ seed.organization_id,
+ sync_id,
+ job_id,
+ attempt_id=uuid4(),
+ attempt_number=1,
+ )
+ await capture(database, service, fence, observation())
+ if case == "tombstone":
+ await capture(
+ database,
+ service,
+ fence,
+ observation().model_copy(update={"kind": "delete", "removal_reason": "deleted"}),
+ )
+ async with database() as db:
+ record = await db.scalar(select(Entity).where(Entity.sync_id == sync_id))
+ if case == "failed":
+ record.projection_error = "ConversionFailed"
+ elif case == "indexed":
+ record.indexed_revision = record.record_revision
+ record.indexed_pipeline_version = 1
+ record.indexed_generation = uuid4()
+ elif case == "foreign-record":
+ record.organization_id = foreign
+ elif case == "withdrawn":
+ record.removal_reason = "access_revoked"
+ await db.commit()
+ if case in ("fresh", "provider", "tombstone"):
+ expected.add(sync_id)
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ first = await store.pending_sources(db, ("almanac", "slack"), limit=2)
+ assert len(first.sources) == 2
+ assert first.next_cursor == first.sources[-1].sync_id
+ second = await store.pending_sources(
+ db, ("almanac", "slack"), after_id=first.next_cursor, limit=2
+ )
+ assert len(second.sources) == 1 and second.next_cursor is None
+ assert {item.sync_id for item in (*first.sources, *second.sources)} == expected
+ assert all(
+ item.organization_id == seed.organization_id
+ for item in (*first.sources, *second.sources)
+ )
+ assert not (await store.pending_sources(db, ())).sources
+
+
+async def test_recovery_execution_skips_failed_until_original_changes(database, source):
+ service, fence = source
+ await bind_projection(database, fence, "almanac")
+ await capture(database, service, fence, observation())
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ await store.fail(db, work, "ConversionFailed")
+ async with database() as db:
+ assert len(await store.pending(db, fence.organization_id, fence.sync_id)) == 1
+ assert not await store.pending(db, fence.organization_id, fence.sync_id, skip_failed=True)
+ # Isolate only the expensive external projection; exercise the actual batch selector.
+ projector = CanonicalProjector(store, database, None, None)
+ projector.project_one = AsyncMock(return_value=True)
+ logger = AsyncMock()
+ assert (
+ await projector.batch(
+ fence.organization_id, fence.sync_id, "almanac", None, logger, skip_failed=True
+ )
+ == ProjectionBatchResult()
+ )
+ projector.project_one.assert_not_awaited()
+ await capture(database, service, fence, observation(payload={"summary": "Changed original"}))
+ result = await projector.batch(
+ fence.organization_id, fence.sync_id, "almanac", None, logger, skip_failed=True
+ )
+ assert result.published == 1
+ projector.project_one.assert_awaited_once()
diff --git a/backend/airweave/domains/entities/canonical/tests/test_provisioning_fence.py b/backend/airweave/domains/entities/canonical/tests/test_provisioning_fence.py
new file mode 100644
index 0000000000..3d41de139c
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_provisioning_fence.py
@@ -0,0 +1,75 @@
+"""Connection changes invalidate capture independently of worker cancellation."""
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select, update
+
+from airweave.domains.entities.canonical.store import StaleWriter
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.models import Sync, SyncJob
+
+pytestmark = pytest.mark.integration
+
+
+async def test_reconnect_rejects_existing_writer_without_losing_history(database, source):
+ service, fence = source
+ original = await capture(database, service, fence, observation())
+ async with database() as db:
+ await db.execute(
+ update(Sync).where(Sync.id == fence.sync_id).values(provisioning_generation=1)
+ )
+ await db.commit()
+
+ with pytest.raises(StaleWriter):
+ await capture(database, service, fence, observation("stale-credential-result"))
+ async with database() as db:
+ with pytest.raises(StaleWriter):
+ await service.store.admit_job(db, fence.organization_id, fence.sync_id, fence.job_id)
+ page = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert page.high_watermark == original.sequence
+ assert [change.record.identity.native_id for change in page.changes] == ["one"]
+
+
+async def test_managed_job_requires_verified_active_generation_even_on_identical_retry(
+ database, source
+):
+ service, fence = source
+ async with database() as db:
+ await db.execute(
+ update(Sync).where(Sync.id == fence.sync_id).values(
+ provisioning_generation=1, status="active"
+ )
+ )
+ await db.execute(
+ update(SyncJob).where(SyncJob.id == fence.job_id).values(provisioning_generation=1)
+ )
+ await db.commit()
+
+ async with database() as db:
+ with pytest.raises(StaleWriter):
+ await service.store.admit_job(db, fence.organization_id, fence.sync_id, fence.job_id)
+ with pytest.raises(StaleWriter):
+ await service.store.activate_writer(
+ db, fence.organization_id, fence.sync_id, fence.job_id,
+ attempt_id=fence.attempt_id, attempt_number=fence.attempt_number,
+ )
+ await db.rollback()
+ await db.execute(
+ update(Sync).where(Sync.id == fence.sync_id).values(provisioning_ready_generation=1)
+ )
+ await db.commit()
+ assert await service.store.admit_job(
+ db, fence.organization_id, fence.sync_id, fence.job_id
+ ) == 1
+ with pytest.raises(StaleWriter):
+ await service.store.admit_job(db, uuid4(), fence.sync_id, fence.job_id)
+
+ await capture(database, service, fence, observation("verified"))
+ async with database() as db:
+ await db.execute(update(Sync).where(Sync.id == fence.sync_id).values(status="paused"))
+ await db.commit()
+ with pytest.raises(StaleWriter):
+ await capture(database, service, fence, observation("after-disconnect"))
+ async with database() as db:
+ assert await db.scalar(select(Sync.observed_change_sequence)) == 1
diff --git a/backend/airweave/domains/entities/canonical/tests/test_query.py b/backend/airweave/domains/entities/canonical/tests/test_query.py
new file mode 100644
index 0000000000..af7396b574
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_query.py
@@ -0,0 +1,225 @@
+"""Exact query behavior against committed records, plus signed cursor isolation."""
+
+from uuid import uuid4
+
+import pytest
+
+from airweave.domains.entities.canonical.query import CanonicalQueryService, InvalidRecordCursor
+from airweave.domains.entities.canonical.query_models import RecordFilters, RecordListQuery
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.requests import CaptureBatch
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, SourceNotFound
+from airweave.domains.entities.canonical.tests.helpers import observation
+
+
+def query_service():
+ return CanonicalQueryService(
+ CanonicalRecordStore(), CanonicalQueryStore(), "synthetic-test-key"
+ )
+
+
+async def test_list_filters_and_cursor_scope(database, source):
+ capture, fence = source
+ async with database() as db:
+ await capture.capture(
+ db,
+ CaptureBatch(
+ fence=fence,
+ records=(
+ observation("one", "inbox"),
+ observation("two", "inbox"),
+ observation("three", "elsewhere"),
+ ),
+ ),
+ )
+ service = query_service()
+ filters = RecordFilters(container_id="inbox")
+ async with database() as db:
+ first = await service.list_records(
+ db, fence.organization_id, fence.sync_id, RecordListQuery(filters=filters, limit=1)
+ )
+ second = await service.list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordListQuery(filters=filters, limit=1, cursor=first.next_cursor),
+ )
+ assert first.has_more and not second.has_more
+ assert {row.identity.native_id for row in (*first.records, *second.records)} == {
+ "one",
+ "two",
+ }
+ assert first.records[0].id < second.records[0].id
+ assert first.consistency == "live"
+ with pytest.raises(InvalidRecordCursor):
+ await service.list_records(
+ db, fence.organization_id, fence.sync_id, RecordListQuery(cursor=first.next_cursor)
+ )
+ with pytest.raises(InvalidRecordCursor):
+ await service.list_records(
+ db,
+ uuid4(),
+ fence.sync_id,
+ RecordListQuery(filters=filters, cursor=first.next_cursor),
+ )
+ with pytest.raises(SourceNotFound):
+ await service.list_records(db, uuid4(), fence.sync_id, RecordListQuery())
+
+
+async def test_change_window_continues_then_polls_new_commits(database, source):
+ capture, fence = source
+ service = query_service()
+ async with database() as db:
+ await capture.capture(
+ db, CaptureBatch(fence=fence, records=(observation("one"), observation("two")))
+ )
+ async with database() as db:
+ first = await service.changes(db, fence.organization_id, fence.sync_id, limit=1)
+ async with database() as db:
+ await capture.capture(db, CaptureBatch(fence=fence, records=(observation("three"),)))
+ async with database() as db:
+ second = await service.changes(
+ db, fence.organization_id, fence.sync_id, cursor=first.next_cursor, limit=1
+ )
+ third = await service.changes(
+ db, fence.organization_id, fence.sync_id, cursor=second.next_cursor, limit=1
+ )
+ assert first.high_watermark == second.high_watermark == 2
+ assert [first.changes[0].sequence, second.changes[0].sequence, third.changes[0].sequence] == [
+ 1,
+ 2,
+ 3,
+ ]
+ assert first.has_more and not second.has_more and not third.has_more
+ assert third.high_watermark == 3
+
+
+async def test_deleted_record_is_readable_but_not_in_active_list(database, source):
+ capture, fence = source
+ service = query_service()
+ async with database() as db:
+ result = await capture.capture(
+ db,
+ CaptureBatch(
+ fence=fence,
+ records=(observation("deleted", kind="delete", removal_reason="provider_deleted"),),
+ ),
+ )
+ record_id = result.changes[0].record.id
+ async with database() as db:
+ record = await service.read(db, fence.organization_id, fence.sync_id, record_id)
+ active = await service.list_records(
+ db, fence.organization_id, fence.sync_id, RecordListQuery()
+ )
+ deleted = await service.list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordListQuery(filters=RecordFilters(state="deleted")),
+ )
+ assert record.deleted_at is not None
+ assert not active.records
+ assert [row.id for row in deleted.records] == [record_id]
+
+
+async def test_parent_filter_is_exact_visible_and_bound_to_continuation(database, source):
+ from airweave.domains.entities.canonical.query import RecordNotFound
+ from airweave.domains.entities.canonical.requests import RecordIdentity
+ from airweave.domains.entities.canonical.tests.helpers import capture as store
+
+ capture, fence = source
+ parent = observation(
+ identity=RecordIdentity(record_type="message", native_id="parent", container_id="C1")
+ ).model_copy(
+ update={"payload": {"files": ["F1", "F2"]}, "descendant_visibility_fields": ("files",)}
+ )
+ other = observation(
+ identity=RecordIdentity(record_type="message", native_id="parent", container_id="C2")
+ )
+ children = tuple(
+ observation(
+ identity=RecordIdentity(record_type="file", native_id=name, container_id="owned"),
+ parent=parent.identity,
+ )
+ for name in ("F1", "F2")
+ )
+ unrelated = observation(
+ identity=RecordIdentity(record_type="file", native_id="F3", container_id="other"),
+ parent=other.identity,
+ )
+ saved = await store(database, capture, fence, parent, other, *children, unrelated)
+ ids = {change.record.identity: change.record.id for change in saved.changes}
+ service = query_service()
+ filters = RecordFilters(parent_record_id=ids[parent.identity])
+ # Exercise the actual query-string boundary with explicit fixture authentication.
+ from types import SimpleNamespace
+
+ from fastapi import FastAPI
+ from httpx import ASGITransport, AsyncClient
+
+ from airweave.api import deps
+ from airweave.api.v1.endpoints.records import router
+ from airweave.db.session import get_db
+
+ app = FastAPI()
+ app.include_router(router, prefix="/sync")
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: SimpleNamespace(
+ organization=SimpleNamespace(id=fence.organization_id)
+ )
+ app.dependency_overrides[deps.get_canonical_query_service] = lambda: service
+ async with AsyncClient(transport=ASGITransport(app), base_url="http://fixture") as client:
+ response = await client.get(
+ f"/sync/{fence.sync_id}/records", params={"parent_record_id": str(ids[parent.identity])}
+ )
+ assert response.status_code == 200
+ assert {item["id"] for item in response.json()["records"]} == {
+ str(ids[child.identity]) for child in children
+ }
+ async with database() as db:
+ first = await service.list_records(
+ db, fence.organization_id, fence.sync_id, RecordListQuery(filters=filters, limit=1)
+ )
+ assert first.has_more and first.records[0].parent == parent.identity
+ with pytest.raises(InvalidRecordCursor):
+ await service.list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordListQuery(
+ filters=RecordFilters(parent_record_id=ids[other.identity]),
+ cursor=first.next_cursor,
+ ),
+ )
+ for organization, sync in ((uuid4(), fence.sync_id), (fence.organization_id, uuid4())):
+ with pytest.raises(RecordNotFound):
+ await service.list_records(db, organization, sync, RecordListQuery(filters=filters))
+ # Attachment inventory changes hide previously attested children on the next page.
+ await store(database, capture, fence, parent.model_copy(update={"payload": {"files": []}}))
+ async with database() as db:
+ page = await service.list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordListQuery(filters=filters, cursor=first.next_cursor),
+ )
+ assert not page.records
+ await store(
+ database,
+ capture,
+ fence,
+ parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ async with database() as db:
+ with pytest.raises(RecordNotFound):
+ await service.list_records(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ RecordListQuery(filters=filters, cursor=first.next_cursor),
+ )
diff --git a/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py
new file mode 100644
index 0000000000..9d17b5b346
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py
@@ -0,0 +1,220 @@
+"""Actual capture/blob/project/feed/publication/retrieval; fixed vectors, not quality."""
+
+import hashlib
+import os
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+from uuid import uuid4
+
+import httpx
+import pytest
+from fastapi import FastAPI
+from sqlalchemy import select, update
+from vespa.application import Vespa
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.api import deps
+from airweave.api.v1.endpoints.records import router
+from airweave.core.logging import logger
+from airweave.db.session import get_db
+from airweave.domains.converters.registry import ConverterRegistry
+from airweave.domains.embedders.fakes.embedder import FakeDenseEmbedder, FakeSparseEmbedder
+from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose, SparseEmbedding
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.domains.entities.canonical.tests.vespa_helpers import deploy_schema
+from airweave.domains.search.adapters.vector_db.filter_translator import FilterTranslator
+from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB
+from airweave.domains.search.executor import SearchPlanExecutor
+from airweave.domains.search.owned import OwnedSearchService
+from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor
+from airweave.models import Entity
+from airweave.models.collection import Collection
+from airweave.models.projection_generation import ProjectionGeneration
+from airweave.models.source_connection import SourceConnection
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+from airweave.platform.destinations.vespa.destination import VespaDestination
+from airweave.platform.sources.google_drive import GoogleDriveSource
+
+pytestmark = pytest.mark.skipif(
+ os.environ.get("OWNED_VESPA_TEST") != "1", reason="requires disposable real Vespa"
+)
+
+
+class FixedDense(FakeDenseEmbedder):
+ """Only model inference is substituted; all processing stays production code."""
+
+ async def embed_many(
+ self, texts: list[str], *, purpose: EmbeddingPurpose = "document"
+ ) -> list[DenseEmbedding]:
+ return [DenseEmbedding(vector=[1.0] + [0.0] * 383) for _ in texts]
+
+
+class FixedSparse(FakeSparseEmbedder):
+ async def embed(self, text):
+ return SparseEmbedding(indices=[1], values=[1.0])
+
+ async def embed_many(self, texts):
+ return [await self.embed(text) for text in texts]
+
+
+async def test_captured_blob_to_published_search_and_withdrawal(
+ database, source, tmp_path, monkeypatch
+):
+ """Search cannot expose another org, stale revision or withdrawn publication."""
+ from airweave.core.config import settings
+
+ monkeypatch.setattr(settings, "VESPA_URL", "http://localhost")
+ monkeypatch.setattr(settings, "VESPA_PORT", 8081)
+ async with httpx.AsyncClient(timeout=120) as http:
+ await deploy_schema(http)
+ service, fence = source
+ log = logger.with_context(request_id="synthetic-complete-projection")
+ storage = FilesystemBackend(tmp_path)
+ content = b"Synthetic fundraising discussion. The team plans a fundraising meeting."
+ digest = hashlib.sha256(content).hexdigest()
+ blob = BlobReference(
+ key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}",
+ sha256=digest,
+ size_bytes=len(content),
+ media_type="text/plain",
+ )
+ await storage.write_file(blob.key, content)
+ original = observation(
+ identity=RecordIdentity(record_type="file", native_id="fixture-file"),
+ payload={"id": "fixture-file", "name": "Fundraising.txt", "mimeType": "text/plain"},
+ blobs=(blob,),
+ )
+ await capture(database, service, fence, original)
+ async with database() as db:
+ deployment = VectorDbDeploymentMetadata(
+ dense_embedder="fixed", embedding_dimensions=384, sparse_embedder="fixed"
+ )
+ db.add(deployment)
+ await db.flush()
+ collection = Collection(
+ name="Synthetic",
+ readable_id="projection-" + uuid4().hex,
+ organization_id=fence.organization_id,
+ vector_db_deployment_metadata_id=deployment.id,
+ )
+ db.add(collection)
+ await db.flush()
+ connection = SourceConnection(
+ name="Synthetic",
+ short_name="google_drive",
+ organization_id=fence.organization_id,
+ readable_collection_id=collection.readable_id,
+ sync_id=fence.sync_id,
+ is_authenticated=True,
+ )
+ db.add(connection)
+ await db.commit()
+ dense, sparse = FixedDense(dimensions=384), FixedSparse()
+ store = CanonicalProjectionStore()
+ projector = CanonicalProjector(
+ store, database, ChunkEmbedProcessor(ConverterRegistry(), dense, sparse), storage
+ )
+ destination = await VespaDestination.create(
+ collection_id=collection.id, organization_id=fence.organization_id, logger=log
+ )
+ registry = SimpleNamespace(get=lambda _: SimpleNamespace(source_class_ref=GoogleDriveSource))
+ blocked = AsyncMock(side_effect=AssertionError("No provider or model-agent call allowed"))
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=log,
+ filter_translator=FilterTranslator(logger=log),
+ )
+ executor = SearchPlanExecutor(dense, sparse, engine, blocked, registry, blocked, blocked)
+ owned = OwnedSearchService(executor, registry)
+ ctx = SimpleNamespace(organization=SimpleNamespace(id=fence.organization_id), logger=log)
+ app = FastAPI()
+ app.include_router(router, prefix="/sync")
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: ctx
+ app.dependency_overrides[deps.get_owned_search_context] = lambda: ctx
+ app.dependency_overrides[deps.get_search_session_factory] = lambda: database
+ app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(owned_search=owned)
+ request = {"query": "fundraising", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}
+ try:
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app), base_url="http://test"
+ ) as api:
+ before = await api.post("/sync/search", json=request)
+ assert before.status_code == 200 and before.json()["items"] == []
+ assert before.json()["sources"][0]["pending_records"] == 1
+ batch = await projector.batch(
+ fence.organization_id, fence.sync_id, "google_drive", destination, log
+ )
+ assert batch.published == 1 and batch.failed == 0
+ async with database() as db:
+ row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id))
+ assert row.indexed_generation is not None and row.indexed_chunk_count > 0
+ record_id = str(row.id)
+ manifest = await db.scalar(select(ProjectionGeneration))
+ assert manifest.documents
+ for mode in ("keyword", "semantic", "hybrid"):
+ result = await api.post("/sync/search", json={**request, "mode": mode})
+ assert result.status_code == 200, result.text
+ assert [hit["record_id"] for hit in result.json()["items"]] == [record_id]
+ excerpt = " ".join(result.json()["items"][0]["excerpts"])
+ # Lexical summaries may select the matching filename rather than
+ # the body sentence. Semantic mode retains the full chunk fallback.
+ assert "fundraising" in excerpt.casefold()
+ assert "" not in excerpt and "" not in excerpt
+ if mode == "semantic":
+ assert "fundraising meeting" in excerpt
+ assert result.json()["sources"][0]["pending_records"] == 0
+ changed = original.model_copy(
+ update={"payload": {**original.payload, "name": "Updated fundraising.txt"}}
+ )
+ await capture(database, service, fence, changed)
+ stale = await api.post("/sync/search", json=request)
+ assert stale.status_code == 200 and stale.json()["items"] == []
+ assert stale.json()["excluded_candidates"] > 0
+ replacement = await projector.batch(
+ fence.organization_id, fence.sync_id, "google_drive", destination, log
+ )
+ assert replacement.published == 1 and replacement.failed == 0
+ current = await api.post("/sync/search", json=request)
+ assert current.status_code == 200
+ assert len(current.json()["items"]) == 1
+ assert current.json()["items"][0]["revision"] == 2
+ ctx.organization.id = uuid4()
+ assert (await api.post("/sync/search", json=request)).status_code == 404
+ ctx.organization.id = fence.organization_id
+ async with database() as db:
+ await db.execute(
+ update(SourceConnection)
+ .where(SourceConnection.id == connection.id)
+ .values(is_authenticated=False)
+ )
+ await db.commit()
+ assert (await api.post("/sync/search", json=request)).status_code == 404
+ async with database() as db:
+ await db.execute(
+ update(SourceConnection)
+ .where(SourceConnection.id == connection.id)
+ .values(is_authenticated=True)
+ )
+ await db.commit()
+ await capture(
+ database,
+ service,
+ fence,
+ changed.model_copy(update={"kind": "delete", "removal_reason": "provider_deleted"}),
+ )
+ # Old physical Vespa documents still exist: SQL authority must hide them now.
+ withdrawn = await api.post("/sync/search", json=request)
+ assert withdrawn.status_code == 200 and withdrawn.json()["items"] == []
+ assert withdrawn.json()["excluded_candidates"] > 0
+ finally:
+ await destination.delete_by_sync_id(fence.sync_id)
+ await destination.close_connection()
diff --git a/backend/airweave/domains/entities/canonical/tests/test_record_forest.py b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py
new file mode 100644
index 0000000000..2a7c0e1d06
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py
@@ -0,0 +1,405 @@
+"""Real PostgreSQL proof of record-forest visibility, independent of provider type depth."""
+
+import pytest
+from sqlalchemy import select, text
+
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalStoreError, content_is_available
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.models.entity import Entity
+
+
+def node(name, parent=None):
+ return observation(identity=RecordIdentity(record_type="block", native_id=name), parent=parent)
+
+
+async def available(database, ids):
+ async with database() as db:
+ return dict(
+ (
+ await db.execute(
+ select(Entity.id, content_is_available()).where(Entity.id.in_(ids))
+ )
+ ).all()
+ )
+
+
+async def test_nested_loss_revival_requires_fresh_attestation_without_changing_bytes(
+ database, source
+):
+ service, fence = source
+ root = node("root")
+ child = node("child", root.identity)
+ leaf = node("leaf", child.identity)
+ result = await capture(database, service, fence, root, child, leaf)
+ ids = [change.record.id for change in result.changes]
+ assert all((await available(database, ids)).values())
+ await capture(
+ database,
+ service,
+ fence,
+ root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ assert not any((await available(database, ids)).values())
+ await capture(database, service, fence, root)
+ assert list((await available(database, ids)).values()).count(True) == 1
+ refreshed = await capture(database, service, fence, child)
+ assert refreshed.changes[0].record.payload == child.payload
+ assert refreshed.changes[0].record.revision == 2
+ assert (await available(database, ids))[ids[2]] is False
+ await capture(database, service, fence, leaf)
+ assert all((await available(database, ids)).values())
+ async with database() as db:
+ epochs = list((await db.scalars(select(Entity.visibility_epoch))).all())
+ assert epochs == [2, 2, 2]
+ unchanged = await capture(database, service, fence, child)
+ assert unchanged.unchanged == 1
+ await capture(
+ database, service, fence, root.model_copy(update={"payload": {"new": "metadata"}})
+ )
+ assert all((await available(database, ids)).values())
+ async with database() as db:
+ assert (await db.get(Entity, ids[0])).visibility_epoch == 2
+
+
+async def test_record_cycle_rejected_and_transaction_rolls_back(database, source):
+ service, fence = source
+ root = node("root")
+ child = node("child", root.identity)
+ result = await capture(database, service, fence, root, child)
+ with pytest.raises(CanonicalStoreError, match="cycle"):
+ await capture(database, service, fence, root.model_copy(update={"parent": child.identity}))
+ assert all((await available(database, [item.record.id for item in result.changes])).values())
+ with pytest.raises(CanonicalStoreError, match="cycle"):
+ await capture(
+ database,
+ service,
+ fence,
+ root.model_copy(
+ update={
+ "parent": child.identity,
+ "kind": "delete",
+ "removal_reason": "provider_deleted",
+ }
+ ),
+ )
+ with pytest.raises(CanonicalStoreError, match="parent itself"):
+ await capture(database, service, fence, root.model_copy(update={"parent": root.identity}))
+
+
+async def test_missing_parent_rejected_and_corrupt_record_cycle_fails_closed(database, source):
+ service, fence = source
+ with pytest.raises(CanonicalStoreError, match="must exist"):
+ await capture(
+ database,
+ service,
+ fence,
+ node("orphan", RecordIdentity(record_type="block", native_id="missing")),
+ )
+ root = node("root")
+ child = node("child", root.identity)
+ result = await capture(database, service, fence, root, child)
+ async with database() as db:
+ row = await db.get(Entity, result.changes[0].record.id)
+ row.parent_record_type = "block"
+ row.parent_native_id = "child"
+ row.parent_visibility_epoch = 1
+ await db.commit()
+ async with database() as db:
+ await db.execute(text("SET LOCAL statement_timeout='2s'"))
+ assert not any(
+ (await db.execute(select(content_is_available()).select_from(Entity))).scalars().all()
+ )
+
+
+async def test_deep_same_kind_forest_has_no_twenty_level_cutoff(database, source):
+ service, fence = source
+ records = [node("0")]
+ for index in range(1, 35):
+ records.append(node(str(index), records[-1].identity))
+ result = await capture(database, service, fence, *records)
+ assert all((await available(database, [item.record.id for item in result.changes])).values())
+
+
+@pytest.mark.parametrize("database", ["forest_upgrade"], indirect=True)
+async def test_upgrade_attests_only_existing_available_flat_links(database):
+ async with database() as db:
+ rows = list((await db.scalars(select(Entity).where(Entity.record_revision > 0))).all())
+ states = dict(
+ (
+ await db.execute(
+ select(Entity.native_id, content_is_available()).where(
+ Entity.record_revision > 0
+ )
+ )
+ ).all()
+ )
+ assert states == {
+ "active-root": True,
+ "active-child": True,
+ "withdrawn-root": False,
+ "withdrawn-child": False,
+ "hidden-child": False,
+ "orphan": False,
+ }
+ assert next(row for row in rows if row.native_id == "active-child").parent_visibility_epoch == 1
+ assert all(
+ row.source_payload == {"native": row.native_id} and row.record_revision == 1 for row in rows
+ )
+
+
+async def test_cleanup_propagates_ancestor_reason_even_before_intermediate_cleanup(
+ database, source
+):
+ service, fence = source
+ root = node("root")
+ child = node("child", root.identity)
+ leaf = node("leaf", child.identity)
+ await capture(database, service, fence, root, child, leaf)
+ await capture(
+ database,
+ service,
+ fence,
+ root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ async with database() as db:
+ removed = await service.reconcile_parents(db, fence)
+ assert len(removed.capture.changes) == 2
+ assert {change.record.removal_reason for change in removed.capture.changes} == {
+ "access_revoked"
+ }
+
+
+async def test_reparent_invalidates_descendants_and_failed_batch_rolls_back(database, source):
+ service, fence = source
+ first, second = node("first"), node("second")
+ child = node("child", first.identity)
+ leaf = node("leaf", child.identity)
+ result = await capture(database, service, fence, first, second, child, leaf)
+ leaf_id = result.changes[-1].record.id
+ moved = child.model_copy(update={"parent": second.identity, "allow_reparent": True})
+ with pytest.raises(CanonicalStoreError, match="different parent"):
+ await capture(database, service, fence, moved.model_copy(update={"allow_reparent": False}))
+ assert (await available(database, [leaf_id]))[leaf_id] is True
+ with pytest.raises(CanonicalStoreError):
+ await capture(database, service, fence, moved, node("bad", node("missing").identity))
+ assert (await available(database, [leaf_id]))[leaf_id] is True
+ await capture(database, service, fence, moved)
+ replay = await capture(
+ database, service, fence, moved.model_copy(update={"allow_reparent": False})
+ )
+ assert replay.unchanged == 1 and replay.changes == ()
+ assert (await available(database, [leaf_id]))[leaf_id] is False
+ await capture(database, service, fence, leaf)
+ assert (await available(database, [leaf_id]))[leaf_id] is True
+
+
+@pytest.mark.parametrize("shape,size", [("deep", 96), ("wide", 1024)])
+async def test_real_recursive_query_plan(shape, size, database, source):
+ import json
+
+ from sqlalchemy.dialects import postgresql
+
+ service, fence = source
+ records = [node("root")]
+ for index in range(size):
+ parent = records[-1] if shape == "deep" else records[0]
+ records.append(node(str(index), parent.identity))
+ for offset in range(0, len(records), 500):
+ await capture(database, service, fence, *records[offset : offset + 500])
+ query = select(Entity.id).where(
+ Entity.organization_id == fence.organization_id,
+ Entity.sync_id == fence.sync_id,
+ content_is_available(),
+ )
+ compiled = query.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})
+ async with database() as db:
+ await db.execute(text("ANALYZE entity"))
+ plan = await db.scalar(text("EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) " + str(compiled)))
+ assert plan[0]["Plan"]["Actual Rows"] == size + 1
+ assert "Recursive Union" in json.dumps(plan)
+ print(
+ json.dumps(
+ {"forest_shape": shape, "records": size + 1, "execution_ms": plan[0]["Execution Time"]}
+ )
+ )
+
+
+@pytest.mark.parametrize("withdrawal", ["delete", "access_change"])
+async def test_nested_withdrawal_gates_original_blob_history_search_and_publication(
+ database, source, withdrawal
+):
+ from types import SimpleNamespace
+ from unittest.mock import AsyncMock
+ from uuid import uuid4
+
+ from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+ from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+ from airweave.domains.entities.canonical.query import CanonicalQueryService, RecordNotFound
+ from airweave.domains.entities.canonical.query_models import RecordFilters
+ from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+ from airweave.domains.entities.canonical.requests import BlobReference
+ from airweave.domains.entities.canonical.tests.helpers import publish_prepared
+ from airweave.domains.entities.canonical.tests.test_search_visibility import hit
+ from airweave.domains.search.canonical_visibility import visible_results
+ from airweave.models.collection import Collection
+ from airweave.models.source_connection import SourceConnection
+ from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+ service, fence = source
+ root = node("root").model_copy(
+ update={"payload": {"accessRole": "owner"}, "descendant_visibility_fields": ("accessRole",)}
+ )
+ child = node("child", root.identity)
+ leaf = node("leaf", child.identity).model_copy(
+ update={"blobs": (BlobReference(key="not-read", sha256="a" * 64, size_bytes=1),)}
+ )
+ captured = await capture(database, service, fence, root, child, leaf)
+ leaf_id = captured.changes[-1].record.id
+ async with database() as db:
+ deployment = VectorDbDeploymentMetadata(
+ dense_embedder="test", embedding_dimensions=3, sparse_embedder="test"
+ )
+ db.add(deployment)
+ await db.flush()
+ db.add(
+ Collection(
+ name="Test",
+ readable_id="test",
+ organization_id=fence.organization_id,
+ vector_db_deployment_metadata_id=deployment.id,
+ )
+ )
+ await db.flush()
+ db.add(
+ SourceConnection(
+ name="Test",
+ short_name="gmail",
+ organization_id=fence.organization_id,
+ readable_collection_id="test",
+ sync_id=fence.sync_id,
+ is_authenticated=True,
+ )
+ )
+ await db.commit()
+ projection = CanonicalProjectionStore()
+ async with database() as db:
+ work = next(
+ item
+ for item in await projection.pending(db, fence.organization_id, fence.sync_id)
+ if item.record.id == leaf_id
+ )
+ generation = uuid4()
+ assert await publish_prepared(projection, db, work, generation, 1)
+ locator = ProjectionLocator(
+ record_id=leaf_id, revision=1, pipeline_version=1, generation=generation, part_index=0
+ )
+ candidate = hit(fence, locator.encode())
+ registry = SimpleNamespace(
+ get=lambda _: SimpleNamespace(
+ source_class_ref=SimpleNamespace(canonical_record_types=("block",))
+ )
+ )
+ async with database() as db:
+ assert await visible_results(db, fence.organization_id, "test", [candidate], registry) == [
+ candidate
+ ]
+ await capture(
+ database,
+ service,
+ fence,
+ root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"})
+ if withdrawal == "delete"
+ else root.model_copy(update={"payload": {"accessRole": "reader"}}),
+ )
+ query = CanonicalQueryService(service.store, CanonicalQueryStore(), "test-signing-key")
+ storage = AsyncMock()
+ async with database() as db:
+ exact = await query.read(db, fence.organization_id, fence.sync_id, leaf_id)
+ assert exact.content_access == "unavailable" and exact.payload == {} and not exact.blobs
+ listed = await query.queries.list_records(
+ db, fence.organization_id, fence.sync_id, RecordFilters(), after_id=None, limit=100
+ )
+ assert {item.identity.native_id for item in listed} == (
+ {"root"} if withdrawal == "access_change" else set()
+ )
+ history = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert all(
+ item.record.content_access == "unavailable"
+ for item in history.changes
+ if withdrawal == "delete" or item.record.identity.native_id != "root"
+ )
+ with pytest.raises(RecordNotFound):
+ await query.blob(
+ db, fence.organization_id, fence.sync_id, leaf_id, 1, "a" * 64, storage
+ )
+ assert not storage.mock_calls
+ assert not await visible_results(db, fence.organization_id, "test", [candidate], registry)
+ assert not await projection.publish(db, work, generation, 1)
+
+ if withdrawal == "access_change":
+ await capture(database, service, fence, child)
+ restored = await capture(database, service, fence, leaf)
+ assert restored.changes[0].record.id == leaf_id
+ async with database() as db:
+ exact = await query.read(db, fence.organization_id, fence.sync_id, leaf_id)
+ assert exact.content_access == "available" and exact.blobs
+
+
+async def test_access_field_transition_is_atomic_and_not_metadata_revision(database, source):
+ from airweave.domains.entities.canonical.store import capture_fingerprint
+
+ service, fence = source
+ root = node("root").model_copy(
+ update={"payload": {"role": "owner"}, "descendant_visibility_fields": ("role",)}
+ )
+ child = node("child", root.identity)
+ initial = await capture(database, service, fence, root, child)
+ root_id, child_id = [item.record.id for item in initial.changes]
+ assert capture_fingerprint(root) == capture_fingerprint(
+ root.model_copy(update={"descendant_visibility_fields": ()})
+ )
+ metadata = root.model_copy(update={"payload": {"role": "owner", "summary": "Renamed"}})
+ await capture(database, service, fence, metadata)
+ async with database() as db:
+ assert (await db.get(Entity, root_id)).visibility_epoch == 1
+ downgrade = metadata.model_copy(update={"payload": {"role": "reader"}})
+ with pytest.raises(CanonicalStoreError):
+ await capture(database, service, fence, downgrade, node("orphan", node("missing").identity))
+ assert (await available(database, [child_id]))[child_id]
+ await capture(database, service, fence, downgrade)
+ assert not (await available(database, [child_id]))[child_id]
+ unchanged = await capture(database, service, fence, downgrade)
+ assert unchanged.unchanged == 1
+ async with database() as db:
+ assert (await db.get(Entity, root_id)).visibility_epoch == 2
+ missing = downgrade.model_copy(update={"payload": {}})
+ await capture(database, service, fence, missing)
+ async with database() as db:
+ assert (await db.get(Entity, root_id)).visibility_epoch == 3
+
+
+async def test_first_access_context_capture_of_legacy_null_payload(database, source):
+ service, fence = source
+ root = node("legacy").model_copy(
+ update={"payload": {"role": "reader"}, "descendant_visibility_fields": ("role",)}
+ )
+ async with database() as db:
+ row = Entity(
+ organization_id=fence.organization_id,
+ sync_id=fence.sync_id,
+ sync_job_id=fence.job_id,
+ entity_id=root.identity.entity_key,
+ entity_definition_short_name="block",
+ hash="legacy",
+ source_payload=None,
+ )
+ db.add(row)
+ await db.commit()
+ record_id = row.id
+ result = await capture(database, service, fence, root)
+ assert result.changes[0].record.id == record_id
+ assert result.changes[0].record.revision == 1
+ async with database() as db:
+ row = await db.get(Entity, record_id)
+ assert row.visibility_epoch == 2 and row.source_payload == {"role": "reader"}
diff --git a/backend/airweave/domains/entities/canonical/tests/test_reprojection.py b/backend/airweave/domains/entities/canonical/tests/test_reprojection.py
new file mode 100644
index 0000000000..11e9aff798
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_reprojection.py
@@ -0,0 +1,100 @@
+"""Real SQL operator scope, dry-run and retry/version fencing."""
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select, update
+
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.reprojection import plan_reprojection
+from airweave.domains.entities.canonical.tests.test_owned_search import (
+ indexed as search_indexed, # noqa: F401, F811
+)
+from airweave.models.sync import Sync
+
+
+async def test_dry_run_then_apply_and_identical_resume(database, search_indexed): # noqa: F811
+ fence, _, _ = search_indexed
+ kwargs = {"expected_version": 2, "target_version": 3}
+ async with database() as db:
+ plan = await plan_reprojection(db, fence.organization_id, fence.sync_id, **kwargs)
+ assert not plan.applied and plan.target_pending_records == 1
+ await db.commit()
+ async with database() as db:
+ assert await db.scalar(select(Sync.index_pipeline_version)) == 2
+ applied = await plan_reprojection(
+ db, fence.organization_id, fence.sync_id, apply=True, **kwargs
+ )
+ await db.commit()
+ async with database() as db:
+ resumed = await plan_reprojection(
+ db, fence.organization_id, fence.sync_id, apply=True, **kwargs
+ )
+ assert resumed.current_version == 3 and resumed.workflow_id == applied.workflow_id
+ await db.commit()
+ async with database() as db:
+ pending = await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id)
+ assert len(pending) == 1 and pending[0].pipeline_version == 3
+
+
+async def test_wrong_scope_and_changed_version_fail(database, search_indexed): # noqa: F811
+ fence, _, _ = search_indexed
+ async with database() as db:
+ with pytest.raises(ValueError, match="unavailable"):
+ await plan_reprojection(
+ db, uuid4(), fence.sync_id, expected_version=2, target_version=3
+ )
+ await db.rollback()
+ await db.execute(update(Sync).values(index_pipeline_version=4))
+ await db.commit()
+ async with database() as db:
+ with pytest.raises(ValueError, match="changed"):
+ await plan_reprojection(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ expected_version=2,
+ target_version=3,
+ apply=True,
+ )
+
+
+async def test_launch_failure_keeps_queue_and_identical_retry_starts_once(
+ database,
+ search_indexed, # noqa: F811
+ monkeypatch,
+):
+ from argparse import Namespace
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.entities.canonical.projection_store import publication_matches
+ from airweave.models.entity import Entity
+ from scripts import reproject_canonical
+
+ fence, locator, _ = search_indexed
+ args = Namespace(
+ organization=fence.organization_id, sync=fence.sync_id, expect=2, target=3, apply=True
+ )
+ monkeypatch.setattr(reproject_canonical, "AsyncSessionLocal", database)
+ monkeypatch.setattr(
+ reproject_canonical, "get_client", AsyncMock(side_effect=RuntimeError("offline"))
+ )
+ with pytest.raises(RuntimeError, match="offline"):
+ await reproject_canonical.run(args)
+ async with database() as db:
+ assert await db.scalar(select(Sync.index_pipeline_version)) == 3
+ visible = await db.scalar(
+ select(Entity.id)
+ .join(Sync, Sync.id == Entity.sync_id)
+ .where(publication_matches(locator))
+ )
+ assert visible is None
+ client = AsyncMock()
+ monkeypatch.setattr(reproject_canonical, "get_client", AsyncMock(return_value=client))
+ await reproject_canonical.run(args)
+ assert client.start_workflow.await_count == 1
+ assert client.start_workflow.call_args.kwargs["id"] == (
+ f"canonical-projection:{fence.organization_id}:{fence.sync_id}"
+ )
+ async with database() as db:
+ assert await db.scalar(select(Sync.index_pipeline_version)) == 3
diff --git a/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py
new file mode 100644
index 0000000000..a52865796e
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py
@@ -0,0 +1,253 @@
+"""Local subprocess verification for the opt-in live harness, with no provider credentials."""
+
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+
+BACKEND = Path(__file__).resolve().parents[5]
+
+
+async def subprocess_script(script, environment):
+ process = await asyncio.create_subprocess_exec(
+ sys.executable,
+ "-c",
+ script,
+ cwd=BACKEND,
+ env={**os.environ, **environment, "PYTHONPATH": str(BACKEND)},
+ stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.PIPE,
+ )
+ try:
+ stdout, stderr = await asyncio.wait_for(process.communicate(), 30)
+ finally:
+ if process.returncode is None:
+ process.kill()
+ await process.wait()
+ return process.returncode, stdout.decode(), stderr.decode()
+
+
+PROCESS_SCRIPT = r"""
+import asyncio, json, os, sys
+from pathlib import Path
+sys.path.insert(0, "tests/live")
+import conftest
+from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
+from airweave.domains.entities.canonical.requests import WriterFence, CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.test_slack_recovery import (
+ runner, run, ROOT, HISTORY, REPLIES,
+)
+from capture_resume import PageResumeProbe, inspect_resume_progress
+from provider_lifecycle import BoundedPageSource
+
+async def main():
+ engine = create_async_engine(os.environ["CANONICAL_TEST_DATABASE_URL"],
+ connect_args={"server_settings": {"search_path": os.environ["TEST_SCHEMA"]}})
+ sessions = async_sessionmaker(engine, expire_on_commit=False)
+ fence = WriterFence.model_validate_json(os.environ["TEST_FENCE"])
+ mode = os.environ["TEST_MODE"]
+ service = CanonicalCaptureService(CanonicalRecordStore())
+ from airweave.platform.sources.slack import SlackApiError
+ retry = mode in ("resume", "withdraw", "expired")
+ responses = [ROOT, REPLIES] if retry else [ROOT, HISTORY, REPLIES]
+ if mode == "withdraw":
+ responses = [ROOT, SlackApiError("not_in_channel")]
+ if mode == "expired":
+ responses = [ROOT, SlackApiError("invalid_cursor"), HISTORY, REPLIES]
+ instance, connector, pipeline = await runner(sessions, (service, fence), responses,
+ attempt=2 if retry else 1)
+ counters = {"provider_requests": 0, "records_observed": 0}
+ get = connector._get
+ async def counted_get(*args):
+ counters["provider_requests"] += 1
+ return await get(*args)
+ connector._get = counted_get
+ probe = PageResumeProbe(sessions, fence.organization_id, fence.sync_id, fence.job_id,
+ pipeline._attempt.id, Path(os.environ["TEST_TARGET"]),
+ "resume" if retry else "interrupt", counters)
+ if mode == "uncommitted":
+ await probe.before_page(CompletedScope(record_type="message", container_id="C1"),
+ ScanContinuation(value={"pending_threads": ["1"]}))
+ raise AssertionError("Uncommitted progress was accepted")
+ pipeline.page_source = BoundedPageSource(connector, counters, 100, probe)
+ await run(instance)
+ evidence = await inspect_resume_progress(sessions, fence.organization_id, fence.sync_id,
+ fence.job_id, Path(os.environ["TEST_TARGET"]))
+ print(json.dumps({**evidence, "resumed": probe.resumed_saved_page,
+ "expired": probe.saved_cursor_expired, **counters}), flush=True)
+ await engine.dispose()
+asyncio.run(main())
+"""
+
+
+async def test_process_loss_after_commit_then_fresh_same_job_resume(database, source, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ environment = {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(tmp_path / "target.json"),
+ "TEST_MODE": "uncommitted",
+ }
+ code, stdout, stderr = await subprocess_script(PROCESS_SCRIPT, environment)
+ assert code != 75 and "Durable interruption state is missing" in stderr
+ assert not (tmp_path / "target.json").exists()
+ code, stdout, stderr = await subprocess_script(
+ PROCESS_SCRIPT, {**environment, "TEST_MODE": "interrupt"}
+ )
+ assert code == 75, stderr
+ stopped = json.loads(stdout.splitlines()[-1])
+ assert stopped["durable_pending_threads_verified"] and stopped["job_status"] == "running"
+ assert stopped["provider_requests"] == 2 and stopped["records_observed"] == 2
+ assert stopped["full_scope_completed"] is False
+ assert (tmp_path / "target.json").stat().st_mode & 0o077 == 0
+ code, stdout, stderr = await subprocess_script(
+ PROCESS_SCRIPT, {**environment, "TEST_MODE": "resume"}
+ )
+ assert code == 0, stderr
+ resumed = json.loads(stdout.splitlines()[-1])
+ assert resumed == {
+ "resumed_page_committed": True,
+ "resumed": True,
+ "expired": False,
+ "provider_requests": 2,
+ "records_observed": 3,
+ }
+
+
+@pytest.mark.parametrize("mode", ["withdraw", "expired"])
+async def test_withdrawal_or_restarted_sweep_does_not_prove_saved_reply_commit(
+ database, source, tmp_path, mode
+):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ environment = {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(tmp_path / "target.json"),
+ "TEST_MODE": "interrupt",
+ }
+ code, _, stderr = await subprocess_script(PROCESS_SCRIPT, environment)
+ assert code == 75, stderr
+ code, stdout, stderr = await subprocess_script(
+ PROCESS_SCRIPT, {**environment, "TEST_MODE": mode}
+ )
+ assert code == 0, stderr
+ evidence = json.loads(stdout.splitlines()[-1])
+ assert evidence["resumed_page_committed"] is False
+ assert evidence["expired"] is (mode == "expired")
+
+
+BUDGET_SCRIPT = r"""
+import asyncio, json, os, sys
+from pathlib import Path
+sys.path.insert(0, "tests/live")
+import conftest
+import provider_lifecycle as lifecycle
+
+async def main():
+ snapshots = []
+ outputs = [
+ {"failed": False, "intentional_interruption": True,
+ "provider_requests": 2, "records_observed": 5},
+ {"failed": False, "full_scope_completed": True, "resumed_saved_page": True,
+ "provider_requests": 3, "records_observed": 7},
+ {"failed": False, "full_scope_completed": True,
+ "provider_requests": 1, "records_observed": 2},
+ ]
+ killed = []
+ class Process:
+ def __init__(self, index):
+ self.index = index
+ self.returncode = (None if (os.environ["TEST_MODE"] == "timeout" or
+ (os.environ["TEST_MODE"] == "partial" and index == 1))
+ else (75 if index == 0 else 0))
+ async def communicate(self):
+ if self.returncode is None:
+ raise asyncio.TimeoutError()
+ return json.dumps(outputs[self.index]).encode(), b""
+ def kill(self):
+ killed.append(True)
+ self.returncode = -9
+ async def wait(self):
+ return self.returncode
+ async def spawn(*args, **kwargs):
+ snapshots.append(json.loads(Path(args[2]).read_text()))
+ return Process(len(snapshots)-1)
+ lifecycle.asyncio.create_subprocess_exec = spawn
+ async def evidence():
+ return {"resumed_page_committed": True}
+ result = await lifecycle.run_trials({"provider": "slack", "root": os.environ["TEST_ROOT"],
+ "timeout": 100, "request_limit": 20, "record_limit": 50},
+ evidence)
+ if os.environ["TEST_MODE"] in ("timeout", "partial"):
+ assert len(snapshots) == (2 if os.environ["TEST_MODE"] == "partial" else 1) and killed
+ if os.environ["TEST_MODE"] == "partial":
+ assert result[-1]["resumed_partial"] is True
+ assert result[-1]["full_scope_completed"] is False
+ summary = lifecycle.trial_counter_summary(result)
+ assert summary["aggregate_counters_complete"] is False
+ assert summary["aggregate_provider_requests"] is None
+ assert summary["aggregate_records_observed"] is None
+ else:
+ assert [item["request_limit"] for item in snapshots] == [20,18,15]
+ assert [item["record_limit"] for item in snapshots] == [50,45,38]
+ assert snapshots[0]["job_id"] == snapshots[1]["job_id"] != snapshots[2]["job_id"]
+ assert [item["attempt_number"] for item in snapshots] == [1,2,1]
+ assert 100 >= snapshots[0]["timeout"] >= snapshots[1]["timeout"] >= snapshots[2]["timeout"]
+ assert lifecycle.trial_counter_summary(result)["aggregate_provider_requests"] == 6
+ print("verified")
+asyncio.run(main())
+"""
+
+
+async def test_parent_trial_keeps_aggregate_budget_and_marks_killed_usage_unknown(tmp_path):
+ for mode in ("normal", "timeout", "partial"):
+ code, stdout, stderr = await subprocess_script(
+ BUDGET_SCRIPT,
+ {
+ "TEST_MODE": mode,
+ "TEST_ROOT": str(tmp_path),
+ },
+ )
+ assert code == 0, stderr
+ assert stdout.splitlines()[-1] == "verified"
+
+
+async def test_trial_summary_distinguishes_observation_from_full_scope_evidence():
+ script = r"""
+import sys
+sys.path.insert(0, 'tests/live')
+import provider_lifecycle as lifecycle
+first = dict(failed=False, full_scope_completed=True, payload_revision_digest='same')
+second = dict(failed=False, full_scope_completed=False, payload_revision_digest='same',
+ loaded_durable_checkpoint=True, sync_token_requests=1)
+assert lifecycle.trial_read_summary([first, second]) == {
+ 'saved_cursor_reused': True, 'identical_second_read': True}
+for field in ('sync_token_requests', 'resumed_changes_requests', 'resumed_saved_page'):
+ delta = {**second, 'sync_token_requests': 0, field: 1}
+ assert lifecycle.trial_read_summary([first, delta])['saved_cursor_reused']
+result = lifecycle.trial_read_summary([first, {**second, 'sync_token_requests': 0}])
+assert not result['saved_cursor_reused']
+result = lifecycle.trial_read_summary([first, {**second, 'saved_cursor_expired': True}])
+assert not result['saved_cursor_reused']
+result = lifecycle.trial_read_summary([first, {**second, 'failed': True}])
+assert not result['identical_second_read']
+result = lifecycle.trial_read_summary(
+ [first, {**second, 'payload_revision_digest': 'changed'}])
+assert not result['identical_second_read']
+result = lifecycle.trial_read_summary([{'failed': False}, {'failed': False}])
+assert not result['identical_second_read']
+assert not lifecycle.trial_read_summary([])['identical_second_read']
+print('verified')
+"""
+ code, stdout, stderr = await subprocess_script(script, {})
+ assert code == 0, stderr
+ assert stdout.strip().endswith("verified")
diff --git a/backend/airweave/domains/entities/canonical/tests/test_retained_text.py b/backend/airweave/domains/entities/canonical/tests/test_retained_text.py
new file mode 100644
index 0000000000..75d3f944fe
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_retained_text.py
@@ -0,0 +1,289 @@
+"""Real PDF conversion/SQL/blob reads; fixed chunking/embedding, no provider calls."""
+
+import hashlib
+from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import fitz
+import httpx
+import pytest
+from fastapi import FastAPI
+from sqlalchemy import select, update
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.api import deps
+from airweave.api.v1.endpoints.records import record_error_response, router
+from airweave.db.session import get_db
+from airweave.domains.converters.registry import ConverterRegistry
+from airweave.domains.embedders.fakes.embedder import FakeDenseEmbedder, FakeSparseEmbedder
+from airweave.domains.entities.canonical.projection_gc import ProjectionGCStore
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.query import CanonicalQueryService
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError
+from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation
+from airweave.domains.entities.canonical.text_models import TextArtifact
+from airweave.domains.entities.canonical.text_query import CanonicalTextReader
+from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor
+from airweave.models.projection_generation import ProjectionGeneration
+from airweave.models.sync import Sync
+from airweave.platform.destinations.vespa.transformer import EntityTransformer
+
+
+@pytest.fixture
+async def publication(database, source, tmp_path):
+ service, fence = source
+ binding = await bind_projection(database, fence, "google_drive")
+ storage = FilesystemBackend(tmp_path)
+ pdf = fitz.open()
+ for index in range(3):
+ page = pdf.new_page()
+ page.insert_textbox(
+ fitz.Rect(40, 40, 560, 780),
+ (
+ f"Page content {index}. "
+ + "Retained complete document words. " * 150
+ + "\n# Content\nThis marker belongs to the source, not our metadata."
+ ),
+ fontsize=9,
+ )
+ raw = pdf.tobytes()
+ pdf.close()
+ digest = hashlib.sha256(raw).hexdigest()
+ blob = BlobReference(
+ key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}",
+ sha256=digest,
+ size_bytes=len(raw),
+ media_type="application/pdf",
+ )
+ await storage.write_file(blob.key, raw)
+ parent = observation(
+ identity=RecordIdentity(record_type="folder", native_id="parent"),
+ payload={"id": "parent", "name": "Folder"},
+ )
+ original = observation(
+ identity=RecordIdentity(record_type="file", native_id="pdf"),
+ parent=parent.identity,
+ payload={"id": "pdf", "name": "Complete.pdf", "mimeType": "application/pdf"},
+ blobs=(blob,),
+ )
+ await capture(database, service, fence, parent, original)
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = next(
+ item
+ for item in await store.pending(db, fence.organization_id, fence.sync_id)
+ if item.record.identity.record_type == "file"
+ )
+ registry = ConverterRegistry()
+ converter = registry.for_extension(".pdf")
+ converter.convert_batch = AsyncMock(wraps=converter.convert_batch)
+ processor = ChunkEmbedProcessor(registry, FakeDenseEmbedder(), FakeSparseEmbedder())
+
+ async def fixed_chunks(entities, context, runtime):
+ return processor._multiply_entities(
+ entities, [[{"text": entity.textual_representation}] for entity in entities], context
+ )
+
+ processor._chunk_entities = fixed_chunks
+ destination = MagicMock(collection_id=binding.collection_id, feed_prepared=AsyncMock())
+ destination.prepare_documents = lambda chunks: {
+ "base_entity": [
+ EntityTransformer(collection_id=destination.collection_id).transform(chunk)
+ for chunk in chunks
+ ]
+ }
+ records = CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-key")
+ reader = CanonicalTextReader(records, storage)
+ projector = CanonicalProjector(store, database, processor, storage)
+ return SimpleNamespace(
+ service=service,
+ fence=fence,
+ storage=storage,
+ work=work,
+ reader=reader,
+ projector=projector,
+ destination=destination,
+ converter=converter,
+ original=original,
+ parent=parent,
+ records=records,
+ blob=blob,
+ raw=raw,
+ )
+
+
+async def test_pdf_full_content_read_is_single_conversion_and_published_with_index(
+ database, publication
+):
+ p = publication
+
+ async def feed(documents):
+ async with database() as db:
+ manifest = await db.scalar(
+ select(ProjectionGeneration).where(
+ ProjectionGeneration.record_id == p.work.record.id
+ )
+ )
+ assert manifest.text_representations
+ listed = await p.reader.list(
+ db, p.fence.organization_id, p.fence.sync_id, p.work.record.id, 1
+ )
+ assert listed.status == "unavailable" # bytes exist but publication has not happened
+
+ p.destination.feed_prepared.side_effect = feed
+ assert await p.projector.project_one(p.work, "google_drive", p.destination, MagicMock())
+ p.converter.convert_batch.assert_awaited_once()
+ app = FastAPI()
+ app.include_router(router, prefix="/sync")
+ app.add_exception_handler(CanonicalStoreError, record_error_response)
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: SimpleNamespace(
+ organization=SimpleNamespace(id=p.fence.organization_id)
+ )
+ app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(
+ storage_backend=p.storage
+ )
+ app.dependency_overrides[deps.get_canonical_query_service] = lambda: p.records
+ base = f"/sync/{p.fence.sync_id}/records/{p.work.record.id}/text-representations"
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app), base_url="http://test"
+ ) as client:
+ listed = await client.get(base, params={"revision": 1})
+ assert listed.status_code == 200
+ descriptor = listed.json()["representations"][0]
+ assert (
+ descriptor["kind"] == "extracted_text" and descriptor["source_anchors"] == "unavailable"
+ )
+ params = {"revision": 1, "generation": descriptor["generation"], "limit": 100000}
+ result = await client.get(f"{base}/{descriptor['id']}", params=params)
+ assert result.status_code == 200
+ text = result.json()["text"]
+ assert len(text) > 1000 and "Page content 2" in text and "marker belongs" in text
+ assert not text.startswith("# Metadata")
+ assert result.json()["next_offset"] is None
+ index = await client.get(f"{base}/{descriptor['id']}", params={**params, "view": "index"})
+ assert index.json()["text"].startswith("# Metadata")
+ assert index.json()["text"].endswith(text)
+ await capture(
+ database,
+ p.service,
+ p.fence,
+ p.parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ )
+ denied = await client.get(f"{base}/{descriptor['id']}", params=params)
+ assert denied.status_code == 404
+ assert await p.storage.read_file(p.blob.key) == p.raw
+
+
+async def test_blob_corruption_and_midread_revision_change_fail_closed(database, publication):
+ from airweave.domains.entities.canonical.query import BlobUnavailable, StaleRecordRevision
+
+ p = publication
+ assert await p.projector.project_one(p.work, "google_drive", p.destination, MagicMock())
+ async with database() as db:
+ descriptor = (
+ await p.reader.list(db, p.fence.organization_id, p.fence.sync_id, p.work.record.id, 1)
+ ).representations[0]
+ row = await db.get(ProjectionGeneration, descriptor.generation)
+ artifact = TextArtifact.model_validate(row.text_representations[0])
+ key = artifact.storage_key(p.fence.sync_id, descriptor.generation)
+ original_bytes = await p.storage.read_file(key)
+ await p.storage.write_file(key, b"corrupt")
+ args = (
+ p.fence.organization_id,
+ p.fence.sync_id,
+ p.work.record.id,
+ 1,
+ descriptor.generation,
+ descriptor.id,
+ )
+ async with database() as db:
+ with pytest.raises(BlobUnavailable):
+ await p.reader.read(db, *args)
+ await p.storage.write_file(key, original_bytes)
+ read_file = p.storage.read_file
+
+ async def pipeline_changes(path, **kwargs):
+ content = await read_file(path, **kwargs)
+ async with database() as db:
+ await db.execute(
+ update(Sync).where(Sync.id == p.fence.sync_id).values(index_pipeline_version=2)
+ )
+ await db.commit()
+ return content
+
+ from airweave.domains.entities.canonical.text_query import TextUnavailable
+
+ p.storage.read_file = pipeline_changes
+ async with database() as db:
+ with pytest.raises(TextUnavailable):
+ await p.reader.read(db, *args)
+ async with database() as db:
+ await db.execute(
+ update(Sync).where(Sync.id == p.fence.sync_id).values(index_pipeline_version=1)
+ )
+ await db.commit()
+
+ async def changing_read(path, **kwargs):
+ content = await read_file(path, **kwargs)
+ await capture(
+ database,
+ p.service,
+ p.fence,
+ p.original.model_copy(
+ update={"payload": {**p.original.payload, "name": "Changed.pdf"}}
+ ),
+ )
+ return content
+
+ p.storage.read_file = changing_read
+ async with database() as db:
+ with pytest.raises(StaleRecordRevision):
+ await p.reader.read(db, *args)
+
+
+async def test_abandoned_generation_text_is_in_bounded_repeatable_gc(database, publication):
+ p = publication
+ p.destination.feed_prepared.side_effect = RuntimeError("feed unavailable")
+ with pytest.raises(RuntimeError):
+ await p.projector.project_one(p.work, "google_drive", p.destination, MagicMock())
+ async with database() as db:
+ row = await db.scalar(
+ select(ProjectionGeneration).where(ProjectionGeneration.record_id == p.work.record.id)
+ )
+ generation = row.id
+ now = datetime.now(timezone.utc) + timedelta(hours=2)
+ gc = ProjectionGCStore()
+ async with database() as db:
+ first = await gc.claim(db, generation, now=now, limit=1)
+ assert len(first.documents) == 1 and not first.artifact_keys
+ await gc.acknowledge(db, first, now=now)
+ async with database() as db:
+ second = await gc.claim(db, generation, now=now, limit=1)
+ assert not second.documents and len(second.artifact_keys) == 1
+ await gc.acknowledge(db, second, now=now, error="StorageUnavailable")
+ async with database() as db:
+ retry = await gc.claim(db, generation, now=now + timedelta(minutes=6), limit=1)
+ assert retry.artifact_keys == second.artifact_keys
+ for key in retry.artifact_keys:
+ retained = await p.storage.read_file(key)
+ await p.storage.delete_file(key)
+ await gc.acknowledge(db, retry, now=now + timedelta(minutes=6))
+ # A timed-out writer can finish after deletion; recurring cleanup retains exact keys.
+ await p.storage.write_file(retry.artifact_keys[0], retained)
+ async with database() as db:
+ repeated = await gc.claim(db, generation, now=now + timedelta(hours=1), limit=100)
+ assert repeated.artifact_keys == retry.artifact_keys
+ for key in repeated.artifact_keys:
+ await p.storage.delete_file(key)
+ await gc.acknowledge(db, repeated, now=now + timedelta(hours=1))
+ assert await p.storage.read_file(p.blob.key) == p.raw
diff --git a/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py b/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py
new file mode 100644
index 0000000000..054b874129
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py
@@ -0,0 +1,266 @@
+"""Real SQL: incomplete discovery cannot acquire exhaustive deletion authority."""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+
+from airweave.domains.entities.canonical.coverage import capture_coverage
+from airweave.domains.entities.canonical.cycle_models import (
+ BeginCycle,
+ CompleteCycle,
+ CycleConfiguration,
+)
+from airweave.domains.entities.canonical.cycle_store import CycleConflict
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ CommitOmission,
+ CommitScanPage,
+ ReconcileScan,
+ ScanContinuation,
+)
+from airweave.domains.entities.canonical.scan_store import ScanConflict
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+
+
+async def start(database, service, fence, policy):
+ config = CycleConfiguration(
+ fingerprint="c" * 64, parents={"event": (None,)}, completion_policies={"event": policy}
+ )
+ async with database() as db:
+ cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config))
+ async with database() as db:
+ state = await service.begin_scan(
+ db,
+ BeginScan(
+ fence=fence,
+ scope=CompletedScope(record_type="event"),
+ cycle_id=cycle.version.cycle_id,
+ fingerprint=config.fingerprint,
+ ),
+ )
+ async with database() as db:
+ result = await service.commit_scan_page(
+ db,
+ CommitScanPage(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ records=(),
+ continuation=ScanContinuation(),
+ final=True,
+ ),
+ )
+ return result.state, config
+
+
+async def finish(database, service, fence, state):
+ async with database() as db:
+ return await service.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ observed_at=datetime.now(timezone.utc),
+ ),
+ )
+
+
+@pytest.mark.parametrize("policy", ["discovery_only", "discovery_with_validation"])
+async def test_incomplete_discovery_preserves_known_record_and_reports_policy(
+ database, source, policy
+):
+ service, fence = source
+ original = (await capture(database, service, fence, observation())).changes[0].record
+ state, config = await start(database, service, fence, policy)
+ assert state.completion_policy == policy
+ if policy == "discovery_with_validation":
+ with pytest.raises(ScanConflict, match="still require"):
+ await finish(database, service, fence, state)
+ async with database() as db:
+ request = CommitOmission(
+ fence=fence,
+ state=state,
+ expected_record=original,
+ observation=observation(payload={"summary": "Fresh"}),
+ )
+ result = await service.commit_omission(db, request)
+ assert result.capture.changes[0].record.payload == {"summary": "Fresh"}
+ with pytest.raises(ScanConflict):
+ async with database() as db:
+ await service.commit_omission(db, request) # Lost acknowledgement.
+ state = result.state
+ complete = await finish(database, service, fence, state)
+ assert complete.state.phase == "complete" and complete.capture.changes == ()
+ async with database() as db:
+ coverage = await capture_coverage(db, fence.organization_id, (fence.sync_id,))
+ assert coverage[fence.sync_id].discovery == "incomplete"
+ assert coverage[fence.sync_id].policies == {"event": policy}
+ assert await capture_coverage(db, uuid4(), (fence.sync_id,)) == {}
+ with pytest.raises(CycleConflict):
+ async with database() as db:
+ await service.begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence,
+ configuration=config.model_copy(
+ update={"completion_policies": {"event": "exhaustive"}}
+ ),
+ ),
+ )
+
+ async with database() as db:
+ current = await service.read_cycle(db, fence)
+ async with database() as db:
+ await service.complete_cycle(db, CompleteCycle(fence=fence, expected=current.version))
+ async with database() as db:
+ completed_coverage = await capture_coverage(db, fence.organization_id, (fence.sync_id,))
+ assert completed_coverage[fence.sync_id].phase == "complete"
+ assert completed_coverage[fence.sync_id].discovery == "incomplete"
+
+
+async def test_exact_validation_restores_tombstone_but_rejects_stale_or_wrong_identity(
+ database, source
+):
+ service, fence = source
+ deleted = observation(kind="delete", removal_reason="scope_removed")
+ original = (await capture(database, service, fence, deleted)).changes[0].record
+ state, _ = await start(database, service, fence, "discovery_with_validation")
+ async with database() as db:
+ missing = await service.scan_missing(db, fence, state)
+ assert missing[0].id == original.id
+ for change in [
+ observation("other"),
+ observation(container_id="other"),
+ observation(kind="delete", removal_reason="absent"),
+ ]:
+ with pytest.raises(ScanConflict):
+ async with database() as db:
+ await service.commit_omission(
+ db,
+ CommitOmission(
+ fence=fence, state=state, expected_record=original, observation=change
+ ),
+ )
+ newer = (
+ (await capture(database, service, fence, observation(payload={"v": 2}))).changes[0].record
+ )
+ with pytest.raises(ScanConflict):
+ async with database() as db:
+ await service.commit_omission(
+ db,
+ CommitOmission(
+ fence=fence, state=state, expected_record=original, observation=observation()
+ ),
+ )
+ async with database() as db:
+ result = await service.commit_omission(
+ db,
+ CommitOmission(
+ fence=fence, state=state, expected_record=newer, observation=observation()
+ ),
+ )
+ assert result.capture.changes[0].record.deleted_at is None
+ await finish(database, service, fence, result.state)
+
+
+async def test_restore_by_exact_validation_and_resume_after_lost_ack(database, source):
+ service, fence = source
+ original = (
+ (
+ await capture(
+ database, service, fence, observation(kind="delete", removal_reason="scope_removed")
+ )
+ )
+ .changes[0]
+ .record
+ )
+ state, _ = await start(database, service, fence, "discovery_with_validation")
+ async with database() as db:
+ restored = await service.commit_omission(
+ db,
+ CommitOmission(
+ fence=fence, state=state, expected_record=original, observation=observation()
+ ),
+ )
+ assert restored.capture.changes[0].record.deleted_at is None
+ assert restored.capture.changes[0].record.content_access == "available"
+ async with database() as db:
+ fresh_fence = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ async with database() as db:
+ resumed = await service.read_scan(db, fresh_fence, state.scope)
+ assert resumed.version == restored.state.version
+ complete = await finish(database, service, fresh_fence, resumed)
+ assert complete.state.phase == "complete"
+
+
+async def test_exhaustive_default_cannot_use_omission_mutation(database, source):
+ service, fence = source
+ original = (await capture(database, service, fence, observation())).changes[0].record
+ state, _ = await start(database, service, fence, "exhaustive")
+ with pytest.raises(ScanConflict, match="does not accept"):
+ async with database() as db:
+ await service.commit_omission(
+ db,
+ CommitOmission(
+ fence=fence, state=state, expected_record=original, observation=observation()
+ ),
+ )
+ finished = await finish(database, service, fence, state)
+ assert finished.capture.changes[0].record.deleted_at is not None
+
+
+async def test_driver_retries_failed_exact_read_without_absence_removal(database, source):
+ from unittest.mock import AsyncMock, Mock
+
+ from airweave.domains.entities.canonical.page_source import CapturePage
+ from airweave.domains.storage.file_service import FileService
+ from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver
+
+ service, fence = source
+ original = (await capture(database, service, fence, observation())).changes[0].record
+
+ class Discovery:
+ capture_cycle_configuration = CycleConfiguration(
+ fingerprint="a" * 64,
+ parents={"event": (None,)},
+ completion_policies={"event": "discovery_with_validation"},
+ )
+ failed = False
+
+ async def capture_page(self, scope, continuation, *, files, parent=None):
+ return CapturePage(records=(), continuation=ScanContinuation(), final=True)
+
+ async def refresh_known(self, record, *, files):
+ if not self.failed:
+ self.failed = True
+ raise ConnectionError("Synthetic provider unavailable")
+ assert record.id == original.id
+ return observation(payload={"fresh": True})
+
+ provider = Discovery()
+ driver = CanonicalScanDriver(
+ service, database, fence, provider, AsyncMock(), AsyncMock(), Mock(spec=FileService)
+ )
+ with pytest.raises(ConnectionError):
+ await driver.run()
+ async with database() as db:
+ pending = await service.read_scan(db, fence, CompletedScope(record_type="event"))
+ assert pending.phase == "reconciling"
+ remaining = await service.scan_missing(db, fence, pending)
+ assert remaining[0].deleted_at is None
+ await driver.run()
+ async with database() as db:
+ completed = await service.read_scan(db, fence, pending.scope)
+ assert completed.phase == "complete"
diff --git a/backend/airweave/domains/entities/canonical/tests/test_scans.py b/backend/airweave/domains/entities/canonical/tests/test_scans.py
new file mode 100644
index 0000000000..4ed96de94c
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_scans.py
@@ -0,0 +1,275 @@
+"""Actual PostgreSQL recovery and transaction tests, without provider traffic."""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+from pydantic import ValidationError
+from sqlalchemy import func, select
+
+from airweave.domains.entities.canonical.cycle_models import (
+ BeginCycle,
+ CompleteCycle,
+ CycleConfiguration,
+)
+from airweave.domains.entities.canonical.cycle_store import CycleConflict
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ CommitScanPage,
+ ReconcileScan,
+ ScanContinuation,
+)
+from airweave.domains.entities.canonical.scan_store import ScanConflict
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, StaleWriter
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+from airweave.models.entity_change import EntityChange
+
+SCOPE = CompletedScope(record_type="event")
+FINGERPRINT = "a" * 64
+
+
+async def begin(database, service, fence, *, cycle=None, **kwargs):
+ if cycle is None:
+ async with database() as db:
+ active = await service.begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence,
+ configuration=CycleConfiguration(
+ fingerprint=FINGERPRINT, root_record_type="event"
+ ),
+ ),
+ )
+ cycle = active.version.cycle_id
+ async with database() as db:
+ return await service.begin_scan(
+ db,
+ BeginScan(
+ fence=fence,
+ scope=SCOPE,
+ cycle_id=cycle or uuid4(),
+ fingerprint=FINGERPRINT,
+ **kwargs,
+ ),
+ )
+
+
+async def page(database, service, fence, state, *records, final=False):
+ async with database() as db:
+ return await service.commit_scan_page(
+ db,
+ CommitScanPage(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ records=records,
+ continuation=ScanContinuation(value={"next": "page-2"}),
+ final=final,
+ ),
+ )
+
+
+async def reconcile(database, service, fence, state, limit=250):
+ async with database() as db:
+ return await service.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ observed_at=datetime.now(timezone.utc),
+ limit=limit,
+ ),
+ )
+
+
+async def test_resume_new_writer_preserves_prior_page_sightings(database, source):
+ service, fence = source
+ await capture(database, service, fence, observation("removed"))
+ state = await begin(database, service, fence)
+ first = await page(database, service, fence, state, observation("one"))
+ async with database() as db:
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ resumed = await begin(database, service, newer, cycle=state.cycle_id)
+ assert resumed == first.state
+ with pytest.raises(StaleWriter):
+ await page(database, service, fence, first.state, observation("stale"))
+ final = await page(database, service, newer, resumed, observation("two"), final=True)
+ complete = await reconcile(database, service, newer, final.state)
+ assert complete.state.phase == "complete"
+ assert [change.record.identity.native_id for change in complete.capture.changes] == ["removed"]
+ async with database() as db:
+ active = (
+ await db.scalars(select(Entity.native_id).where(Entity.deleted_at.is_(None)))
+ ).all()
+ assert set(active) == {"one", "two"}
+
+
+async def test_page_commit_without_ack_reload_and_stale_cas(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+ committed = await page(database, service, fence, state, observation())
+ with pytest.raises(ScanConflict):
+ await page(database, service, fence, state, observation("must-not-appear"))
+ async with database() as db:
+ loaded = await service.read_scan(db, fence, SCOPE)
+ assert loaded == committed.state
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 1
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 1
+
+
+async def test_page_failure_rolls_back_records_journal_and_progress(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+
+ class FailingStore(CanonicalRecordStore):
+ async def _capture_locked(self, *args, **kwargs):
+ await super()._capture_locked(*args, **kwargs)
+ raise RuntimeError("synthetic post-write failure")
+
+ failing = CanonicalCaptureService(FailingStore())
+ with pytest.raises(RuntimeError, match="synthetic"):
+ await page(database, failing, fence, state, observation(), final=True)
+ async with database() as db:
+ loaded = await service.read_scan(db, fence, SCOPE)
+ assert loaded == state
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 0
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 0
+
+
+async def test_partial_cannot_reconcile_and_expired_cursor_restarts_sweep(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+ partial = await page(database, service, fence, state, observation("prior-page"))
+ with pytest.raises(ScanConflict, match="fully collected"):
+ await reconcile(database, service, fence, partial.state)
+ restarted = await begin(
+ database, service, fence, cycle=state.cycle_id, expected=partial.state.version, restart=True
+ )
+ assert restarted.version.sweep_id != state.version.sweep_id
+ assert restarted.continuation.value == {}
+ with pytest.raises(ScanConflict):
+ await page(database, service, fence, partial.state, observation("old-reply"))
+ final = await page(database, service, fence, restarted, observation("now-present"), final=True)
+ complete = await reconcile(database, service, fence, final.state)
+ assert [change.record.identity.native_id for change in complete.capture.changes] == [
+ "prior-page"
+ ]
+
+
+async def test_completed_scope_reused_in_cycle_but_next_cycle_refreshes(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+ with pytest.raises(CycleConflict):
+ await begin(database, service, fence, cycle=uuid4(), expected=state.version)
+ final = await page(database, service, fence, state, final=True)
+ complete = (await reconcile(database, service, fence, final.state)).state
+ assert await begin(database, service, fence, cycle=state.cycle_id) == complete
+ async with database() as db:
+ active = await service.read_cycle(db, fence)
+ async with database() as db:
+ finished = await service.complete_cycle(
+ db, CompleteCycle(fence=fence, expected=active.version)
+ )
+ async with database() as db:
+ following = await service.begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence, configuration=finished.configuration, expected=finished.version
+ ),
+ )
+ next_cycle = await begin(
+ database, service, fence, cycle=following.version.cycle_id, expected=complete.version
+ )
+ assert next_cycle.cycle_id != complete.cycle_id
+ assert next_cycle.version.sweep_id != complete.version.sweep_id
+ assert next_cycle.phase == "collecting"
+ with pytest.raises(CycleConflict):
+ await begin(
+ database,
+ service,
+ fence,
+ cycle=complete.cycle_id,
+ expected=next_cycle.version,
+ restart=True,
+ )
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(CaptureScan)) == 1
+
+
+async def test_reconcile_is_bounded_and_resumes_after_worker_loss(database, source):
+ service, fence = source
+ await capture(database, service, fence, *(observation(str(n)) for n in range(3)))
+ state = await begin(database, service, fence)
+ final = await page(database, service, fence, state, final=True)
+ first = await reconcile(database, service, fence, final.state, limit=1)
+ assert first.state.phase == "reconciling" and len(first.capture.changes) == 1
+ resumed = await begin(database, service, fence, cycle=state.cycle_id)
+ assert resumed == first.state
+ last = await reconcile(database, service, fence, resumed)
+ assert last.state.phase == "complete" and len(last.capture.changes) == 2
+ with pytest.raises(ScanConflict):
+ await page(database, service, fence, last.state, observation())
+
+
+async def test_scope_identity_null_and_empty_are_distinct_and_pages_cannot_cross(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+ with pytest.raises(ScanConflict, match="outside"):
+ await page(database, service, fence, state, observation(container_id=""))
+ from airweave.domains.entities.canonical.scan_store import scope_key
+
+ assert scope_key(SCOPE) != scope_key(CompletedScope(record_type="event", container_id=""))
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(CaptureScan)) == 1
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 0
+
+
+def test_progress_is_bounded_and_unknown_fields_rejected():
+ with pytest.raises(ValidationError, match="64 KiB"):
+ ScanContinuation(value={"cursor": "x" * 65536})
+ with pytest.raises(ValidationError):
+ ScanContinuation(value={}, range_start="unsupported")
+
+
+async def test_concurrent_same_page_has_one_commit_and_one_cas_conflict(database, source):
+ import asyncio
+
+ service, fence = source
+ state = await begin(database, service, fence)
+ outcomes = await asyncio.gather(
+ page(database, service, fence, state, observation("one")),
+ page(database, service, fence, state, observation("two")),
+ return_exceptions=True,
+ )
+ assert sum(isinstance(result, ScanConflict) for result in outcomes) == 1
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Entity)) == 1
+ assert await db.scalar(select(func.count()).select_from(EntityChange)) == 1
+ assert (await db.scalar(select(CaptureScan))).revision == 2
+
+
+@pytest.mark.parametrize("database", ["legacy"], indirect=True)
+async def test_scan_migration_preserves_existing_legacy_records(database, source):
+ service, fence = source
+ state = await begin(database, service, fence)
+ final = await page(database, service, fence, state, observation(), final=True)
+ assert (await reconcile(database, service, fence, final.state)).state.phase == "complete"
+ async with database() as db:
+ legacy = await db.scalar(select(Entity).where(Entity.record_revision == 0))
+ assert legacy is not None and legacy.entity_id == "legacy-native-key"
diff --git a/backend/airweave/domains/entities/canonical/tests/test_search_metadata.py b/backend/airweave/domains/entities/canonical/tests/test_search_metadata.py
new file mode 100644
index 0000000000..e2848a4b93
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_search_metadata.py
@@ -0,0 +1,44 @@
+"""Canonical dates keep source ownership and microsecond precision through feed."""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata
+from airweave.domains.entities.canonical.tests.helpers import observation
+from airweave.platform.destinations.vespa.transformer import EntityTransformer
+from airweave.platform.entities._base import AirweaveSystemMetadata
+from airweave.platform.entities.slack import SlackChannelEntity
+
+
+def test_metadata_never_uses_mapper_dates():
+ native = datetime(2026, 1, 1, microsecond=1, tzinfo=timezone.utc)
+ record = SourceRecord(
+ id=uuid4(),
+ sync_id=uuid4(),
+ identity=observation().identity,
+ revision=1,
+ payload={},
+ payload_schema_version=1,
+ capture_hash="hash",
+ content_hash=None,
+ completeness="complete",
+ observed_at=native,
+ source_created_at=native,
+ source_updated_at=None,
+ deleted_at=None,
+ removal_reason=None,
+ blobs=(),
+ indexed_revision=None,
+ indexed_pipeline_version=None,
+ )
+ entity = SlackChannelEntity(channel_id="one", title="One", purpose="", topic="", breadcrumbs=[])
+ entity.created_at = entity.updated_at = datetime(2030, 1, 1, tzinfo=timezone.utc)
+ entity.airweave_system_metadata = AirweaveSystemMetadata()
+ stamp_search_metadata(entity.airweave_system_metadata, record)
+ fields = EntityTransformer()._build_system_metadata(entity)
+ assert fields["source_created_us"] == 1767225600000001
+ assert fields["source_created_known"] == 1
+ assert fields["source_updated_us"] is None
+ assert fields["source_updated_known"] == 0
+ assert fields["canonical_record_type"] == "event"
diff --git a/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py b/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py
new file mode 100644
index 0000000000..c322b93416
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py
@@ -0,0 +1,277 @@
+"""Real PostgreSQL: search owns short reads, even when remote work stops or fails."""
+
+# Imported pytest fixtures are requested by parameter name.
+# ruff: noqa: F811
+
+import asyncio
+from datetime import datetime, timedelta, timezone
+from unittest.mock import Mock
+from uuid import uuid4
+
+import pytest
+from cryptography.fernet import Fernet
+from fastapi import Request
+from fastapi_auth0 import Auth0User
+from sqlalchemy import select, text, update
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+from airweave.adapters.cache.fake import FakeContextCache
+from airweave.adapters.rate_limiter.fake import FakeRateLimiter
+from airweave.api import deps
+from airweave.core import credentials
+from airweave.core.config import AuthMode, settings
+from airweave.core.shared_models import AuthMethod
+from airweave.db import session as db_session
+from airweave.domains.entities.canonical.tests.test_owned_search import ( # noqa: F401
+ http_search,
+ indexed,
+)
+from airweave.domains.entities.canonical.tests.test_search_visibility import hit
+from airweave.domains.search.types import SearchResults
+from airweave.models.api_key import APIKey
+from airweave.models.collection import Collection
+from airweave.models.entity import Entity
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+from airweave.models.user import User
+from airweave.models.user_organization import UserOrganization
+
+
+@pytest.fixture
+async def single_connection(database):
+ async with database() as db:
+ schema = await db.scalar(text("SELECT current_schema()"))
+ url = db.bind.url
+ engine = create_async_engine(
+ url,
+ pool_size=1,
+ max_overflow=0,
+ pool_timeout=0.25,
+ connect_args={"server_settings": {"search_path": schema}},
+ )
+ try:
+ yield async_sessionmaker(engine, expire_on_commit=False)
+ finally:
+ await engine.dispose()
+
+
+async def connection_is_free(sessions):
+ # A second checkout of this one-connection pool would fail if search held it.
+ async with sessions() as probe:
+ assert await probe.scalar(text("SELECT 1")) == 1
+
+
+@pytest.mark.parametrize("stage", ["embedding", "vespa"])
+async def test_waiting_search_releases_connection_and_cancels_cleanly(
+ indexed, http_search, single_connection, stage, monkeypatch
+):
+ fence, _, _ = indexed
+ client, vector, _, executor, _ = http_search
+ client._transport.app.dependency_overrides[deps.get_search_session_factory] = (
+ lambda: single_connection
+ )
+ entered, release = asyncio.Event(), asyncio.Event()
+ target, method = (
+ (executor._sparse_embedder, "embed") if stage == "embedding" else (vector, "execute_query")
+ )
+ original = getattr(target, method)
+
+ async def paused(*args, **kwargs):
+ entered.set()
+ await release.wait()
+ return await original(*args, **kwargs)
+
+ monkeypatch.setattr(target, method, paused)
+ task = asyncio.create_task(
+ client.post(
+ "/sync/search",
+ json={
+ "query": "budget",
+ "sync_ids": [str(fence.sync_id)],
+ "mode": "keyword",
+ },
+ )
+ )
+ try:
+ await asyncio.wait_for(entered.wait(), timeout=5)
+ await connection_is_free(single_connection)
+ finally:
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ release.set()
+ await connection_is_free(single_connection)
+
+
+async def test_revoked_source_during_vespa_cannot_return_cached_hit(
+ indexed, http_search, single_connection, monkeypatch
+):
+ fence, locator, _ = indexed
+ client, vector, _, _, _ = http_search
+ client._transport.app.dependency_overrides[deps.get_search_session_factory] = (
+ lambda: single_connection
+ )
+ vector.seed_results(SearchResults(results=[hit(fence, locator.encode())]))
+ original = vector.execute_query
+
+ async def revoke_then_return(query):
+ # This mutation also proves the network phase has released the only connection.
+ async with single_connection() as db:
+ await db.execute(update(SourceConnection).values(is_authenticated=False))
+ await db.commit()
+ return await original(query)
+
+ monkeypatch.setattr(vector, "execute_query", revoke_then_return)
+ response = await client.post(
+ "/sync/search",
+ json={
+ "query": "budget",
+ "sync_ids": [str(fence.sync_id)],
+ "mode": "keyword",
+ },
+ )
+ assert response.status_code == 404
+ assert "Private original text" not in response.text
+ await connection_is_free(single_connection)
+
+
+@pytest.mark.parametrize("outcome", ["failure", "cancellation"])
+async def test_read_phase_releases_acquired_connection(
+ indexed, http_search, single_connection, outcome
+):
+ fence, locator, _ = indexed
+ client, vector, _, _, _ = http_search
+ app = client._transport.app
+ app.dependency_overrides[deps.get_search_session_factory] = lambda: single_connection
+ service = app.dependency_overrides[deps.get_container]().owned_search
+ vector.seed_results(SearchResults(results=[hit(fence, locator.encode())]))
+ entered = asyncio.Event()
+
+ async def interrupt_after_sql(db, *args):
+ await db.scalar(select(Entity.id))
+ entered.set()
+ if outcome == "failure":
+ raise RuntimeError("synthetic enrichment failure")
+ await asyncio.Event().wait()
+
+ service._enrich = interrupt_after_sql
+ task = asyncio.create_task(
+ client.post(
+ "/sync/search",
+ json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"},
+ )
+ )
+ waiting = asyncio.create_task(entered.wait())
+ try:
+ await asyncio.wait({task, waiting}, timeout=5, return_when=asyncio.FIRST_COMPLETED)
+ if not entered.is_set():
+ if task.done():
+ response = await task # Surface an early exception instead of hiding it as timeout.
+ pytest.fail(f"Search returned HTTP {response.status_code} before enrichment")
+ pytest.fail("Search did not reach enrichment within five seconds")
+ if outcome == "cancellation":
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ else:
+ with pytest.raises(RuntimeError, match="synthetic enrichment failure"):
+ await task
+ finally:
+ # Failed setup must not leave a request borrowing from the fixture's pool.
+ waiting.cancel()
+ task.cancel()
+ await asyncio.gather(waiting, task, return_exceptions=True)
+ await connection_is_free(single_connection)
+
+
+@pytest.mark.parametrize("mode", [AuthMode.API_KEY, AuthMode.AUTH0])
+async def test_search_auth_closes_own_transaction_without_committing_auth0_activity(
+ source, single_connection, mode, monkeypatch
+):
+ _, fence = source
+ monkeypatch.setattr(settings, "AUTH_MODE", mode)
+ monkeypatch.setattr(settings, "ENCRYPTION_KEY", Fernet.generate_key().decode())
+ # No external analytics in this real SQL authentication proof.
+ monkeypatch.setattr("airweave.api.context_resolver.analytics.identify_user", Mock())
+ old_active = datetime(2020, 1, 1)
+ async with single_connection() as db:
+ user = User(email="search@example.com", auth0_id="auth0|search", last_active_at=old_active)
+ db.add(user)
+ await db.flush()
+ user_id = user.id
+ db.add(
+ UserOrganization(
+ user_id=user.id,
+ organization_id=fence.organization_id,
+ role="owner",
+ is_primary=True,
+ )
+ )
+ db.add(
+ APIKey(
+ organization_id=fence.organization_id,
+ encrypted_key=credentials.encrypt({"key": "synthetic-owned-search-key"}),
+ expiration_date=datetime.now(timezone.utc).replace(tzinfo=None)
+ + timedelta(hours=1),
+ )
+ )
+ await db.commit()
+ monkeypatch.setattr(db_session, "AsyncSessionLocal", single_connection)
+ assert deps.get_search_session_factory() is single_connection
+ auth0_user = Auth0User.model_construct(id="auth0|search", email="search@example.com")
+ ctx = await deps.get_owned_search_context(
+ request=Request(
+ {
+ "type": "http",
+ "method": "POST",
+ "path": "/sync/search",
+ "headers": [],
+ "client": ("127.0.0.1", 1),
+ "scheme": "http",
+ }
+ ),
+ sessions=deps.get_search_session_factory(),
+ x_api_key="synthetic-owned-search-key" if mode == AuthMode.API_KEY else None,
+ x_organization_id=str(fence.organization_id),
+ auth0_user=auth0_user if mode == AuthMode.AUTH0 else None,
+ cache=FakeContextCache(),
+ rate_limiter=FakeRateLimiter(),
+ )
+ assert ctx.organization.id == fence.organization_id
+ assert ctx.auth_method == (AuthMethod.API_KEY if mode == AuthMode.API_KEY else AuthMethod.AUTH0)
+ async with single_connection() as db:
+ assert await db.scalar(select(User.last_active_at).where(User.id == user_id)) == old_active
+ if mode == AuthMode.AUTH0:
+ assert ctx.user.last_active_at > old_active
+ await connection_is_free(single_connection)
+
+
+@pytest.mark.parametrize("change", ["collection", "pipeline"])
+async def test_scope_identity_change_during_retrieval_rejects_response(
+ indexed, http_search, single_connection, monkeypatch, change
+):
+ fence, locator, _ = indexed
+ client, vector, _, _, _ = http_search
+ client._transport.app.dependency_overrides[deps.get_search_session_factory] = (
+ lambda: single_connection
+ )
+ vector.seed_results(SearchResults(results=[hit(fence, locator.encode())]))
+ original = vector.execute_query
+
+ async def change_then_return(query):
+ async with single_connection() as db:
+ if change == "collection":
+ await db.execute(update(Collection).values(id=uuid4()))
+ else:
+ await db.execute(update(Sync).values(index_pipeline_version=1))
+ await db.commit()
+ return await original(query)
+
+ monkeypatch.setattr(vector, "execute_query", change_then_return)
+ response = await client.post(
+ "/sync/search",
+ json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"},
+ )
+ assert response.status_code == 404
+ assert "Private original text" not in response.text
+ await connection_is_free(single_connection)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py b/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py
new file mode 100644
index 0000000000..a48cdb9fc7
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py
@@ -0,0 +1,109 @@
+"""Stale index content cannot survive authoritative changes or cross-account scope."""
+
+from types import SimpleNamespace
+from uuid import uuid4
+
+from sqlalchemy import update
+
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared
+from airweave.domains.search.canonical_visibility import visible_results
+from airweave.domains.search.types.results import (
+ SearchAccessControl,
+ SearchResult,
+ SearchSystemMetadata,
+)
+from airweave.models.collection import Collection
+from airweave.models.source_connection import SourceConnection
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+def hit(fence, locator):
+ return SearchResult(
+ entity_id="index-chunk",
+ name="Synthetic",
+ relevance_score=1,
+ breadcrumbs=[],
+ textual_representation="Private original text",
+ airweave_system_metadata=SearchSystemMetadata(
+ source_name="gmail",
+ entity_type="message",
+ sync_id=str(fence.sync_id),
+ chunk_index=0,
+ original_entity_id=locator,
+ ),
+ access=SearchAccessControl(),
+ web_url="https://example.test",
+ raw_source_fields={},
+ )
+
+
+async def test_search_rejects_stale_foreign_malformed_and_legacy_publications(database, source):
+ service, fence = source
+ await capture(database, service, fence, observation())
+ async with database() as db:
+ deployment = VectorDbDeploymentMetadata(
+ dense_embedder="test", embedding_dimensions=3, sparse_embedder="test"
+ )
+ db.add(deployment)
+ await db.flush()
+ db.add(
+ Collection(
+ name="Test",
+ readable_id="test",
+ organization_id=fence.organization_id,
+ vector_db_deployment_metadata_id=deployment.id,
+ )
+ )
+ await db.flush()
+ db.add(
+ SourceConnection(
+ name="Gmail",
+ short_name="gmail",
+ organization_id=fence.organization_id,
+ readable_collection_id="test",
+ sync_id=fence.sync_id,
+ is_authenticated=True,
+ )
+ )
+ await db.commit()
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ generation = uuid4()
+ async with database() as db:
+ assert await publish_prepared(store, db, work, generation, 1)
+ locator = ProjectionLocator(
+ record_id=work.record.id,
+ revision=1,
+ pipeline_version=1,
+ generation=generation,
+ part_index=0,
+ )
+ valid = hit(fence, locator.encode())
+ registry = SimpleNamespace(
+ get=lambda _: SimpleNamespace(
+ source_class_ref=SimpleNamespace(canonical_record_types=("message",))
+ )
+ )
+ wrong_revision = locator.model_copy(update={"revision": 2})
+ inputs = [
+ valid,
+ hit(fence, wrong_revision.encode()),
+ hit(fence, "native-id"),
+ hit(fence, "canonical:invalid"),
+ ]
+ async with database() as db:
+ assert await visible_results(db, fence.organization_id, "test", inputs, registry) == [valid]
+ assert await visible_results(db, uuid4(), "test", inputs, registry) == []
+ assert await visible_results(db, fence.organization_id, "other", inputs, registry) == []
+ async with database() as db:
+ await db.execute(update(SourceConnection).values(is_authenticated=False))
+ await db.commit()
+ assert await visible_results(db, fence.organization_id, "test", inputs, registry) == []
+ await db.execute(update(SourceConnection).values(is_authenticated=True))
+ await db.commit()
+ await capture(database, service, fence, observation(payload={"id": "one", "summary": "Edited"}))
+ async with database() as db:
+ assert await visible_results(db, fence.organization_id, "test", inputs, registry) == []
diff --git a/backend/airweave/domains/entities/canonical/tests/test_service_key_http.py b/backend/airweave/domains/entities/canonical/tests/test_service_key_http.py
new file mode 100644
index 0000000000..48361749a7
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_service_key_http.py
@@ -0,0 +1,98 @@
+"""Real persisted-key authentication through the original-record HTTP boundary."""
+
+from datetime import timedelta
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+from uuid import uuid4
+
+from cryptography.fernet import Fernet
+from fastapi import FastAPI
+from httpx import ASGITransport, AsyncClient
+from sqlalchemy import delete, update
+
+from airweave.adapters.rate_limiter.null import NullRateLimiter
+from airweave.api.v1.endpoints.records import record_error_response, router
+from airweave.core import container as container_mod
+from airweave.core import credentials
+from airweave.core.config import AuthMode, settings
+from airweave.core.datetime_utils import utc_now_naive
+from airweave.db.session import get_db
+from airweave.domains.entities.canonical.requests import CaptureBatch
+from airweave.domains.entities.canonical.store import CanonicalStoreError
+from airweave.domains.entities.canonical.tests.helpers import observation
+from airweave.models import Organization
+from airweave.models.api_key import APIKey
+
+
+async def test_persisted_keys_scope_expiry_and_revocation_over_http(database, source, monkeypatch):
+ """Only Redis cache/rate limiting are substituted; auth and SQL stay real."""
+ monkeypatch.setattr(settings, "AUTH_MODE", AuthMode.API_KEY)
+ monkeypatch.setattr(settings, "ENCRYPTION_KEY", Fernet.generate_key().decode())
+ cache = SimpleNamespace(
+ get_organization=AsyncMock(return_value=None), set_organization=AsyncMock()
+ )
+ monkeypatch.setattr(
+ container_mod,
+ "container",
+ SimpleNamespace(context_cache=cache, rate_limiter=NullRateLimiter()),
+ )
+ capture, fence = source
+ other_org = uuid4()
+ owner_key, foreign_key = uuid4().hex, uuid4().hex
+ async with database() as db:
+ db.add(Organization(id=other_org, name="Other synthetic owner"))
+ await db.flush()
+ keys = [
+ APIKey(
+ organization_id=org,
+ encrypted_key=credentials.encrypt({"key": key}),
+ expiration_date=utc_now_naive() + timedelta(days=1),
+ created_by_email="synthetic@example.com",
+ modified_by_email="synthetic@example.com",
+ )
+ for org, key in ((fence.organization_id, owner_key), (other_org, foreign_key))
+ ]
+ db.add_all(keys)
+ await db.commit()
+ result = await capture.capture(db, CaptureBatch(fence=fence, records=(observation("one"),)))
+ app = FastAPI()
+ app.include_router(router, prefix="/sync")
+ app.add_exception_handler(CanonicalStoreError, record_error_response)
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ base = f"/sync/{fence.sync_id}/records"
+ paths = (base, f"{base}/{result.changes[0].record.id}", base + "/changes")
+ async with AsyncClient(transport=ASGITransport(app), base_url="http://test") as client:
+ for path in paths:
+ assert (await client.get(path)).status_code == 401
+ assert (await client.get(path, headers={"X-API-Key": "invalid"})).status_code == 403
+ valid = await client.get(path, headers={"X-API-Key": owner_key})
+ assert valid.status_code == 200, valid.text
+ assert (await client.get(path, headers={"X-API-Key": foreign_key})).status_code == 404
+ assert (
+ await client.get(
+ path,
+ headers={
+ "X-API-Key": foreign_key,
+ "X-Organization-ID": str(fence.organization_id),
+ },
+ )
+ ).status_code == 403
+ async with database() as db:
+ await db.execute(
+ update(APIKey)
+ .where(APIKey.id == keys[0].id)
+ .values(
+ expiration_date=utc_now_naive() - timedelta(seconds=1),
+ )
+ )
+ await db.commit()
+ assert (await client.get(base, headers={"X-API-Key": owner_key})).status_code == 403
+ async with database() as db:
+ await db.execute(delete(APIKey).where(APIKey.id == keys[1].id))
+ await db.commit()
+ assert (await client.get(base, headers={"X-API-Key": foreign_key})).status_code == 403
diff --git a/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py b/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py
new file mode 100644
index 0000000000..59ea3fa124
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py
@@ -0,0 +1,232 @@
+"""Real SQL publication/revision/access gates over native spreadsheet blobs."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.domains.entities.canonical.projection_mappers import _drive
+from airweave.domains.entities.canonical.query import (
+ CanonicalQueryService,
+ RecordNotFound,
+ StaleRecordRevision,
+)
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.helpers import capture
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.sources.records.google_drive import file_record
+from airweave.platform.sources.records.google_drive_content import capture_file_content
+from airweave.platform.sources.records.sheets_manifest import SHEETS_MIME, GridBounds
+
+
+async def test_sheets_exact_revision_projection_and_removal(
+ database, source, tmp_path, monkeypatch
+):
+ service, fence = source
+ monkeypatch.setattr(
+ "airweave.domains.storage.file_service.paths.temp_sync_dir", lambda _: str(tmp_path / "tmp")
+ )
+ files = FileService(fence.job_id, FilesystemBackend(tmp_path / "store"), sync_id=fence.sync_id)
+ grid = {
+ "spreadsheetId": "book",
+ "sheets": [
+ {
+ "properties": {
+ "sheetId": 0,
+ "title": "Sheet",
+ "gridProperties": {"rowCount": 2, "columnCount": 2},
+ }
+ }
+ ],
+ }
+
+ async def latest(*args, **kwargs):
+ return {"version": "7"}
+
+ def response(request):
+ if request.url.path.endswith("/export"):
+ return httpx.Response(200, content=b"exact-export")
+ if "ranges" in request.url.params:
+ return httpx.Response(
+ 200,
+ json={
+ "spreadsheetId": "book",
+ "sheets": [
+ {
+ **grid["sheets"][0],
+ "data": [
+ {
+ "rowData": [
+ {
+ "values": [
+ {
+ "formattedValue": "42",
+ "userEnteredValue": {"formulaValue": "=6*7"},
+ }
+ ]
+ }
+ ]
+ }
+ ],
+ }
+ ],
+ },
+ )
+ return httpx.Response(200, json=grid)
+
+ original = file_record({"id": "book", "version": "7", "mimeType": SHEETS_MIME})
+ async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client:
+ retained = await capture_file_content(
+ original,
+ files=files,
+ get=latest,
+ client=client,
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ capture_native_sheets=True,
+ )
+ record = (await capture(database, service, fence, retained)).changes[0].record
+ query = CanonicalQueryService(
+ CanonicalRecordStore(), CanonicalQueryStore(), "synthetic-signing"
+ )
+ async with database() as db:
+ read = await query.spreadsheet(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ record.id,
+ record.revision,
+ files.storage,
+ GridBounds(sheet_id=0, start_row=0, end_row=2, start_column=0, end_column=2),
+ )
+ assert len(read.cells) == 1
+ assert read.cells[0].value["userEnteredValue"] == {"formulaValue": "=6*7"}
+ assert read.grid_status == "complete" and not read.missing
+ with pytest.raises(StaleRecordRevision):
+ await query.spreadsheet(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ record.id,
+ record.revision + 1,
+ files.storage,
+ )
+ entities = (await _drive(record, files.storage, tmp_path / "projection")).entities
+ assert len(entities) == 1
+ from pathlib import Path
+
+ assert "42" in Path(entities[0].local_path).read_text()
+ assert "=6*7" not in Path(entities[0].local_path).read_text()
+ monkeypatch.setattr(
+ "airweave.platform.sources.records.google_sheets_content.MAX_GRID_REQUESTS", 0
+ )
+ async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client:
+ partial = await capture_file_content(
+ original,
+ files=files,
+ get=latest,
+ client=client,
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ capture_native_sheets=True,
+ )
+ newer = (await capture(database, service, fence, partial)).changes[0].record
+ assert newer.id == record.id and newer.revision == record.revision + 1
+ projection = await _drive(newer, files.storage, tmp_path / "fallback")
+ fallback = projection.entities
+ assert len(projection.parts) == 2
+ assert projection.parts[1].entity is None
+ assert projection.parts[1].part.key == "/native/sheets/0/rows/0:2/columns/0:2"
+ assert newer.completeness == "partial"
+ assert len(fallback) == 1
+ assert Path(fallback[0].local_path).read_bytes() == b"exact-export"
+ assert fallback[0].file_type == "xlsx"
+ async with database() as db:
+ native_partial = await query.spreadsheet(
+ db, fence.organization_id, fence.sync_id, newer.id, newer.revision, files.storage
+ )
+ assert native_partial.grid_status == "partial" and len(native_partial.missing) == 1
+ assert not native_partial.captured and not native_partial.cells
+ await capture(
+ database,
+ service,
+ fence,
+ original.model_copy(update={"kind": "delete", "removal_reason": "scope_removed"}),
+ )
+ async with database() as db:
+ with pytest.raises(RecordNotFound):
+ await query.spreadsheet(
+ db, fence.organization_id, fence.sync_id, record.id, record.revision, files.storage
+ )
+
+
+async def test_failed_native_tile_retries_same_file_without_advancing_checkpoint(
+ database, source, tmp_path, monkeypatch
+):
+ from contextlib import aclosing
+
+ from sqlalchemy import select
+
+ from airweave.domains.entities.canonical.tests import test_drive_recovery as recovery
+ from airweave.domains.entities.canonical.tests.test_workspace_capture import NativeDocs
+ from airweave.models.entity import Entity
+ from airweave.platform.configs.config import GoogleDriveConfig
+
+ class NativeSheets(NativeDocs):
+ async def handle(self, request):
+ if request.url.host == "sheets.googleapis.com":
+ value = {
+ "spreadsheetId": "doc",
+ "sheets": [
+ {
+ "properties": {
+ "sheetId": 0,
+ "title": "Sheet",
+ "gridProperties": {"rowCount": 1, "columnCount": 1},
+ }
+ }
+ ],
+ }
+ if "ranges" in request.url.params and self.fail:
+ raise ConnectionError("interrupted native tile")
+ return httpx.Response(200, json=value)
+ response = await super().handle(request)
+ if request.url.path.endswith("/files/doc"):
+ return httpx.Response(
+ 200, json={"id": "doc", "version": "7", "mimeType": SHEETS_MIME}
+ )
+ return response
+
+ monkeypatch.setattr(
+ recovery,
+ "GoogleDriveConfig",
+ lambda **kwargs: GoogleDriveConfig(capture_native_sheets=True, **kwargs),
+ )
+ monkeypatch.setattr(
+ "airweave.domains.storage.file_service.paths.temp_sync_dir", lambda _: str(tmp_path / "tmp")
+ )
+ service, fence = source
+ files = FileService(fence.job_id, FilesystemBackend(tmp_path / "store"), sync_id=fence.sync_id)
+ native = NativeSheets()
+ pipeline, ctx, runtime, client = await recovery.setup(database, source, native)
+ pipeline.files = files
+ async with aclosing(client):
+ with pytest.raises(ConnectionError, match="interrupted native tile"):
+ await recovery.run(pipeline, ctx, runtime)
+ async with database() as db:
+ assert not (await db.scalars(select(Entity))).all()
+ cycle = await service.read_cycle(db, pipeline._writer())
+ assert cycle.promoted_checkpoint is None
+ native.fail = False
+ pipeline, ctx, runtime, client = await recovery.setup(database, source, native, attempt=2)
+ pipeline.files = files
+ async with aclosing(client):
+ await recovery.run(pipeline, ctx, runtime)
+ async with database() as db:
+ records = (await db.scalars(select(Entity))).all()
+ assert len(records) == 1
+ cycle = await service.read_cycle(db, pipeline._writer())
+ assert cycle.promoted_checkpoint.checkpoint.value == {"page_token": "end"}
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py
new file mode 100644
index 0000000000..7afbc97f97
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py
@@ -0,0 +1,346 @@
+"""Actual child frontier commits/retry with synthetic Slack HTTP and original bytes."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.domains.entities.canonical.page_source import (
+ InvalidScanContinuation,
+ RequiredScopeAccessLost,
+)
+from airweave.domains.entities.canonical.requests import CaptureBatch, RecordIdentity
+from airweave.domains.entities.canonical.tests.test_capture_pipeline import components, orchestrator
+from airweave.domains.entities.canonical.tests.test_slack_recovery import run
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.domains.storage.file_service import FileService
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+from airweave.models.sync_cursor import SyncCursor
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackPrincipal, SlackSource
+
+ROOT = {"channels": [{"id": "C1", "name": "Synthetic"}]}
+MESSAGE = {
+ "ts": "1",
+ "text": "Original message",
+ "files": [
+ {
+ "id": f"F{i}",
+ "size": 3,
+ "mimetype": "text/plain",
+ "url_private": f"https://files.slack.com/F{i}",
+ }
+ for i in range(1, 4)
+ ],
+}
+
+
+async def attempt(
+ database, source, storage, calls, *, number, fail_last=False, responses=None, connector_out=None
+):
+ service, fence = source
+ ctx, _, runtime, bus = components(database, source)
+ connector = SlackSource(
+ auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock()
+ )
+ connector.slack_config = SlackConfig(
+ expected_team_id="T1", expected_user_id="U1", capture_files=True
+ )
+ connector._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1")
+ # A fresh message page attests its complete inventory in the same SQL commit.
+ # A retry with no fresh owner page must still perform exact owner refresh.
+ defaults = [ROOT, {"messages": [MESSAGE]}]
+ connector._get = AsyncMock(side_effect=responses if responses is not None else defaults)
+ if connector_out is not None:
+ connector_out.append(connector)
+ files = FileService(uuid4(), storage, sync_id=fence.sync_id)
+
+ async def download(url, *args, **kwargs):
+ identity = url.rsplit("/", 1)[-1]
+ calls.append(identity)
+ if fail_last and identity == "F3":
+ raise ConnectionError("Synthetic last-file failure")
+ return await files.store_canonical_blob(b"abc", media_type="text/plain")
+
+ files.capture_canonical_url = download
+ pipeline = CanonicalCapturePipeline(
+ service,
+ database,
+ bus,
+ connector.canonical_record_types,
+ CaptureAttempt(id=fence.attempt_id if number == 1 else uuid4(), number=number),
+ connector.canonical_container_parents,
+ page_source=connector,
+ files=files,
+ )
+ runtime.source = connector
+ runtime.canonical_capture = pipeline
+ instance = orchestrator(ctx, pipeline, runtime, None, bus)
+ instance.stream = None
+ try:
+ await run(instance)
+ finally:
+ await files.cleanup_sync_directory(MagicMock())
+ return pipeline, connector
+
+
+async def test_last_file_retry_preserves_committed_children_and_inventory_hides_old_bytes(
+ database, source, tmp_path
+):
+ service, fence = source
+ storage = FilesystemBackend(tmp_path)
+ calls = []
+ with pytest.raises(ConnectionError, match="last-file"):
+ await attempt(database, source, storage, calls, number=1, fail_last=True)
+ assert calls == ["F1", "F2", "F3"]
+ async with database() as db:
+ rows = (await db.scalars(select(Entity).where(Entity.sync_id == fence.sync_id))).all()
+ assert {row.native_id for row in rows} == {"C1", "1", "F1", "F2"}
+ message = next(row for row in rows if row.native_id == "1")
+ assert message.source_payload == MESSAGE and message.payload_schema_version == 2
+ pipeline, connector = await attempt(database, source, storage, calls, number=2)
+ assert calls == ["F1", "F2", "F3", "F3"]
+ assert connector._get.await_count == 2
+ async with database() as db:
+ rows = (await db.scalars(select(Entity).where(Entity.sync_id == fence.sync_id))).all()
+ children = [row for row in rows if row.entity_definition_short_name == "file"]
+ assert len(children) == 3 and all(row.record_revision == 1 for row in children)
+ current = connector._capture_message(
+ {**MESSAGE, "files": [MESSAGE["files"][2]]}, "C1"
+ ).model_copy(
+ update={
+ "payload_schema_version": 2,
+ "descendant_visibility_fields": ("files",),
+ "parent": RecordIdentity(record_type="channel", native_id="C1"),
+ "completeness": "complete",
+ }
+ )
+ await service.capture(db, CaptureBatch(fence=pipeline._writer(), records=(current,)))
+ async with database() as db:
+ for row in children:
+ stored = await service.store.read(db, fence.organization_id, fence.sync_id, row.id)
+ assert stored.content_access == "unavailable"
+
+
+async def test_omitted_accessible_message_does_not_remove_retained_file_children(
+ database, source, tmp_path
+):
+ from airweave.models.sync_job import SyncJob
+
+ service, fence = source
+ storage = FilesystemBackend(tmp_path)
+ await attempt(database, source, storage, [], number=1)
+ new_job = uuid4()
+ async with database() as db:
+ (await db.get(SyncJob, fence.job_id)).status = "completed"
+ db.add(
+ SyncJob(
+ id=new_job,
+ sync_id=fence.sync_id,
+ organization_id=fence.organization_id,
+ status="running",
+ )
+ )
+ await db.commit()
+ following = (service, fence.model_copy(update={"job_id": new_job, "attempt_id": uuid4()}))
+ connectors = []
+ with pytest.raises(ValueError, match="omitted an accessible prior message"):
+ await attempt(
+ database,
+ following,
+ storage,
+ [],
+ number=1,
+ responses=[ROOT, {"messages": []}, {"messages": [MESSAGE]}],
+ connector_out=connectors,
+ )
+ call = connectors[0]._get.call_args_list[-1]
+ assert call.args[0].endswith("conversations.history")
+ assert call.args[1] == {
+ "channel": "C1",
+ "oldest": "1",
+ "latest": "1",
+ "inclusive": "true",
+ "limit": 1,
+ }
+ async with database() as db:
+ rows = (await db.scalars(select(Entity).where(Entity.sync_id == fence.sync_id))).all()
+ assert len(rows) == 5
+ for row in rows:
+ stored = await service.store.read(db, fence.organization_id, fence.sync_id, row.id)
+ assert stored.deleted_at is None and stored.content_access == "available"
+ if row.entity_definition_short_name == "file":
+ assert stored.blobs
+
+
+@pytest.mark.parametrize("repeat_owner", [False, True])
+async def test_files_progress_before_history_finishes_and_resume_both_cursors(
+ database, source, tmp_path, repeat_owner
+):
+ storage = FilesystemBackend(tmp_path)
+ downloads = []
+ histories = []
+ owner_reads = []
+ interrupt = True
+
+ async def respond(url, params):
+ if url.endswith("conversations.list"):
+ return ROOT
+ if "oldest" in params:
+ owner_reads.append(params["oldest"])
+ return {"messages": [MESSAGE]}
+ cursor = params.get("cursor")
+ histories.append(cursor)
+ if cursor is None:
+ return {"messages": [MESSAGE], "response_metadata": {"next_cursor": "H2"}}
+ if cursor == "H2":
+ if interrupt:
+ raise asyncio.CancelledError()
+ return {
+ "messages": [MESSAGE] if repeat_owner else [],
+ "response_metadata": {"next_cursor": "H3"},
+ }
+ assert cursor == "H3"
+ return {"messages": []}
+
+ with pytest.raises(asyncio.CancelledError):
+ await attempt(database, source, storage, downloads, number=1, responses=respond)
+ assert downloads == ["F1"]
+ assert owner_reads == []
+ async with database() as db:
+ scans = (await db.scalars(select(CaptureScan))).all()
+ history = next(row for row in scans if row.record_type == "message")
+ files = next(row for row in scans if row.record_type == "file")
+ assert history.phase == files.phase == "collecting"
+ assert history.continuation["history_cursor"] == "H2"
+ sweeps = (history.sweep_id, files.sweep_id)
+ assert "canonical_checkpoint" not in (await db.scalar(select(SyncCursor))).cursor_data
+
+ interrupt = False
+ await attempt(database, source, storage, downloads, number=2, responses=respond)
+ assert downloads == ["F1", "F2", "F3"]
+ assert histories == [None, "H2", "H2", "H3"]
+ assert owner_reads == ([] if repeat_owner else ["1"])
+ async with database() as db:
+ scans = (await db.scalars(select(CaptureScan))).all()
+ history = next(row for row in scans if row.record_type == "message")
+ files = next(row for row in scans if row.record_type == "file")
+ assert (history.sweep_id, files.sweep_id) == sweeps
+ assert history.phase == files.phase == "complete"
+ rows = (await db.scalars(select(Entity))).all()
+ assert next(row for row in rows if row.native_id == "1").source_payload == MESSAGE
+ assert all(row.record_revision == 1 for row in rows if row.native_id.startswith("F"))
+ cursor = (await db.scalar(select(SyncCursor))).cursor_data
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+
+
+async def test_child_cursor_restart_bound_survives_interleaved_turns(
+ database, source, tmp_path, monkeypatch
+):
+ original = SlackSource.capture_page
+ file_calls = 0
+
+ async def capture(connector, scope, continuation, **kwargs):
+ nonlocal file_calls
+ if scope.record_type == "file":
+ file_calls += 1
+ if file_calls != 2:
+ raise InvalidScanContinuation("Synthetic child cursor expiry")
+ return await original(connector, scope, continuation, **kwargs)
+
+ monkeypatch.setattr(SlackSource, "capture_page", capture)
+
+ async def respond(url, params):
+ if url.endswith("conversations.list"):
+ return ROOT
+ if "oldest" in params:
+ return {"messages": [MESSAGE]}
+ return {
+ "messages": [MESSAGE] if "cursor" not in params else [],
+ "response_metadata": {"next_cursor": "H3" if "cursor" in params else "H2"},
+ }
+
+ downloads = []
+ with pytest.raises(InvalidScanContinuation, match="child cursor"):
+ await attempt(
+ database, source, FilesystemBackend(tmp_path), downloads, number=1, responses=respond
+ )
+ assert file_calls == 3 and downloads == ["F1"]
+ async with database() as db:
+ cursor = (await db.scalar(select(SyncCursor))).cursor_data
+ assert "canonical_checkpoint" not in cursor
+ assert cursor["canonical_cycle"]["phase"] == "active"
+
+
+async def test_required_child_access_loss_does_not_withdraw_enclosing_channel(
+ database, source, tmp_path, monkeypatch
+):
+ original = SlackSource.capture_page
+
+ async def capture(connector, scope, continuation, **kwargs):
+ if scope.record_type == "file":
+ # Exercise the shared driver contract; Slack does not currently emit this subtype.
+ raise RequiredScopeAccessLost("Synthetic required child access loss")
+ return await original(connector, scope, continuation, **kwargs)
+
+ monkeypatch.setattr(SlackSource, "capture_page", capture)
+ with pytest.raises(RequiredScopeAccessLost, match="required child"):
+ await attempt(database, source, FilesystemBackend(tmp_path), [], number=1)
+ service, fence = source
+ async with database() as db:
+ rows = (await db.scalars(select(Entity))).all()
+ assert {row.native_id for row in rows} == {"C1", "1"}
+ for row in rows:
+ stored = await service.store.read(db, fence.organization_id, fence.sync_id, row.id)
+ expected = "available" if row.native_id == "C1" else "unavailable"
+ assert stored.content_access == expected
+ assert "canonical_checkpoint" not in (await db.scalar(select(SyncCursor))).cursor_data
+
+
+async def test_fresh_empty_inventory_reconciles_old_files_without_exact_owner_call(
+ database, source, tmp_path
+):
+ from airweave.models.sync_job import SyncJob
+
+ service, fence = source
+ storage = FilesystemBackend(tmp_path)
+ downloads = []
+ await attempt(database, source, storage, downloads, number=1)
+ assert downloads == ["F1", "F2", "F3"]
+ job = uuid4()
+ async with database() as db:
+ (await db.get(SyncJob, fence.job_id)).status = "completed"
+ db.add(
+ SyncJob(
+ id=job,
+ sync_id=fence.sync_id,
+ organization_id=fence.organization_id,
+ status="running",
+ )
+ )
+ await db.commit()
+ following = (service, fence.model_copy(update={"job_id": job, "attempt_id": uuid4()}))
+ pipeline, connector = await attempt(
+ database,
+ following,
+ storage,
+ downloads,
+ number=1,
+ responses=[ROOT, {"messages": [{**MESSAGE, "files": []}]}],
+ )
+ assert connector._get.await_count == 2
+ assert downloads == ["F1", "F2", "F3"]
+ async with database() as db:
+ children = (
+ await db.scalars(select(Entity).where(Entity.entity_definition_short_name == "file"))
+ ).all()
+ assert len(children) == 3 and all(child.deleted_at is not None for child in children)
+ scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "file"))
+ assert scan.phase == "complete"
+ assert scan.parent_verified_attempt_id == pipeline._writer().attempt_id
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_identity_writes.py b/backend/airweave/domains/entities/canonical/tests/test_slack_identity_writes.py
new file mode 100644
index 0000000000..390d79a886
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_identity_writes.py
@@ -0,0 +1,37 @@
+"""A different native Slack user must never begin source capture."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.models.entity import Entity
+from airweave.models.sync_cursor import SyncCursor
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackSource
+
+
+async def test_wrong_user_leaves_existing_cursor_and_originals_unchanged(database, source):
+ async with database() as db:
+ before = [(row.id, row.cursor_data) for row in (await db.scalars(select(SyncCursor))).all()]
+ calls = []
+
+ def profile(request):
+ calls.append(request.url.path)
+ return httpx.Response(200, json={"ok": True, "team_id": "T1", "user_id": "U2"})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(profile)) as client:
+ with pytest.raises(ValueError, match="does not match"):
+ await SlackSource.create(
+ auth=StaticTokenProvider("same-broker"),
+ logger=MagicMock(),
+ http_client=client,
+ config=SlackConfig(expected_team_id="T1", expected_user_id="U1"),
+ )
+ assert calls == ["/api/auth.test"]
+ async with database() as db:
+ assert (await db.scalars(select(Entity))).all() == []
+ after = [(row.id, row.cursor_data) for row in (await db.scalars(select(SyncCursor))).all()]
+ assert after == before
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py
new file mode 100644
index 0000000000..b9e4d93826
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py
@@ -0,0 +1,41 @@
+"""Retention-limited history cannot erase prior originals or publish a new baseline."""
+
+import pytest
+
+from airweave.domains.entities.canonical.tests.test_mixed_scopes import next_job
+from airweave.domains.entities.canonical.tests.test_slack_recovery import ROOT, run, runner, saved
+from airweave.domains.sources.exceptions import SourceError
+
+
+async def test_limited_final_history_preserves_prior_originals_and_committed_progress(
+ database, source
+):
+ baseline, _, _ = await runner(
+ database, source, [ROOT, {"messages": [{"ts": "1", "text": "old"}]}]
+ )
+ await run(baseline)
+ _, before, _ = await saved(database)
+ source = await next_job(database, source)
+ refresh, _, _ = await runner(
+ database,
+ source,
+ [
+ ROOT,
+ {
+ "messages": [{"ts": "3", "text": "new"}],
+ "response_metadata": {"next_cursor": "older"},
+ },
+ {"messages": [{"ts": "2", "text": "limited page"}], "is_limited": True},
+ ],
+ )
+ with pytest.raises(SourceError, match="remains incomplete"):
+ await run(refresh)
+ rows, after, scans = await saved(database)
+ assert {row.native_id for row in rows} == {"C1", "1", "3"}
+ assert all(row.deleted_at is None for row in rows)
+ assert after["canonical_checkpoint"] == before["canonical_checkpoint"]
+ assert after["canonical_cycle"]["phase"] != "complete"
+ current = next(scan for scan in scans if scan.record_type == "message")
+ assert current.phase == "collecting"
+ assert current.continuation["history_cursor"] == "older"
+ assert current.continuation["history_done"] is False
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py
new file mode 100644
index 0000000000..5efad75762
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py
@@ -0,0 +1,90 @@
+"""Lost membership hides records captured before the first durable checkpoint."""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.domains.entities.canonical.tests.test_capture_pipeline import components
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackApiError, SlackSource
+
+pytestmark = pytest.mark.integration
+
+
+@pytest.mark.parametrize("listing_fails", [False, True])
+@pytest.mark.parametrize("capture_files", [False, True])
+async def test_first_completed_membership_hides_prior_partial_channel(
+ database, source, listing_fails, capture_files
+):
+ service, fence = source
+ parent = RecordIdentity(record_type="channel", native_id="lost-channel")
+ seeded = await capture(
+ database,
+ service,
+ fence,
+ observation(identity=parent),
+ observation(
+ identity=RecordIdentity(
+ record_type="message", native_id="1", container_id="lost-channel"
+ ),
+ parent=parent,
+ ),
+ )
+ child_id = seeded.changes[1].record.id
+ # The failed prior capture saved records but no channel_ids checkpoint.
+ ctx, _, runtime, bus = components(database, source)
+ connector = SlackSource(
+ auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock()
+ )
+ connector.slack_config = SlackConfig(
+ expected_team_id="T1", expected_user_id="U1", capture_files=capture_files
+ )
+ connector._get = AsyncMock(return_value={"ok": True, "team_id": "T1", "user_id": "U1"})
+ await connector.validate()
+ connector._get = AsyncMock(
+ side_effect=(
+ ConnectionError("membership page failed")
+ if listing_fails
+ else [{"channels": []}, SlackApiError("channel_not_found")]
+ ),
+ )
+ pipeline = CanonicalCapturePipeline(
+ service,
+ database,
+ bus,
+ connector.canonical_record_types,
+ CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number),
+ connector.canonical_container_parents,
+ page_source=connector,
+ files=MagicMock(),
+ )
+ await pipeline.start(ctx)
+
+ async def scan():
+ await pipeline.run_scans(ctx, runtime, AsyncMock())
+
+ if listing_fails:
+ with pytest.raises(ConnectionError, match="membership page failed"):
+ await scan()
+ else:
+ await scan()
+ await pipeline.cleanup_orphaned_entities(ctx, runtime)
+ async with database() as db:
+ record = await service.store.read(db, fence.organization_id, fence.sync_id, child_id)
+ if listing_fails:
+ assert record.content_access == "available" and record.payload
+ return
+ assert record.content_access == "unavailable"
+ assert record.payload == {} and record.blobs == ()
+ parent_record = await service.store.read(
+ db, fence.organization_id, fence.sync_id, seeded.changes[0].record.id
+ )
+ assert parent_record.removal_reason == "scope_removed"
+ assert parent_record.content_access == "unavailable" and parent_record.payload == {}
+ changes = await service.store.changes(db, fence.organization_id, fence.sync_id)
+ assert all(change.record.payload == {} for change in changes.changes)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py
new file mode 100644
index 0000000000..b6eec492eb
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py
@@ -0,0 +1,379 @@
+"""Retained Slack file projection: no provider calls or live authorization claims."""
+
+import hashlib
+from datetime import datetime, timezone
+from pathlib import Path
+from unittest.mock import AsyncMock
+from uuid import uuid4
+
+import pytest
+
+from airweave.domains.entities.canonical.blob_materializer import BlobIntegrityError
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.projection_mappers import map_record
+from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity
+from airweave.platform.entities.slack import SlackAttachmentEntity
+
+
+def message(files):
+ return SourceRecord(
+ id=uuid4(),
+ sync_id=uuid4(),
+ identity=RecordIdentity(record_type="message", native_id="1.000001", container_id="C1"),
+ revision=1,
+ payload={"ts": "1.000001", "text": "Intact message", "files": files},
+ payload_schema_version=1,
+ capture_hash="capture",
+ content_hash=None,
+ completeness="partial",
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=None,
+ source_updated_at=None,
+ deleted_at=None,
+ removal_reason=None,
+ blobs=(),
+ indexed_revision=None,
+ indexed_pipeline_version=None,
+ )
+
+
+def blob(record, data, index):
+ digest = hashlib.sha256(data).hexdigest()
+ return BlobReference(
+ key=f"canonical/{record.sync_id}/blobs/sha256/{digest}",
+ sha256=digest,
+ size_bytes=len(data),
+ source_path=f"/files/{index}",
+ )
+
+
+@pytest.mark.asyncio
+async def test_body_retained_pdf_duplicate_bytes_and_unavailable_file():
+ record = message(
+ [
+ {"id": "F1", "name": "one.pdf", "mimetype": "application/pdf"},
+ {"id": "F2", "name": "two.pdf", "mimetype": "application/pdf"},
+ {"id": "F3", "name": "clip.mp4", "mimetype": "video/mp4"},
+ ]
+ )
+ content = b"%PDF-1.4 retained fixture, converter not exercised"
+ record = record.model_copy(
+ update={"blobs": (blob(record, content, 0), blob(record, content, 1))}
+ )
+ before = record.model_dump()
+ storage = AsyncMock()
+ storage.read_file.return_value = content
+ async with map_record(record, "slack", storage) as mapped:
+ assert [part.part.key for part in mapped.parts] == ["body", "file:F1", "file:F2", "file:F3"]
+ assert mapped.parts[0].entity.text == "Intact message"
+ assert mapped.parts[3].entity is None
+ first, second = mapped.parts[1].entity, mapped.parts[2].entity
+ assert isinstance(first, SlackAttachmentEntity)
+ assert first.entity_id != second.entity_id
+ assert first.local_path == second.local_path
+ assert Path(first.local_path).read_bytes() == content
+ assert mapped.parts[1].part.media_type == "application/pdf"
+ assert mapped.parts[1].part.extension == ".pdf"
+ assert record.model_dump() == before
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "files", [[{"id": "F1"}, {"id": "F1"}], [{"name": "missing-id"}], "bad-list"]
+)
+async def test_ambiguous_or_malformed_files_never_become_complete(files):
+ storage = AsyncMock()
+ with pytest.raises(ValueError, match="Slack"):
+ async with map_record(message(files), "slack", storage):
+ pass
+ storage.read_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_wrong_blob_owner_stale_path_and_corrupt_bytes_fail_before_conversion():
+ record = message([{"id": "F1", "name": "safe.pdf"}])
+ valid = blob(record, b"expected", 0)
+ storage = AsyncMock()
+ for bad in (
+ valid.model_copy(update={"key": f"canonical/{uuid4()}/blobs/sha256/{valid.sha256}"}),
+ valid.model_copy(update={"source_path": "/files/9"}),
+ ):
+ with pytest.raises(ValueError):
+ async with map_record(record.model_copy(update={"blobs": (bad,)}), "slack", storage):
+ pass
+ storage.read_file.assert_not_called()
+ storage.read_file.return_value = b"wrong"
+ with pytest.raises(BlobIntegrityError):
+ async with map_record(record.model_copy(update={"blobs": (valid,)}), "slack", storage):
+ pass
+ with pytest.raises(ValueError, match="lacks a retained file"):
+ async with map_record(
+ record.model_copy(update={"completeness": "complete"}), "slack", storage
+ ):
+ pass
+
+
+async def test_slack_pdf_publication_records_partial_extraction(database, source, tmp_path):
+ """Real PDF conversion/SQL publication, with synthetic embedding and feed only."""
+ import fitz
+ from sqlalchemy import select
+
+ from airweave.adapters.storage.filesystem import FilesystemBackend
+ from airweave.core.logging import logger
+ from airweave.domains.entities.canonical.projection_store import current_extraction
+ from airweave.domains.entities.canonical.tests.helpers import (
+ bind_projection,
+ capture,
+ observation,
+ )
+ from airweave.domains.entities.canonical.tests.test_extraction_coverage import (
+ destination,
+ projector,
+ )
+ from airweave.domains.storage.file_service import FileService
+ from airweave.models import Entity
+
+ service, fence = source
+ binding = await bind_projection(database, fence, "slack")
+ storage = FilesystemBackend(tmp_path)
+ pdf = fitz.open()
+ page = pdf.new_page()
+ page.insert_text(
+ (40, 40), "Retained Slack PDF with useful searchable investment discussion. " * 4
+ )
+ content = pdf.tobytes()
+ pdf.close()
+ files = FileService(uuid4(), storage, sync_id=fence.sync_id)
+ retained = await files.store_canonical_blob(content, media_type="application/pdf")
+ item = observation(
+ identity=RecordIdentity(record_type="message", native_id="1.000001", container_id="C1"),
+ payload={
+ "ts": "1.000001",
+ "text": "Intact Slack discussion.",
+ "files": [
+ {"id": "F1", "name": "discussion.pdf", "mimetype": "application/pdf"},
+ {"id": "F2", "name": "missing.pdf", "mimetype": "application/pdf"},
+ ],
+ },
+ completeness="partial",
+ blobs=(retained.model_copy(update={"source_path": "/files/0"}),),
+ )
+ await capture(database, service, fence, item)
+ result = await projector(database, storage).batch(
+ fence.organization_id, fence.sync_id, "slack", destination(binding.collection_id), logger
+ )
+ assert result.published == 1 and result.failed == 0
+ async with database() as db:
+ row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id))
+ coverage = await current_extraction(
+ db, fence.organization_id, fence.sync_id, row.id, row.record_revision
+ )
+ assert coverage.status == "partial"
+ assert [part.outcome for part in coverage.parts] == [
+ "indexed",
+ "indexed",
+ "unavailable_original",
+ ]
+ assert coverage.parts[1].key == "file:F1"
+ assert row.completeness == "partial"
+ old_revision = row.record_revision
+
+ # Switch one retained inline original to explicit child ownership. The parent
+ # revision invalidates its previous PDF publication before children publish.
+ from airweave.domains.entities.canonical.requests import parent_container_key
+ from airweave.domains.entities.canonical.slack_files import SlackFileManifest, SlackFileOutcome
+
+ parent = item.identity
+ children = []
+ for index, native in enumerate(item.payload["files"]):
+ manifest = SlackFileManifest(
+ files=(
+ SlackFileOutcome(
+ index=0,
+ native_id=native["id"],
+ outcome="captured" if index == 0 else "unavailable",
+ reason=None if index == 0 else "access_denied",
+ ),
+ )
+ )
+ evidence = await files.store_canonical_blob(
+ manifest.model_dump_json().encode(), media_type="application/json"
+ )
+ refs = (evidence.model_copy(update={"role": "representation_manifest"}),)
+ if index == 0:
+ refs += (retained.model_copy(update={"source_path": ""}),)
+ children.append(
+ observation(
+ identity=RecordIdentity(
+ record_type="file",
+ native_id=native["id"],
+ container_id=parent_container_key(parent),
+ ),
+ parent=parent,
+ payload=native,
+ completeness="complete" if index == 0 else "partial",
+ blobs=refs,
+ )
+ )
+ await capture(
+ database,
+ service,
+ fence,
+ item.model_copy(
+ update={
+ "payload_schema_version": 2,
+ "blobs": (),
+ "completeness": "complete",
+ "descendant_visibility_fields": ("files",),
+ }
+ ),
+ *children,
+ )
+ async with database() as db:
+ parent_row = await db.scalar(select(Entity).where(Entity.native_id == parent.native_id))
+ assert parent_row.record_revision > old_revision
+ assert parent_row.indexed_revision != parent_row.record_revision
+ result = await projector(database, storage).batch(
+ fence.organization_id, fence.sync_id, "slack", destination(binding.collection_id), logger
+ )
+ assert result.published == 3 and result.failed == 0
+ async with database() as db:
+ rows = (await db.scalars(select(Entity).where(Entity.sync_id == fence.sync_id))).all()
+ coverage_by_id = {
+ row.native_id: await current_extraction(
+ db, fence.organization_id, fence.sync_id, row.id, row.record_revision
+ )
+ for row in rows
+ }
+ assert [part.key for part in coverage_by_id[parent.native_id].parts] == ["body"]
+ assert [part.outcome for part in coverage_by_id["F1"].parts] == ["indexed"]
+ assert [part.outcome for part in coverage_by_id["F2"].parts] == ["unavailable_original"]
+
+
+@pytest.mark.asyncio
+async def test_manifest_enrichment_requires_exact_native_identity_and_blob_evidence():
+ """Slack Connect stubs retain their native JSON; acquisition evidence supplies format."""
+ import json
+
+ record = message([{"id": "F1", "file_access": "check_file_info"}, {"id": "F2"}])
+ content = b"retained PDF fixture"
+ original = blob(record, content, 0)
+ payload = {
+ "version": 1,
+ "files": [
+ {
+ "index": 0,
+ "native_id": "F1",
+ "outcome": "captured",
+ "file": {"id": "F1", "name": "verified.pdf", "mimetype": "application/pdf"},
+ },
+ {"index": 1, "native_id": "F2", "outcome": "unavailable", "reason": "access_denied"},
+ ],
+ }
+
+ async def run(value, *, include_original=True):
+ data = json.dumps(value).encode()
+ evidence = blob(record, data, 9).model_copy(
+ update={"source_path": None, "role": "representation_manifest"}
+ )
+ source = record.model_copy(
+ update={"blobs": ((original,) if include_original else ()) + (evidence,)}
+ )
+ storage = AsyncMock()
+ storage.read_file.side_effect = lambda key, **_: data if key == evidence.key else content
+ async with map_record(source, "slack", storage) as mapped:
+ assert mapped.parts[1].entity.filename == "verified.pdf"
+ assert mapped.parts[1].part.media_type == "application/pdf"
+ assert mapped.parts[2].entity is None
+ assert source.payload == record.payload
+
+ await run(payload)
+ with pytest.raises(ValueError, match="contradicts"):
+ await run(payload, include_original=False)
+ payload["files"][0].update(outcome="unavailable", reason="access_denied")
+ with pytest.raises(ValueError, match="contradicts"):
+ await run(payload)
+ payload["files"][0].update(outcome="captured", reason=None)
+ payload["files"][1]["native_id"] = "F3"
+ with pytest.raises(ValueError, match="native message identities"):
+ await run(payload)
+
+
+@pytest.mark.asyncio
+async def test_child_owned_message_never_indexes_inline_attachments():
+ source = message([{"id": "F1", "name": "one.pdf"}]).model_copy(
+ update={"payload_schema_version": 2, "completeness": "complete"}
+ )
+ storage = AsyncMock()
+ async with map_record(source, "slack", storage) as mapped:
+ assert [part.part.key for part in mapped.parts] == ["body"]
+ storage.read_file.assert_not_called()
+ with pytest.raises(ValueError, match="cannot retain inline"):
+ async with map_record(
+ source.model_copy(update={"blobs": (blob(source, b"x", 0),)}), "slack", storage
+ ):
+ pass
+
+
+@pytest.mark.asyncio
+async def test_child_file_retained_and_unavailable_have_one_extraction_owner():
+ import json
+
+ from airweave.domains.entities.canonical.requests import parent_container_key
+
+ parent = message([]).identity
+ source = message([]).model_copy(
+ update={
+ "identity": RecordIdentity(
+ record_type="file", native_id="F1", container_id=parent_container_key(parent)
+ ),
+ "parent": parent,
+ "payload": {"id": "F1", "name": "one.pdf", "mimetype": "application/pdf"},
+ }
+ )
+ storage = AsyncMock()
+ for captured in (True, False):
+ content = b"fixture PDF bytes"
+ original = blob(source, content, 0).model_copy(update={"source_path": ""})
+ data = json.dumps(
+ {
+ "version": 1,
+ "files": [
+ {
+ "index": 0,
+ "native_id": "F1",
+ "outcome": "captured" if captured else "unavailable",
+ "reason": None if captured else "access_denied",
+ }
+ ],
+ }
+ ).encode()
+ manifest = blob(source, data, 0).model_copy(
+ update={"source_path": None, "role": "representation_manifest"}
+ )
+ current = source.model_copy(
+ update={"blobs": (manifest, original) if captured else (manifest,)}
+ )
+ contents = {manifest.key: data, original.key: content}
+ storage.read_file.side_effect = lambda key, contents=contents, **_: contents[key]
+ async with map_record(current, "slack", storage) as mapped:
+ assert len(mapped.parts) == 1 and mapped.parts[0].part.key == "file:F1"
+ assert (mapped.parts[0].entity is not None) == captured
+ with pytest.raises(ValueError, match="exact message parent"):
+ async with map_record(current.model_copy(update={"parent": None}), "slack", storage):
+ pass
+ if captured:
+ with pytest.raises(ValueError, match="root original"):
+ async with map_record(
+ current.model_copy(
+ update={
+ "blobs": (
+ manifest,
+ original.model_copy(update={"source_path": "/files/0"}),
+ )
+ }
+ ),
+ "slack",
+ storage,
+ ):
+ pass
diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py
new file mode 100644
index 0000000000..081d306af0
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py
@@ -0,0 +1,615 @@
+"""Actual Slack page adapter plus orchestrator/capture transactions; synthetic HTTP only."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.tests.test_capture_pipeline import components, orchestrator
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+from airweave.models.sync_cursor import SyncCursor
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackApiError, SlackSource
+
+ROOT = {"channels": [{"id": "C1", "name": "Synthetic channel"}]}
+MESSAGE = {"ts": "1", "reply_count": 1, "text": "Synthetic parent"}
+HISTORY = {"messages": [MESSAGE]}
+REPLIES = {"messages": [MESSAGE, {"ts": "1.1", "thread_ts": "1", "text": "Synthetic reply"}]}
+
+
+async def runner(database, source, responses, attempt=1, service=None, *, capture_files=False):
+ original_service, fence = source
+ service = service or original_service
+ ctx, _, runtime, bus = components(database, source)
+ connector = SlackSource(
+ auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock()
+ )
+ connector.slack_config = SlackConfig(
+ expected_team_id="T1", expected_user_id="U1", capture_files=capture_files
+ )
+ connector._get = AsyncMock(return_value={"ok": True, "team_id": "T1", "user_id": "U1"})
+ await connector.validate()
+ connector._get = AsyncMock(side_effect=responses)
+ pipeline = CanonicalCapturePipeline(
+ service,
+ database,
+ bus,
+ connector.canonical_record_types,
+ CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt),
+ connector.canonical_container_parents,
+ page_source=connector,
+ files=MagicMock(),
+ )
+ runtime.source = connector
+ runtime.canonical_capture = pipeline
+ instance = orchestrator(ctx, pipeline, runtime, None, bus)
+ instance.stream = None
+ return instance, connector, pipeline
+
+
+async def run(instance, *, checkpoint=True):
+ await instance._start_sync()
+ await instance._process_entities()
+ await instance._cleanup_orphaned_entities_if_needed()
+ if checkpoint:
+ await instance._save_cursor_data()
+
+
+async def saved(database):
+ async with database() as db:
+ rows = (await db.scalars(select(Entity))).all()
+ cursor = (await db.scalar(select(SyncCursor))).cursor_data
+ scans = (await db.scalars(select(CaptureScan))).all()
+ return rows, cursor, scans
+
+
+@pytest.mark.parametrize("capture_files", [False, True])
+async def test_cancel_after_history_resumes_pending_reply_without_repeating_history(
+ database, source, capture_files
+):
+ first, _, _ = await runner(
+ database, source, [ROOT, HISTORY, asyncio.CancelledError()], capture_files=capture_files
+ )
+ with pytest.raises(asyncio.CancelledError):
+ await run(first)
+ rows, cursor, scans = await saved(database)
+ assert {row.native_id for row in rows} == {"C1", "1"}
+ assert "canonical_checkpoint" not in cursor
+ message_scan = next(scan for scan in scans if scan.record_type == "message")
+ assert message_scan.continuation["pending_threads"] == ["1"]
+ second, connector, _ = await runner(
+ database, source, [ROOT, REPLIES], attempt=2, capture_files=capture_files
+ )
+ await run(second)
+ operations = [call.args[0].split("/")[-1] for call in connector._get.call_args_list]
+ assert operations[:2] == ["conversations.list", "conversations.replies"]
+ # Both attachment owners were freshly returned by the reply page; its SQL
+ # commit attests empty inventories without a redundant exact lookup.
+ assert len(operations) == 2
+ rows, cursor, _ = await saved(database)
+ assert {row.native_id for row in rows} == {"C1", "1", "1.1"}
+ assert next(row for row in rows if row.native_id == "1").record_revision == 1
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+
+
+async def test_lost_page_ack_reloads_committed_pending_threads(database, source):
+ class LostAck(CanonicalCaptureService):
+ async def commit_scan_page(self, db, request):
+ result = await super().commit_scan_page(db, request)
+ if request.scope.record_type == "message":
+ raise ConnectionError("synthetic lost commit acknowledgement")
+ return result
+
+ first, connector, _ = await runner(
+ database, source, [ROOT, HISTORY], service=LostAck(source[0].store)
+ )
+ with pytest.raises(ConnectionError, match="acknowledgement"):
+ await run(first)
+ assert connector._get.await_count == 2
+ second, connector, _ = await runner(database, source, [ROOT, REPLIES], attempt=2)
+ await run(second)
+ assert connector._get.call_args_list[-1].args[0].endswith("conversations.replies")
+
+
+async def test_expired_reply_cursor_restarts_whole_scope_and_never_partial_reconcile(
+ database, source
+):
+ first_reply = {
+ "messages": [{"ts": "1.old", "thread_ts": "1"}],
+ "response_metadata": {"next_cursor": "expired"},
+ }
+ first, _, _ = await runner(
+ database, source, [ROOT, HISTORY, first_reply, asyncio.CancelledError()]
+ )
+ with pytest.raises(asyncio.CancelledError):
+ await run(first)
+ rows, _, scans = await saved(database)
+ assert next(row for row in rows if row.native_id == "1.old").deleted_at is None
+ old_sweep = next(scan for scan in scans if scan.record_type == "message").sweep_id
+ second, connector, _ = await runner(
+ database, source, [ROOT, SlackApiError("invalid_cursor"), HISTORY, REPLIES], attempt=2
+ )
+ await run(second)
+ calls = connector._get.call_args_list
+ assert calls[1].args[1]["cursor"] == "expired"
+ assert calls[2].args[0].endswith("conversations.history") and "cursor" not in calls[2].args[1]
+ rows, _, scans = await saved(database)
+ assert next(row for row in rows if row.native_id == "1.old").deleted_at is not None
+ assert next(scan for scan in scans if scan.record_type == "message").sweep_id != old_sweep
+
+
+@pytest.mark.parametrize("error_code", ["invalid_cursor", "thread_not_found"])
+async def test_repeated_invalid_cursor_fails_without_checkpoint_or_absence(
+ database, source, error_code
+):
+ first, _, _ = await runner(
+ database,
+ source,
+ [ROOT, HISTORY, SlackApiError(error_code), HISTORY, SlackApiError(error_code)],
+ )
+ from airweave.domains.entities.canonical.page_source import InvalidScanContinuation
+
+ with pytest.raises(InvalidScanContinuation):
+ await run(first)
+ rows, cursor, _ = await saved(database)
+ assert all(row.deleted_at is None for row in rows)
+ assert "canonical_checkpoint" not in cursor
+
+
+@pytest.mark.parametrize("accessible", [False, True])
+async def test_membership_loss_on_retry_hides_pending_children_only_after_confirmation(
+ database, source, accessible
+):
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()])
+ with pytest.raises(asyncio.CancelledError):
+ await run(first)
+ responses = [
+ {"channels": []},
+ {"channel": {"id": "C1"}} if accessible else SlackApiError("channel_not_found"),
+ ]
+ second, connector, pipeline = await runner(database, source, responses, attempt=2)
+ if accessible:
+ with pytest.raises(ValueError, match="omitted"):
+ await run(second)
+ else:
+ await run(second)
+ assert connector._get.await_count == 2
+ rows, cursor, _ = await saved(database)
+ child = next(row for row in rows if row.native_id == "1")
+ async with database() as db:
+ value = await source[0].store.read(
+ db, source[1].organization_id, source[1].sync_id, child.id
+ )
+ assert value.content_access == ("available" if accessible else "unavailable")
+ assert cursor["canonical_cycle"]["phase"] == ("active" if accessible else "complete")
+
+
+async def test_crash_after_scans_then_lost_final_ack_same_job_does_not_start_new_cycle(
+ database, source
+):
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, REPLIES])
+ await run(first, checkpoint=False)
+ _, cursor, _ = await saved(database)
+ cycle_id = cursor["canonical_cycle"]["version"]["cycle_id"]
+ second, connector, _ = await runner(database, source, [ROOT], attempt=2)
+ await run(second)
+ assert connector._get.await_count == 1
+ third, connector, _ = await runner(database, source, [], attempt=3)
+ await run(third)
+ assert connector._get.await_count == 0
+ _, cursor, _ = await saved(database)
+ assert cursor["canonical_cycle"]["version"]["cycle_id"] == cycle_id
+
+
+async def test_new_job_starts_fresh_cycle_without_overriding_successful_old_job(database, source):
+ from airweave.domains.entities.canonical.store import StaleWriter
+ from airweave.models.sync_job import SyncJob
+
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, REPLIES])
+ await run(first)
+ _, cursor, _ = await saved(database)
+ prior_cycle = cursor["canonical_cycle"]["version"]["cycle_id"]
+ service, fence = source
+ new_job_id = uuid4()
+ async with database() as db:
+ old_job = await db.get(SyncJob, fence.job_id)
+ old_job.status = "completed"
+ db.add(
+ SyncJob(
+ id=new_job_id,
+ sync_id=fence.sync_id,
+ organization_id=fence.organization_id,
+ status="running",
+ )
+ )
+ await db.commit()
+ stale, _, _ = await runner(database, source, [], attempt=2)
+ with pytest.raises(StaleWriter, match="active job"):
+ await run(stale)
+ new_source = (service, fence.model_copy(update={"job_id": new_job_id, "attempt_id": uuid4()}))
+ new, connector, _ = await runner(database, new_source, [ROOT, HISTORY, REPLIES])
+ await run(new)
+ assert connector._get.await_count == 3
+ _, cursor, _ = await saved(database)
+ assert cursor["canonical_cycle"]["version"]["cycle_id"] != prior_cycle
+ assert cursor["canonical_cycle"]["completed_job_id"] == str(new_job_id)
+ async with database() as db:
+ assert (await db.get(SyncJob, fence.job_id)).status == "completed"
+
+
+async def test_access_lost_during_reply_withdraws_root_and_children(database, source):
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, SlackApiError("not_in_channel")])
+ await run(first)
+ rows, cursor, _ = await saved(database)
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+ assert all(row.removal_reason == "access_revoked" for row in rows)
+ async with database() as db:
+ for row in rows:
+ value = await source[0].store.read(
+ db, source[1].organization_id, source[1].sync_id, row.id
+ )
+ assert value.content_access == "unavailable" and value.payload == {}
+
+
+@pytest.mark.parametrize("threaded,count", [(False, 501), (True, 101)])
+async def test_oversized_page_or_pending_thread_queue_fails_before_commit(
+ database, source, threaded, count
+):
+ oversized = {
+ "messages": [
+ {
+ "ts": str(i),
+ "reply_count": int(threaded),
+ "text": "private fixture must not enter validation diagnostics",
+ }
+ for i in range(count)
+ ]
+ }
+ first, _, _ = await runner(database, source, [ROOT, oversized])
+ with pytest.raises(ValueError, match="invalid page") as error:
+ await run(first)
+ assert "private fixture" not in str(error.value)
+ rows, cursor, scans = await saved(database)
+ assert {row.native_id for row in rows} == {"C1"}
+ child = next(scan for scan in scans if scan.record_type == "message")
+ assert child.phase == "collecting" and child.revision == 1 and child.continuation == {}
+ assert "canonical_checkpoint" not in cursor
+
+
+async def test_interrupted_omission_confirmations_remove_nothing_then_retry(database, source):
+ from airweave.domains.entities.canonical.requests import RecordIdentity
+ from airweave.domains.entities.canonical.tests.helpers import capture, observation
+
+ await capture(
+ database,
+ source[0],
+ source[1],
+ *(
+ observation(identity=RecordIdentity(record_type="channel", native_id=value))
+ for value in ("C1", "C2")
+ ),
+ )
+ first, _, _ = await runner(
+ database,
+ source,
+ [
+ {"channels": []},
+ SlackApiError("channel_not_found"),
+ ConnectionError("confirmation failed"),
+ ],
+ )
+ with pytest.raises(ConnectionError, match="confirmation failed"):
+ await run(first)
+ rows, cursor, _ = await saved(database)
+ assert all(row.deleted_at is None for row in rows)
+ assert "canonical_checkpoint" not in cursor
+ second, connector, _ = await runner(
+ database,
+ source,
+ [{"channels": []}, SlackApiError("channel_not_found"), SlackApiError("channel_not_found")],
+ attempt=2,
+ )
+ await run(second)
+ assert connector._get.await_count == 3
+ rows, cursor, _ = await saved(database)
+ assert all(row.removal_reason == "scope_removed" for row in rows)
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+
+
+async def test_cleanup_budget_stop_preserves_reconciling_state_and_resumes(database, source):
+ from airweave.domains.entities.canonical.requests import RecordIdentity
+ from airweave.domains.entities.canonical.tests.helpers import capture, observation
+
+ parent = RecordIdentity(record_type="channel", native_id="C1")
+ await capture(
+ database,
+ source[0],
+ source[1],
+ observation(identity=parent),
+ *(
+ observation(
+ identity=RecordIdentity(record_type="message", native_id=str(i), container_id="C1"),
+ parent=parent,
+ )
+ for i in range(251)
+ ),
+ )
+ first, _, _ = await runner(database, source, [ROOT, {"messages": []}])
+ checks = 0
+
+ async def limit():
+ nonlocal checks
+ checks += 1
+ if checks == 5:
+ raise RuntimeError("synthetic cleanup budget reached")
+
+ first._check_capture_limits = limit
+ with pytest.raises(RuntimeError, match="cleanup budget"):
+ await run(first)
+ rows, cursor, scans = await saved(database)
+ messages = [row for row in rows if row.entity_definition_short_name == "message"]
+ assert sum(row.deleted_at is not None for row in messages) == 250
+ assert next(scan for scan in scans if scan.record_type == "message").phase == "reconciling"
+ assert "canonical_checkpoint" not in cursor
+ second, connector, _ = await runner(database, source, [ROOT], attempt=2)
+ await run(second)
+ assert connector._get.await_count == 1
+ rows, cursor, _ = await saved(database)
+ assert all(
+ row.deleted_at is not None for row in rows if row.entity_definition_short_name == "message"
+ )
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+
+
+@pytest.mark.parametrize("completed_child", [False, True])
+async def test_restored_root_restarts_withdrawn_child_scope_in_same_cycle(
+ database, source, completed_child
+):
+ first_responses = [
+ ROOT,
+ HISTORY,
+ REPLIES if completed_child else SlackApiError("not_in_channel"),
+ ]
+ first, _, _ = await runner(database, source, first_responses)
+ await run(first, checkpoint=False)
+ restore_attempt = 2
+ if completed_child:
+ removed, _, _ = await runner(
+ database, source, [{"channels": []}, SlackApiError("channel_not_found")], attempt=2
+ )
+ await run(removed, checkpoint=False)
+ restore_attempt = 3
+ before, cursor, scans = await saved(database)
+ assert all(row.deleted_at is not None for row in before)
+ old_sweep = next(scan for scan in scans if scan.record_type == "message").sweep_id
+ restored, connector, _ = await runner(
+ database, source, [ROOT, HISTORY, REPLIES], attempt=restore_attempt
+ )
+ await run(restored)
+ operations = [call.args[0].split("/")[-1] for call in connector._get.call_args_list]
+ assert operations == ["conversations.list", "conversations.history", "conversations.replies"]
+ after, _, scans = await saved(database)
+ assert {row.native_id for row in after} == {"C1", "1", "1.1"}
+ assert all(row.deleted_at is None for row in after)
+ assert next(scan for scan in scans if scan.record_type == "message").sweep_id != old_sweep
+
+
+async def test_root_revival_and_child_invalidation_roll_back_together(database, source):
+ from airweave.domains.entities.canonical.scan_store import CanonicalScanStore
+
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, REPLIES])
+ await run(first, checkpoint=False)
+ removed, _, _ = await runner(
+ database, source, [{"channels": []}, SlackApiError("channel_not_found")], attempt=2
+ )
+ await run(removed, checkpoint=False)
+ before, _, scans = await saved(database)
+ root_before = next(row for row in before if row.native_id == "C1")
+ child_before = next(scan for scan in scans if scan.record_type == "message")
+
+ class FailAfterPage(CanonicalScanStore):
+ async def page(self, db, request):
+ await super().page(db, request)
+ raise RuntimeError("synthetic failure after revival and invalidation")
+
+ service = CanonicalCaptureService(source[0].store)
+ service.scans = FailAfterPage(service.store)
+ restoring, _, _ = await runner(database, source, [ROOT], attempt=3, service=service)
+ with pytest.raises(RuntimeError, match="after revival"):
+ await run(restoring)
+ after, _, scans = await saved(database)
+ root_after = next(row for row in after if row.native_id == "C1")
+ child_after = next(scan for scan in scans if scan.record_type == "message")
+ assert root_after.deleted_at == root_before.deleted_at
+ assert root_after.record_revision == root_before.record_revision
+ assert child_after.sweep_id == child_before.sweep_id
+ assert child_after.revision == child_before.revision and child_after.phase == "complete"
+
+
+async def test_unavailable_parent_without_tombstone_also_invalidates_child_progress(
+ database, source
+):
+ first, _, _ = await runner(database, source, [ROOT, HISTORY, SlackApiError("not_in_channel")])
+ await run(first, checkpoint=False)
+ async with database() as db:
+ root = await db.scalar(select(Entity).where(Entity.native_id == "C1"))
+ root.deleted_at = None # Exercise the independent content-access gate representation.
+ await db.commit()
+ restored, connector, _ = await runner(database, source, [ROOT, HISTORY, REPLIES], attempt=2)
+ await run(restored)
+ assert connector._get.call_args_list[1].args[0].endswith("conversations.history")
+ rows, _, _ = await saved(database)
+ assert all(row.deleted_at is None and row.removal_reason is None for row in rows)
+
+
+async def test_scope_withdrawal_keeps_uncertain_reason_without_claiming_access_revocation(
+ database, source
+):
+ from airweave.domains.entities.canonical.page_source import ScopeAccessLost
+
+ instance, _, pipeline = await runner(database, source, [ROOT, HISTORY])
+ original_page = pipeline.page_source.capture_page
+
+ async def unavailable(scope, continuation, *, files, parent=None):
+ if scope.record_type == "message" and continuation.value.get("pending_threads"):
+ raise ScopeAccessLost(
+ "Selected scope is no longer available", removal_reason="scope_removed"
+ )
+ return await original_page(scope, continuation, files=files)
+
+ pipeline.page_source.capture_page = unavailable
+ await run(instance)
+ rows, _, _ = await saved(database)
+ assert all(row.removal_reason == "scope_removed" for row in rows)
+ async with database() as db:
+ for row in rows:
+ value = await source[0].store.read(
+ db, source[1].organization_id, source[1].sync_id, row.id
+ )
+ assert value.content_access == "unavailable" and value.payload == {}
+
+
+@pytest.mark.parametrize("change", ["rename", "revive"])
+async def test_old_provider_route_failure_cannot_withdraw_changed_same_fence_parent(
+ database, source, change
+):
+ from datetime import datetime, timezone
+
+ from airweave.domains.entities.canonical.page_source import ScopeAccessLost
+ from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord
+ from airweave.domains.entities.canonical.scan_store import ScanConflict
+
+ service, fence = source
+ instance, connector, pipeline = await runner(database, source, [ROOT])
+ original_page = connector.capture_page
+
+ async def change_during_network(scope, continuation, *, files, parent=None):
+ if parent is None:
+ return await original_page(scope, continuation, files=files, parent=parent)
+ update = CaptureRecord(
+ identity=parent.identity,
+ parent=parent.parent,
+ payload={**parent.payload, "name": "Current renamed channel"},
+ observed_at=datetime.now(timezone.utc),
+ )
+ records = (update,)
+ if change == "revive":
+ records = (
+ update.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}),
+ update,
+ )
+ async with database() as db:
+ await service.capture(db, CaptureBatch(fence=fence, records=records))
+ raise ScopeAccessLost("Old endpoint routing reported access loss")
+
+ connector.capture_page = change_during_network
+ with pytest.raises(ScanConflict, match="changed during provider access"):
+ await run(instance)
+ rows, cursor, scans = await saved(database)
+ current = next(row for row in rows if row.native_id == "C1")
+ assert current.deleted_at is None and current.removal_reason is None
+ assert current.source_payload["name"] == "Current renamed channel"
+ assert current.visibility_epoch == (2 if change == "revive" else 1)
+ assert "canonical_checkpoint" not in cursor
+ assert next(scan for scan in scans if scan.record_type == "message").phase == "collecting"
+
+
+async def test_missing_queued_thread_restarts_inventory_before_absence(database, source):
+ instance, connector, _ = await runner(
+ database,
+ source,
+ [
+ ROOT,
+ HISTORY,
+ SlackApiError("thread_not_found"),
+ {"messages": [{"ts": "2", "text": "Current message"}]},
+ ],
+ )
+ await run(instance)
+ assert [call.args[0].split("/")[-1] for call in connector._get.call_args_list] == [
+ "conversations.list",
+ "conversations.history",
+ "conversations.replies",
+ "conversations.history",
+ ]
+ assert "cursor" not in connector._get.call_args_list[-1].args[1]
+ rows, cursor, scans = await saved(database)
+ old = next(row for row in rows if row.native_id == "1")
+ assert old.deleted_at is not None and old.removal_reason == "absent"
+ assert next(row for row in rows if row.native_id == "2").deleted_at is None
+ assert next(row for row in rows if row.native_id == "C1").deleted_at is None
+ assert cursor["canonical_cycle"]["phase"] == "complete"
+ assert all(scan.phase == "complete" for scan in scans)
+
+
+async def test_native_slack_dates_reach_persisted_search_fields_and_date_filters(database, source):
+ from datetime import datetime, timezone
+
+ from airweave.domains.entities.canonical.projection_mappers import map_record
+ from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata
+ from airweave.domains.entities.canonical.store import source_record
+ from airweave.domains.search.owned import OwnedSearchService
+ from airweave.domains.search.owned_models import OwnedSearchRequest
+ from airweave.platform.destinations.vespa.transformer import EntityTransformer
+ from airweave.platform.entities._base import AirweaveSystemMetadata
+
+ first, _, _ = await runner(
+ database,
+ source,
+ [
+ ROOT,
+ {
+ "messages": [
+ {"ts": "1767225600.000001", "text": "Before editing"},
+ {
+ "ts": "1767225600.000002",
+ "text": "Edited",
+ "edited": {"ts": "1767312000.000003"},
+ },
+ ]
+ },
+ ],
+ )
+ await run(first)
+ rows, _, _ = await saved(database)
+ unchanged = next(row for row in rows if row.native_id == "1767225600.000001")
+ edited = next(row for row in rows if row.native_id == "1767225600.000002")
+ assert unchanged.source_created_at == datetime(2026, 1, 1, microsecond=1, tzinfo=timezone.utc)
+ assert unchanged.source_updated_at is None
+ assert edited.source_updated_at == datetime(2026, 1, 2, microsecond=3, tzinfo=timezone.utc)
+
+ # Exercise the same canonical metadata stamper and Vespa serializer used by projection.
+ async with map_record(source_record(edited), "slack", MagicMock()) as entities:
+ entities = entities.entities
+ (entity,) = entities
+ entity.airweave_system_metadata = AirweaveSystemMetadata()
+ stamp_search_metadata(entity.airweave_system_metadata, source_record(edited))
+ fields = EntityTransformer()._build_system_metadata(entity)
+ assert fields["source_created_us"] == 1767225600000002
+ assert fields["source_updated_us"] == 1767312000000003
+ assert fields["source_updated_known"] == 1
+ assert fields["canonical_record_type"] == "message"
+
+ # Current-SQL postvalidation enforces inclusive lower and exclusive upper bounds.
+ request = OwnedSearchRequest(
+ query="Edited",
+ sync_ids=(edited.sync_id,),
+ created_after=datetime(2026, 1, 1, microsecond=2, tzinfo=timezone.utc),
+ updated_after=datetime(2026, 1, 2, microsecond=3, tzinfo=timezone.utc),
+ )
+ assert OwnedSearchService._matches(edited, request)
+ assert not OwnedSearchService._matches(unchanged, request)
+ assert not OwnedSearchService._matches(
+ edited, request.model_copy(update={"updated_before": edited.source_updated_at})
+ )
+ stamp_search_metadata(entity.airweave_system_metadata, source_record(unchanged))
+ fields = EntityTransformer()._build_system_metadata(entity)
+ assert fields["source_updated_known"] == 0
+ assert fields["source_updated_us"] is None
diff --git a/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py
new file mode 100644
index 0000000000..55eb2877b8
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py
@@ -0,0 +1,70 @@
+"""Failure injection at conversion, chunking, embedding and remote-feed boundaries."""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from airweave.domains.converters.fakes.registry import FakeConverterRegistry
+from airweave.domains.sync_pipeline.exceptions import EntityProcessingError
+from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor
+from airweave.platform.destinations.vespa.destination import VespaDestination
+from airweave.platform.entities._base import AirweaveSystemMetadata
+from airweave.platform.entities.slack import SlackMessageEntity
+
+
+@pytest.mark.parametrize("stage", ["conversion", "chunking", "embedding"])
+async def test_strict_processor_rejects_dropped_required_content(stage):
+ processor = ChunkEmbedProcessor(FakeConverterRegistry(), MagicMock(), MagicMock())
+ entity = SlackMessageEntity.from_api({"ts": "1.0", "text": "Required content"}, breadcrumbs=[])
+ entity.entity_id = "part"
+ entity.airweave_system_metadata = AirweaveSystemMetadata()
+ chunk = entity.model_copy(deep=True)
+ chunk.entity_id = "part__chunk_0"
+ chunk.airweave_system_metadata.original_entity_id = "part"
+ entity.textual_representation = "Required content"
+ processor._text_builder.build_for_batch = AsyncMock(
+ return_value=[] if stage == "conversion" else [entity]
+ )
+ processor._chunk_entities = AsyncMock(return_value=[] if stage == "chunking" else [chunk])
+ processor._embed_entities = AsyncMock(return_value=[] if stage == "embedding" else [chunk])
+ with pytest.raises(EntityProcessingError, match="dropped required"):
+ await processor.process(
+ [entity],
+ SimpleNamespace(logger=MagicMock()),
+ SimpleNamespace(entity_tracker=AsyncMock()),
+ strict=True,
+ )
+
+
+async def test_strict_destination_cannot_ack_dropped_transformations():
+ destination = VespaDestination()
+ destination._client = AsyncMock()
+ destination._transformer = MagicMock()
+ destination._transformer.transform_batch.return_value = {}
+ with pytest.raises(RuntimeError, match="dropped required"):
+ await destination.bulk_insert([MagicMock()], strict=True)
+ destination._client.feed_documents.assert_not_called()
+
+
+async def test_strict_delete_retries_partial_http_failure_and_accepts_absence():
+ from airweave.platform.destinations.vespa.client import VespaClient
+
+ client = VespaClient(MagicMock())
+ transport = MagicMock()
+ transport.delete = AsyncMock(
+ side_effect=[
+ SimpleNamespace(status_code=200),
+ SimpleNamespace(status_code=503),
+ ]
+ )
+ documents = [("base_entity", "first"), ("base_entity", "second")]
+ with pytest.raises(RuntimeError, match="deletion incomplete"):
+ await client._delete_by_doc_ids(documents, transport, strict=True)
+ transport.delete = AsyncMock(
+ side_effect=[
+ SimpleNamespace(status_code=404),
+ SimpleNamespace(status_code=200),
+ ]
+ )
+ assert await client._delete_by_doc_ids(documents, transport, strict=True) == 2
diff --git a/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py
new file mode 100644
index 0000000000..90ad830f82
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py
@@ -0,0 +1,122 @@
+"""Stripe source composition and durable recovery; synthetic HTTP, real PostgreSQL."""
+
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import httpx
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.cycle_models import CompleteCycle
+from airweave.domains.sources.exceptions import SourceEntityForbiddenError
+from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider
+from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+from airweave.platform.configs.config import StripeCaptureConfig, StripeConfig
+from airweave.platform.sources.stripe import StripeSource
+
+
+async def connector(calls, *, fail=False):
+ async def get(url, **kwargs):
+ params = kwargs["params"]
+ calls.append((url, params))
+ status = 200
+ if url.endswith("/account"):
+ data = {"object": "account", "id": "acct_selected"}
+ elif url.endswith("/balance"):
+ data = {"object": "balance", "livemode": False, "available": []}
+ elif url.endswith("/customers"):
+ second = "starting_after" in params
+ if second and fail:
+ status, data = 403, {"error": {"message": "private provider detail"}}
+ else:
+ data = {
+ "object": "list",
+ "has_more": not second,
+ "data": [
+ {
+ "id": "cus_second" if second else "cus_first",
+ "object": "customer",
+ "livemode": False,
+ "unknown": {"retained": [1, None]},
+ }
+ ],
+ }
+ else:
+ data = {"object": "list", "data": [], "has_more": False}
+ return httpx.Response(status, json=data, request=httpx.Request("GET", url))
+
+ http = AsyncMock()
+ http.get.side_effect = get
+
+ @asynccontextmanager
+ async def stream(method, url, **kwargs):
+ assert method == "GET"
+ yield await get(url, **kwargs)
+
+ http.stream = MagicMock(side_effect=stream)
+ source = await StripeSource.create(
+ auth=ManagedAuthProvider(
+ api_key="fixture",
+ connected_account_id="fixture",
+ allowed_hosts=frozenset({"api.stripe.com"}),
+ ),
+ logger=MagicMock(),
+ http_client=http,
+ config=StripeConfig(
+ original_capture=StripeCaptureConfig(
+ expected_account_id="acct_selected",
+ livemode=False,
+ api_version="2025-06-30.basil",
+ )
+ ),
+ )
+ return source.capture_page_source
+
+
+async def test_failed_customer_page_recovers_durable_cursor_and_preserves_original(
+ database, source
+):
+ service, fence = source
+ first = await connector([], fail=True)
+ with pytest.raises(SourceEntityForbiddenError):
+ await CanonicalScanDriver(
+ service, database, fence, first, AsyncMock(), AsyncMock(), MagicMock()
+ ).run()
+ async with database() as db:
+ scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "customer"))
+ assert scan.phase == "collecting"
+ assert scan.continuation["starting_after"] == "cus_first"
+ before = await service.read_cycle(db, fence)
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ calls = []
+ resumed = await connector(calls)
+ after = await CanonicalScanDriver(
+ service, database, newer, resumed, AsyncMock(), AsyncMock(), MagicMock()
+ ).run()
+ async with database() as db:
+ after = await service.complete_cycle(db, CompleteCycle(fence=newer, expected=after.version))
+ records = list(
+ (
+ await db.scalars(
+ select(Entity).where(Entity.entity_definition_short_name == "customer")
+ )
+ ).all()
+ )
+ assert after.version.cycle_id == before.version.cycle_id and after.phase == "complete"
+ assert after.last_full_capture.discovery == "incomplete"
+ assert [
+ params.get("starting_after") for url, params in calls if url.endswith("/customers")
+ ] == ["cus_first"]
+ assert len(records) == 2
+ assert all(row.record_revision == 1 and row.deleted_at is None for row in records)
+ assert all(row.source_payload["unknown"] == {"retained": [1, None]} for row in records)
diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py
new file mode 100644
index 0000000000..26cb650819
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py
@@ -0,0 +1,92 @@
+"""Safe metadata around actual Wispr validation; no network."""
+
+from airweave.domains.entities.canonical.tests.test_resume_probe import subprocess_script
+
+
+async def test_safe_validation_failure_metadata_without_payload_leak():
+ script = r"""
+import asyncio, json, os, sys
+from types import SimpleNamespace
+from uuid import uuid4
+sys.path.insert(0, "tests/live")
+import conftest
+import httpx
+import provider_sample
+from wispr_diagnostics import WisprDiagnostics
+from airweave.domains.sources.exceptions import SourceServerError
+
+async def main():
+ original_client = httpx.AsyncClient
+ for error, kind, rate in [("private-detail rate limit", "string", True),
+ ({"status_code": 429, "private": "secret"}, "object", True),
+ ("private-detail failure", "string", False)]:
+ diagnostics = WisprDiagnostics()
+ calls = []
+ async def handle(request):
+ calls.append(diagnostics.operation)
+ if diagnostics.operation == "metadata":
+ value = {"id": "fixture", "status": "ACTIVE", "user_id": "fixture",
+ "toolkit": {"slug": "wispr_flow_mcp"}}
+ elif diagnostics.operation == "session":
+ value = {"session_id": "fixture"}
+ else:
+ value = {"data": {}, "error": error}
+ return httpx.Response(200, json=value)
+ class MockClient(original_client):
+ def __init__(self, **kwargs):
+ super().__init__(transport=httpx.MockTransport(handle), **kwargs)
+ provider_sample.httpx.AsyncClient = MockClient
+ async def request(req): diagnostics.request(req)
+ os.environ["LIVE_WISPR_USER_ID"] = "fixture"
+ try:
+ async with provider_sample.wispr_source("fixture", "fixture",
+ SimpleNamespace(organization_id=uuid4()), request_hook=request,
+ response_hook=diagnostics.response, envelope_hook=diagnostics.envelope):
+ raise AssertionError("Validation error swallowed")
+ except (ValueError, SourceServerError) as exc:
+ if kind == "string" and rate:
+ assert isinstance(exc, SourceServerError)
+ assert "rate-limit signal" in str(exc)
+ assert exc.status_code is None
+ else:
+ assert str(exc) == "Wispr tool execution failed; capture is incomplete"
+ assert "private" not in str(exc) and "secret" not in str(exc)
+ assert calls == ["metadata", "session", "search"]
+ assert diagnostics.operation == "search" and diagnostics.last_http_status == 200
+ assert diagnostics.tool_error_present and diagnostics.tool_error_kind == kind
+ assert diagnostics.rate_signal_detected == rate
+ assert "private" not in diagnostics.model_dump_json()
+ assert "secret" not in diagnostics.model_dump_json()
+ req = httpx.Request("POST", "https://backend.composio.dev/api/v3.1/tool_router/session/x/execute",
+ json={"tool_slug": "WISPR_FLOW_MCP_GET_MEETING",
+ "arguments": {"id": "private"}})
+ diagnostics.request(req)
+ await diagnostics.response(httpx.Response(429))
+ assert diagnostics.operation == "get" and diagnostics.rate_signal_detected
+ assert diagnostics.tool_error_present is None
+ print("verified")
+asyncio.run(main())
+"""
+ code, output, error = await subprocess_script(script, {})
+ assert code == 0, error
+ assert output.splitlines()[-1] == "verified"
+
+
+async def test_reduced_wispr_budget_validation_before_network():
+ script = r"""
+import sys
+sys.path.insert(0, "tests/live")
+import conftest
+from provider_lifecycle import wispr_request_limit
+for value in ("0", "21", "-1", "no", "1.5"):
+ try: wispr_request_limit(value)
+ except ValueError: pass
+ else: raise AssertionError("Invalid budget accepted")
+assert wispr_request_limit("17") == 17
+assert wispr_request_limit("20") == 20
+assert wispr_request_limit("1") == 1
+print("verified")
+"""
+ code, output, error = await subprocess_script(script, {})
+ assert code == 0, error
+ assert output.splitlines()[-1] == "verified"
diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py
new file mode 100644
index 0000000000..78a3b75eea
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py
@@ -0,0 +1,337 @@
+"""Wispr resumes completed bodies through the real canonical forest, without absence claims."""
+
+from datetime import datetime, timezone
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord, RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalStoreError
+from airweave.domains.sources.exceptions import SourceServerError
+from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider
+from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
+from airweave.platform.sources.wispr import WisprSource
+
+
+async def connector(rows, body_calls, *, fail_at=None, failure=None):
+ result = await WisprSource.create(
+ auth=ManagedToolAuthProvider(
+ api_key="fixture", connected_account_id="fixture", user_id="fixture"
+ ),
+ logger=MagicMock(),
+ http_client=MagicMock(),
+ )
+
+ async def execute(slug, arguments):
+ if slug == "WISPR_FLOW_MCP_SEARCH_SCRATCHPAD_NOTES":
+ return {"notes": [], "has_more": False}
+ if slug == "WISPR_FLOW_MCP_SEARCH_MEETINGS":
+ return {"meetings": rows, "has_more": False}
+ native = arguments["meeting_id"]
+ body_calls.append(native)
+ if len(body_calls) == fail_at:
+ raise failure or ValueError("Synthetic opaque body failure")
+ return {
+ "id": native,
+ "content": "notes",
+ "transcript": "transcript",
+ "modified_at": "2026-09-30T00:00:00Z",
+ "start": "2026-09-20T00:00:00Z",
+ }
+
+ result._execute = execute
+ return result
+
+
+def driver(service, database, fence, source):
+ return CanonicalScanDriver(
+ service, database, fence, source, AsyncMock(), AsyncMock(), MagicMock()
+ )
+
+
+@pytest.mark.parametrize(
+ "failure",
+ [
+ ValueError("Synthetic opaque body failure"),
+ SourceServerError("Synthetic rate signal; retry delay unknown", source_short_name="wispr"),
+ ],
+)
+async def test_failed_body_resumes_without_repeating_completed_siblings(database, source, failure):
+ service, fence = source
+ rows = [
+ {"id": f"m{i}", "title": "Earlier listing", "start": "2026-09-20T00:00:00Z"}
+ for i in range(10)
+ ]
+ calls = []
+ first = await connector(rows, calls, fail_at=10, failure=failure)
+ with pytest.raises(type(failure), match="Synthetic"):
+ await driver(service, database, fence, first).run()
+ assert len(calls) == 10
+ failed = calls[-1]
+ async with database() as db:
+ scans = list(
+ (
+ await db.scalars(select(CaptureScan).where(CaptureScan.record_type == "meeting"))
+ ).all()
+ )
+ assert sum(scan.phase == "complete" for scan in scans) == 9
+ assert sum(scan.phase == "collecting" for scan in scans) == 1
+ before = await service.read_cycle(db, fence)
+ async with database() as db:
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ # Listing changes must not masquerade as body freshness.
+ changed = [
+ {**row, "title": "Updated listing", "modified_at": "2026-10-01T00:00:00Z"}
+ for row in rows
+ if row["id"] != calls[0]
+ ]
+ second = await connector(changed, calls)
+ after = await driver(service, database, newer, second).run()
+ assert before.version.cycle_id == after.version.cycle_id
+ assert len(calls) == 11 and calls[-1] == failed
+ async with database() as db:
+ bodies = list(
+ (
+ await db.scalars(
+ select(Entity).where(Entity.entity_definition_short_name == "meeting")
+ )
+ ).all()
+ )
+ assert len(bodies) == 10
+ assert all(row.deleted_at is None for row in bodies)
+ assert all(row.source_created_at is None for row in bodies)
+ assert all(
+ row.source_updated_at == datetime(2026, 9, 30, tzinfo=timezone.utc) for row in bodies
+ )
+ # Inventory refresh does not rewrite the body captured on the prior attempt.
+ assert (
+ sum(row.source_payload["listing"]["title"] == "Earlier listing" for row in bodies) == 9
+ )
+
+
+async def test_old_parentless_body_requires_explicit_migration(database, source):
+ service, fence = source
+ async with database() as db:
+ await service.capture(
+ db,
+ CaptureBatch(
+ fence=fence,
+ records=(
+ CaptureRecord(
+ identity=RecordIdentity(record_type="meeting", native_id="old"),
+ payload={"prior": True},
+ completeness="partial",
+ observed_at=datetime.now(timezone.utc),
+ ),
+ ),
+ ),
+ )
+ current = await connector([{"id": "old", "start": "2026-09-20T00:00:00Z"}], [])
+ with pytest.raises(CanonicalStoreError, match="different parent"):
+ await driver(service, database, fence, current).run()
+ async with database() as db:
+ body = await db.scalar(
+ select(Entity).where(Entity.entity_definition_short_name == "meeting")
+ )
+ assert body.source_payload == {"prior": True}
+ assert body.parent_record_type is None
+
+
+async def test_scratchpad_failure_resumes_failed_body_without_replaying_completed_notes(
+ database, source
+):
+ service, fence = source
+ listings = [
+ {"id": "same-id", "title": "First", "modified_at": "2026-09-30T00:00:00Z"},
+ {"id": "n2", "title": "Second", "modified_at": "2026-09-30T00:00:00Z"},
+ ]
+ calls = []
+ failing = True
+ completed: set[str] = set()
+ failed_native: str | None = None
+
+ async def execute(slug, arguments):
+ nonlocal failed_native
+ if slug == "WISPR_FLOW_MCP_SEARCH_MEETINGS":
+ # Native IDs may collide between resource kinds without aliasing storage.
+ return {"meetings": [{"id": "same-id"}], "has_more": False}
+ if slug == "WISPR_FLOW_MCP_SEARCH_SCRATCHPAD_NOTES":
+ return {"notes": listings, "has_more": False}
+ if slug == "WISPR_FLOW_MCP_GET_MEETING":
+ return {"id": "same-id", "content": "meeting", "transcript": "verbatim"}
+ assert slug == "WISPR_FLOW_MCP_GET_SCRATCHPAD_NOTE"
+ native = arguments["note_id"]
+ offset = arguments["view_content"]["start_char"]
+ calls.append((native, offset))
+ if failing and len(completed) == 1 and offset == 3:
+ failed_native = native
+ raise SourceServerError("Synthetic Wispr rate signal", source_short_name="wispr")
+ if offset == 3:
+ completed.add(native)
+ return {
+ "id": native,
+ "modified_at": "2026-09-30T00:00:00Z",
+ "content": (
+ "abc\n(...truncated, 3 chars remaining; continue with view_content.start_char=3...)"
+ )
+ if offset == 0
+ else "def",
+ }
+
+ first = await connector([], [])
+ first._execute = execute
+ with pytest.raises(SourceServerError):
+ await driver(service, database, fence, first).run()
+ async with database() as db:
+ notes = list(
+ (
+ await db.scalars(
+ select(Entity).where(Entity.entity_definition_short_name == "scratchpad_note")
+ )
+ ).all()
+ )
+ assert len(notes) == 1 and notes[0].native_id in completed
+ assert notes[0].record_revision == 1
+ before = await service.read_cycle(db, fence)
+ async with database() as db:
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ assert failed_native is not None and failed_native not in completed
+ failing = False
+ calls.clear()
+ second = await connector([], [])
+ second._execute = execute
+ after = await driver(service, database, newer, second).run()
+ assert before.version.cycle_id == after.version.cycle_id
+ # Durable scope order follows stored UUIDs, not the provider listing order.
+ assert calls == [
+ (failed_native, 0),
+ (failed_native, 3),
+ ] # Bodies commit atomically, ranges do not checkpoint separately.
+ async with database() as db:
+ rows = list(
+ (
+ await db.scalars(
+ select(Entity).where(
+ Entity.entity_definition_short_name.in_(["meeting", "scratchpad_note"])
+ )
+ )
+ ).all()
+ )
+ assert len(rows) == 3 and all(
+ row.deleted_at is None and row.record_revision == 1 for row in rows
+ )
+ notes = [row for row in rows if row.entity_definition_short_name == "scratchpad_note"]
+ assert all(len(row.source_payload["responses"]) == 2 for row in notes)
+ assert all(row.completeness == "partial" for row in notes)
+
+
+async def test_active_meeting_only_cycle_requires_explicit_restart_for_scratchpad(database, source):
+ import hashlib
+ import json
+
+ from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration
+ from airweave.domains.entities.canonical.cycle_store import CycleConflict
+
+ service, fence = source
+ old_fingerprint = hashlib.sha256(
+ json.dumps(
+ {
+ "version": 2,
+ "account": "fixture",
+ "inventory": "meeting_listing",
+ "body": "meeting",
+ "policy": "discovery_only",
+ "listing_limit": 200,
+ },
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+ old = CycleConfiguration.from_source(
+ fingerprint=old_fingerprint,
+ record_types=("meeting_listing", "meeting"),
+ container_parents={"meeting": "meeting_listing"},
+ completion_policies={"meeting_listing": "discovery_only"},
+ )
+ async with database() as db:
+ before = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=old))
+ current = await connector([], [])
+ with pytest.raises(CycleConflict, match="configuration changed"):
+ await driver(service, database, fence, current).run()
+ async with database() as db:
+ after = await service.read_cycle(db, fence)
+ assert after == before
+
+
+async def test_body_byte_overflow_preserves_committed_siblings_and_resume(
+ database, source, monkeypatch
+):
+ from airweave.domains.sources.exceptions import SourceError
+ from airweave.platform.sources import wispr
+
+ service, fence = source
+ rows = [{"id": f"bounded{i}", "start": "2026-09-20T00:00:00Z"} for i in range(2)]
+ calls = []
+ first = await connector(rows, calls)
+ execute = first._execute
+
+ async def oversized_second(slug, arguments):
+ result = await execute(slug, arguments)
+ if "meeting_id" in arguments and len(calls) == 2:
+ result["native_metadata"] = "x" * 2048
+ return result
+
+ first._execute = oversized_second
+ monkeypatch.setattr(wispr, "MAX_BODY_BYTES", 1024)
+ with pytest.raises(SourceError, match="body exceeds byte limit"):
+ await driver(service, database, fence, first).run()
+ async with database() as db:
+ states = list(
+ (
+ await db.scalars(select(CaptureScan).where(CaptureScan.record_type == "meeting"))
+ ).all()
+ )
+ assert sorted(state.phase for state in states) == ["collecting", "complete"]
+ bodies = list(
+ (
+ await db.scalars(
+ select(Entity).where(Entity.entity_definition_short_name == "meeting")
+ )
+ ).all()
+ )
+ assert len(bodies) == 1 and bodies[0].native_id == calls[0]
+ newer = await service.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ second = await connector(rows, calls)
+ completed = await driver(service, database, newer, second).run()
+ from airweave.domains.entities.canonical.cycle_models import CompleteCycle
+
+ async with database() as db:
+ completed = await service.complete_cycle(
+ db, CompleteCycle(fence=newer, expected=completed.version)
+ )
+ assert completed.phase == "complete"
+ assert len(calls) == 3 and calls[1] == calls[2]
diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py
new file mode 100644
index 0000000000..1d8b9aa6d2
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py
@@ -0,0 +1,224 @@
+"""Actual process exit/restart around committed Wispr pages; synthetic provider only."""
+
+import json
+
+from sqlalchemy import text
+
+from airweave.domains.entities.canonical.tests.test_resume_probe import subprocess_script
+
+SCRIPT = r"""
+import asyncio, os, sys
+from pathlib import Path
+from uuid import uuid4
+sys.path.insert(0, "tests/live")
+import conftest
+from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
+from airweave.domains.entities.canonical.requests import WriterFence
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.test_wispr_recovery import connector, driver
+from wispr_resume import WisprResumeProbe
+from provider_lifecycle import BoundedPageSource
+from airweave.domains.entities.canonical.tests.test_capture_pipeline import components, orchestrator
+from airweave.domains.entities.canonical.tests.test_slack_recovery import run
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
+from unittest.mock import MagicMock
+
+async def main():
+ engine = create_async_engine(os.environ["CANONICAL_TEST_DATABASE_URL"],
+ connect_args={"server_settings": {"search_path": os.environ["TEST_SCHEMA"]}})
+ sessions = async_sessionmaker(engine, expire_on_commit=False)
+ fence = WriterFence.model_validate_json(os.environ["TEST_FENCE"])
+ service = CanonicalCaptureService(CanonicalRecordStore())
+ mode = os.environ["TEST_MODE"]
+ if mode == "resume":
+ async with sessions() as db:
+ fence = await service.activate_writer(db, fence.organization_id, fence.sync_id,
+ fence.job_id, attempt_id=uuid4(), attempt_number=2)
+ calls = []
+ source = await connector([{"id": str(i)} for i in range(6)], calls,
+ fail_at=2 if mode == "fail" else None)
+ if os.environ.get("TEST_MIXED") == "1":
+ async def mixed(slug, arguments):
+ rows=[{"id":"private-same-id-"+str(i)} for i in range(3)]
+ if slug == "WISPR_FLOW_MCP_SEARCH_MEETINGS":
+ return {"meetings":rows,"has_more":False}
+ if slug == "WISPR_FLOW_MCP_SEARCH_SCRATCHPAD_NOTES":
+ return {"notes":rows,"has_more":False}
+ if slug == "WISPR_FLOW_MCP_GET_MEETING":
+ return {"id":arguments["meeting_id"],"content":"private meeting text",
+ "transcript":"private transcript text"}
+ assert slug == "WISPR_FLOW_MCP_GET_SCRATCHPAD_NOTE"
+ return {"id":arguments["note_id"],"content":"private scratchpad text"}
+ source._execute = mixed
+ counters = {"provider_requests": 0, "records_observed": 0}
+ execute = source._execute
+ async def counted(*args):
+ counters["provider_requests"] += 1
+ return await execute(*args)
+ source._execute = counted
+ probe = WisprResumeProbe(sessions, fence.organization_id, fence.sync_id, fence.job_id,
+ fence.attempt_id, Path(os.environ["TEST_TARGET"]),
+ "interrupt" if mode == "fail" else mode, counters)
+ ctx, _, runtime, bus = components(sessions, (service, fence))
+ wrapped = BoundedPageSource(source, counters, 600, probe)
+ if mode == "wrong_topology":
+ parents = {}
+ else:
+ parents = source.canonical_container_parents
+ pipeline = CanonicalCapturePipeline(service, sessions, bus, source.canonical_record_types,
+ CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), parents,
+ page_source=wrapped, files=MagicMock())
+ runtime.source, runtime.canonical_capture = source, pipeline
+ instance = orchestrator(ctx, pipeline, runtime, None, bus)
+ instance.stream = None
+ await run(instance)
+asyncio.run(main())
+"""
+
+
+async def test_wispr_process_restart_skips_completed_bodies(database, source, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ env = {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(tmp_path / "target.json"),
+ }
+ code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "interrupt"})
+ assert code == 75, error
+ first = json.loads(output.splitlines()[-1])
+ assert first["completed_bodies"] == 3 and first["provider_requests"] == 5
+ code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "resume"})
+ assert code == 76, error
+ second = json.loads(output.splitlines()[-1])
+ assert second["wispr_recovery_verified"] and second["provider_requests"] == 3
+ assert not second["full_scope_completed"]
+ assert (tmp_path / "target.json").stat().st_mode & 0o077 == 0
+
+
+async def test_wispr_failure_does_not_claim_interruption(database, source, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ code, output, error = await subprocess_script(
+ SCRIPT,
+ {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(tmp_path / "target.json"),
+ "TEST_MODE": "fail",
+ },
+ )
+ assert code not in (75, 76)
+ assert "Synthetic opaque body failure" in error
+ assert not (tmp_path / "target.json").exists()
+
+
+async def test_wispr_trial_budget_is_aggregate_and_failure_stops(tmp_path):
+ script = r"""
+import asyncio, os, sys
+sys.path.insert(0, "tests/live")
+import conftest
+import provider_lifecycle as lifecycle
+async def main():
+ for fail in (False, True):
+ snapshots = []
+ async def execute(path, timeout):
+ import json
+ value = json.loads(path.read_text())
+ snapshots.append(value)
+ if fail:
+ return 1, {"failed": True, "provider_requests": 2, "records_observed": 0}
+ second = len(snapshots) == 2
+ return (76 if second else 75), {"failed": False,
+ "intentional_interruption": not second, "wispr_recovery_verified": second,
+ "provider_requests": 7, "records_observed": 203}
+ lifecycle.execute_trial = execute
+ results = await lifecycle.run_trials({"provider": "wispr", "wispr_resume": True,
+ "root": os.environ["TEST_ROOT"], "request_limit": 20,
+ "record_limit": 600, "timeout": 180})
+ if fail:
+ assert len(snapshots) == 1
+ else:
+ assert len(snapshots) == 2
+ assert [v["request_limit"] for v in snapshots] == [20, 13]
+ assert [v["record_limit"] for v in snapshots] == [600, 397]
+ assert snapshots[0]["job_id"] == snapshots[1]["job_id"]
+ assert [v["attempt_number"] for v in snapshots] == [1, 2]
+ assert 180 >= snapshots[0]["timeout"] >= snapshots[1]["timeout"] > 0
+ print("verified")
+asyncio.run(main())
+"""
+ code, output, error = await subprocess_script(script, {"TEST_ROOT": str(tmp_path)})
+ assert code == 0, error
+ assert output.splitlines()[-1] == "verified"
+
+
+async def test_wispr_pipeline_rejects_actual_topology_mismatch(database, source, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ code, output, error = await subprocess_script(
+ SCRIPT,
+ {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(tmp_path / "target.json"),
+ "TEST_MODE": "wrong_topology",
+ },
+ )
+ assert code not in (0, 75, 76)
+ assert "Page source cycle must match its declared container topology" in error
+
+
+async def test_wispr_legacy_target_is_rejected_without_private_output(tmp_path):
+ script = r"""
+import sys, json, os
+from pathlib import Path
+from uuid import uuid4
+sys.path.insert(0,"tests/live")
+import conftest
+from wispr_resume import WisprResumeProbe
+path=Path(os.environ["TARGET"])
+path.write_text(json.dumps({"cycle_id":str(uuid4()),"bodies":[]}))
+try:
+ WisprResumeProbe(None,None,None,None,None,path,"resume",{})
+except ValueError as exc:
+ assert str(exc)=="Incompatible Wispr resume target; start a new private v3 trial"
+else:
+ raise AssertionError("Legacy target accepted")
+print("verified")
+"""
+ code, output, error = await subprocess_script(
+ script, {"TARGET": str(tmp_path / "old-target.json")}
+ )
+ assert code == 0, error
+ assert output.strip() == "verified"
+
+
+async def test_mixed_kind_native_id_collision_survives_real_process_resume(
+ database, source, tmp_path
+):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ target = tmp_path / "mixed-target.json"
+ env = {
+ "TEST_SCHEMA": schema,
+ "TEST_FENCE": source[1].model_dump_json(),
+ "TEST_TARGET": str(target),
+ "TEST_MIXED": "1",
+ }
+ code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "interrupt"})
+ assert code == 75, error
+ saved = json.loads(target.read_text())
+ assert saved["schema_version"] == 3 and len(saved["configuration_digest"]) == 64
+ assert len(saved["bodies"]) == 3
+ assert all("record_type" in body for body in saved["bodies"])
+ assert "private" not in output and "private" not in target.read_text()
+ code, output, error = await subprocess_script(SCRIPT, {**env, "TEST_MODE": "resume"})
+ assert code == 76, error
+ evidence = json.loads(output.splitlines()[-1])
+ assert evidence["wispr_recovery_verified"] and evidence["completed_bodies"] == 4
+ assert all(count > 0 for count in evidence["completed_body_kinds"].values())
+ assert "private" not in output
+ assert target.stat().st_mode & 0o077 == 0
diff --git a/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py b/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py
new file mode 100644
index 0000000000..906c8c85f8
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py
@@ -0,0 +1,162 @@
+"""Actual Drive pipeline publishes Docs references atomically; HTTP is synthetic."""
+
+import hashlib
+from contextlib import aclosing
+from unittest.mock import AsyncMock
+
+import httpx
+import pytest
+from sqlalchemy import select
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.domains.entities.canonical.query import CanonicalQueryService, RecordNotFound
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.helpers import capture
+from airweave.domains.entities.canonical.tests.test_drive_recovery import run, setup
+from airweave.domains.storage.file_service import FileService
+from airweave.models.entity import Entity
+from airweave.platform.sources.records.google_drive import file_record
+from airweave.platform.sources.records.google_drive_pages import DrivePages
+
+
+class NativeDocs:
+ def __init__(self):
+ self.fail = True
+
+ async def handle(self, request):
+ path = request.url.path
+ if path.endswith("/about"):
+ assert request.url.params["fields"] == "user(permissionId,me)"
+ return httpx.Response(
+ 200, json={"user": {"permissionId": "principal-a", "me": True}}
+ )
+ if request.url.host == "docs.googleapis.com":
+ if self.fail:
+ raise ConnectionError("interrupted native acquisition")
+ return httpx.Response(
+ 200,
+ json={
+ "documentId": "doc",
+ "tabs": [
+ {
+ "tabProperties": {"tabId": "one"},
+ "documentTab": {"body": {"content": []}},
+ }
+ ],
+ },
+ )
+ if path.endswith("/export"):
+ return httpx.Response(200, content=b"retained-export")
+ if path.endswith("/startPageToken"):
+ return httpx.Response(200, json={"startPageToken": "start"})
+ if path.endswith("/files"):
+ return httpx.Response(200, json={"files": [{"id": "doc"}]})
+ if path.endswith("/files/doc"):
+ return httpx.Response(
+ 200,
+ json={
+ "id": "doc",
+ "version": "7",
+ "mimeType": "application/vnd.google-apps.document",
+ },
+ )
+ if path.endswith("/changes"):
+ return httpx.Response(200, json={"changes": [], "newStartPageToken": "end"})
+ raise AssertionError("Unexpected synthetic provider endpoint")
+
+
+async def test_recapture_original_uuid_atomically_adds_manifest_and_withdraws_all_parts(
+ database, source, tmp_path, monkeypatch
+):
+ service, fence = source
+ monkeypatch.setattr(
+ "airweave.domains.storage.file_service.paths.temp_sync_dir",
+ lambda _: str(tmp_path / "temp"),
+ )
+ files = FileService(
+ fence.job_id, FilesystemBackend(tmp_path / "storage"), sync_id=fence.sync_id
+ )
+ old_blob = await files.store_canonical_blob(
+ b"retained-export",
+ media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+ )
+ old = file_record(
+ {"id": "doc", "version": "7", "mimeType": "application/vnd.google-apps.document"}
+ ).model_copy(
+ update={"blobs": (old_blob,), "content_hash": old_blob.sha256, "completeness": "complete"}
+ )
+ previous = (await capture(database, service, fence, old)).changes[0].record
+ native = NativeDocs()
+ pipeline, ctx, runtime, client = await setup(database, source, native)
+ pipeline.files = files
+ async with aclosing(client):
+ with pytest.raises(ConnectionError, match="interrupted"):
+ await run(pipeline, ctx, runtime)
+ async with database() as db:
+ unchanged = await db.get(Entity, previous.id)
+ assert unchanged.record_revision == previous.revision
+ assert len(unchanged.blob_references) == 1
+ cycle = await service.read_cycle(db, pipeline._writer())
+ assert cycle.promoted_checkpoint is None
+ native.fail = False
+ pipeline, ctx, runtime, client = await setup(database, source, native, attempt=2)
+ pipeline.files = files
+ async with aclosing(client):
+ await run(pipeline, ctx, runtime)
+ async with database() as db:
+ completed = await service.read_cycle(db, pipeline._writer())
+ # A pre-native completed Drive boundary cannot authorize a quiet delta that
+ # would leave export-only Docs permanently without their native representation.
+ previous_cycle = completed.model_copy(
+ update={
+ "configuration": completed.configuration.model_copy(
+ update={"fingerprint": hashlib.sha256(b"drive-all-accessible-v1").hexdigest()}
+ )
+ }
+ )
+ boundary = AsyncMock(return_value={"startPageToken": "new-native-baseline"})
+ plan = await DrivePages(boundary).prepare(previous_cycle, completed.configuration)
+ assert plan.mode == "full"
+ assert plan.starting_checkpoint.value == {"page_token": "new-native-baseline"}
+ boundary.assert_awaited_once()
+ query = CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-key")
+ async with database() as db:
+ current = await query.read(db, fence.organization_id, fence.sync_id, previous.id)
+ assert current.id == previous.id and current.revision == previous.revision + 1
+ assert len(current.blobs) == 3
+ assert sum(blob.role == "representation_manifest" for blob in current.blobs) == 1
+ for blob in current.blobs:
+ assert await query.blob(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ current.id,
+ current.revision,
+ blob.sha256,
+ files.storage,
+ )
+ await capture(
+ database,
+ service,
+ pipeline._writer(),
+ old.model_copy(
+ update={
+ "kind": "delete",
+ "removal_reason": "scope_removed",
+ }
+ ),
+ )
+ async with database() as db:
+ for blob in current.blobs:
+ with pytest.raises(RecordNotFound):
+ await query.blob(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ current.id,
+ current.revision,
+ blob.sha256,
+ files.storage,
+ )
+ assert len((await db.scalars(select(Entity))).all()) == 1
diff --git a/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py b/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py
new file mode 100644
index 0000000000..ea9442c099
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py
@@ -0,0 +1,113 @@
+"""Exercise the opt-in Docs sample with fake provider HTTP and real SQL/HTTP/process reads."""
+
+import json
+import os
+
+import pytest
+from sqlalchemy import text
+
+from airweave.domains.entities.canonical.tests.test_resume_probe import subprocess_script
+
+SCRIPT = r"""
+import asyncio,json,os,sys
+from contextlib import asynccontextmanager
+from pathlib import Path
+from unittest.mock import MagicMock
+sys.path.insert(0, 'tests/live')
+import canonical_capture as harness
+import workspace_document
+import httpx
+from sqlalchemy.ext.asyncio import create_async_engine,async_sessionmaker
+from provider_sample import verify_rest_identity
+from airweave.platform.sources.google_drive import GoogleDriveSource
+from airweave.platform.configs.config import GoogleDriveConfig
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.domains.storage.paths import StoragePaths
+
+document = {'documentId':'doc','title':'Fixture','tabs':[
+ {'tabProperties':{'tabId':'root'},'documentTab':{'body':{'content':[]}},
+ 'childTabs':[{'tabProperties':{'tabId':'child'},'documentTab':{'body':{'content':[]}}}]}]}
+metadata = {'id':'doc','name':'Fixture','mimeType':'application/vnd.google-apps.document',
+ 'version':'7'}
+seen=[]
+def response(request):
+ assert request.method=='GET'
+ seen.append(request.url.path)
+ if request.url.path.endswith('/about'):
+ return httpx.Response(200,json={'user':{'permissionId':'principal-a','me':True}})
+ if request.url.path=='/drive/v3/files':
+ return httpx.Response(200,json={'files':[metadata]})
+ if request.url.path.endswith('/export'):
+ return httpx.Response(200,content=b'synthetic-docx')
+ if request.url.host=='docs.googleapis.com':
+ return httpx.Response(200,json=document)
+ assert request.url.path=='/drive/v3/files/doc'
+ return httpx.Response(200,json={'version':'7'})
+
+@asynccontextmanager
+async def source(name,account,email,key,fence,**options):
+ async with httpx.AsyncClient(transport=httpx.MockTransport(response),
+ event_hooks={'request':[options['request_hook']]}) as client:
+ connector=await GoogleDriveSource.create(auth=StaticTokenProvider('synthetic'),
+ logger=MagicMock(),http_client=client,config=GoogleDriveConfig(expected_permission_id="principal-a"))
+ yield connector,'provider_email'
+
+async def main():
+ workspace_document.rest_source=source
+ os.environ['LIVE_WORKSPACE_DOCUMENT']='1'
+ root=Path(os.environ['TEST_ROOT']);root.chmod(0o700)
+ StoragePaths.TEMP_PROCESSING=str(root/'temp')
+ engine=create_async_engine(os.environ['CANONICAL_TEST_DATABASE_URL'],
+ connect_args={'server_settings':{'search_path':os.environ['TEST_SCHEMA']}})
+ try:
+ result=await harness.verify('google_drive','synthetic','synthetic@example.test',
+ 'synthetic',async_sessionmaker(engine,expire_on_commit=False),engine,root)
+ assert len(seen)==5
+ print(json.dumps(result))
+ finally:
+ await engine.dispose()
+asyncio.run(main())
+"""
+
+
+async def test_workspace_sample_survives_sql_reopen_and_real_http_reader(database, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ code, output, error = await subprocess_script(
+ SCRIPT, {"TEST_ROOT": str(tmp_path), "TEST_SCHEMA": schema}
+ )
+ assert code == 0, output + error
+ result = json.loads(output.splitlines()[-1])
+ assert result["records"] == result["journal_changes"] == 1
+ assert result["blobs"] == 3 and result["blob_sha_verified"]
+ assert result["replay_new_changes"] == 0 and result["fresh_connections_readback"]
+ assert not result["full_scope_completed"] and not result["checkpoint_saved"]
+ assert result["almanac_reader"] is None
+ assert result["document_reader"]["verified"]
+ assert result["document_reader"]["stale_revision_denied"]
+ assert result["document_reader"]["wrong_source_denied"]
+
+
+@pytest.mark.skipif(
+ not os.environ.get("ALMANAC_PROBE_ROOT"), reason="Explicit companion checkout required"
+)
+async def test_workspace_sample_through_almanac_cli(database, tmp_path):
+ async with database() as db:
+ schema = await db.scalar(text("select current_schema()"))
+ code, output, error = await subprocess_script(
+ SCRIPT,
+ {
+ "TEST_ROOT": str(tmp_path),
+ "TEST_SCHEMA": schema,
+ "LIVE_ALMANAC_ROOT": os.environ["ALMANAC_PROBE_ROOT"],
+ "LIVE_ALMANAC_PYTHON": os.environ["ALMANAC_PROBE_PYTHON"],
+ },
+ )
+ assert code == 0, output + error
+ result = json.loads(output.splitlines()[-1])
+ assert result["document_reader"]["verified"]
+ assert (
+ result["document_reader"]["scope"]
+ == "actual_cli_and_almanac_fixture_route_to_real_store_http"
+ )
+ assert result["document_reader"]["revoked_binding_denied"]
diff --git a/backend/airweave/domains/entities/canonical/tests/vespa_helpers.py b/backend/airweave/domains/entities/canonical/tests/vespa_helpers.py
new file mode 100644
index 0000000000..ad3b5a1fa2
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/tests/vespa_helpers.py
@@ -0,0 +1,47 @@
+"""Disposable real Vespa schema deployment for synthetic integration tests."""
+
+import asyncio
+import io
+import time
+import zipfile
+from pathlib import Path
+
+import httpx
+
+
+async def wait_healthy(client, url, seconds=180):
+ """Bound readiness independently of the enclosing CI job timeout."""
+ deadline = time.monotonic() + seconds
+ while time.monotonic() < deadline:
+ try:
+ response = await client.get(url)
+ if (
+ response.status_code == 200
+ and response.json().get("status", {}).get("code") == "up"
+ ):
+ return
+ except (httpx.TransportError, ValueError):
+ pass
+ await asyncio.sleep(2)
+ raise AssertionError(f"Disposable Vespa did not become ready: {url}")
+
+
+async def deploy_schema(http):
+ """Deploy only to the explicitly disposable engine used by these tests."""
+ package = Path(__file__).resolve().parents[6] / "vespa/app"
+ archive = io.BytesIO()
+ with zipfile.ZipFile(archive, "w") as output:
+ for path in sorted(package.rglob("*")):
+ if path.is_file():
+ content = path.read_text().replace("{{EMBEDDING_DIM}}", "384")
+ output.writestr(
+ str(path.relative_to(package)), content.replace("{{VERSION}}", "test")
+ )
+ await wait_healthy(http, "http://localhost:19071/state/v1/health")
+ deployed = await http.post(
+ "http://localhost:19071/application/v2/tenant/default/prepareandactivate",
+ content=archive.getvalue(),
+ headers={"Content-Type": "application/zip"},
+ )
+ assert deployed.status_code == 200, deployed.text
+ await wait_healthy(http, "http://localhost:8081/state/v1/health")
diff --git a/backend/airweave/domains/entities/canonical/text_artifacts.py b/backend/airweave/domains/entities/canonical/text_artifacts.py
new file mode 100644
index 0000000000..0d4055f0e8
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/text_artifacts.py
@@ -0,0 +1,53 @@
+"""Build immutable text bytes once, then publish alongside the index generation."""
+
+import hashlib
+from uuid import UUID
+
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.text_models import TextArtifact, representation_id
+from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES
+from airweave.domains.sync_pipeline.pipeline.text_models import BuiltText
+
+
+def prepare_text(
+ built: tuple[BuiltText, ...],
+ generation: UUID,
+) -> tuple[tuple[TextArtifact, bytes], ...]:
+ """Keep complete converter output; oversized text fails rather than truncating."""
+ artifacts = []
+ for item in built:
+ locator = ProjectionLocator.parse(item.entity_id)
+ if locator is None or locator.generation != generation:
+ raise ValueError("Built text lost its publication identity")
+ content = item.text.encode("utf-8")
+ if len(content) > MAX_FILE_SIZE_BYTES:
+ raise ValueError("Derived text exceeds retained representation size limit")
+ artifact = TextArtifact(
+ id=representation_id(generation, locator.part_index),
+ part_index=locator.part_index,
+ sha256=hashlib.sha256(content).hexdigest(),
+ size_bytes=len(content),
+ characters=len(item.text),
+ content_start=item.content_start,
+ kind=item.kind,
+ )
+ artifacts.append((artifact, content))
+ return tuple(artifacts)
+
+
+def text_manifest(
+ artifacts: tuple[TextArtifact, ...] | None,
+ indexed_parts: set[int],
+ sync_id: UUID,
+ generation: UUID,
+) -> list[dict] | None:
+ """Validate exact part membership and generation-owned keys before committing descriptors."""
+ if artifacts is None:
+ return None
+ if len({item.id for item in artifacts}) != len(artifacts):
+ raise ValueError("Text representations must be unique")
+ if {item.part_index for item in artifacts} != indexed_parts:
+ raise ValueError("Text representations must cover exactly indexed parts")
+ for item in artifacts:
+ item.storage_key(sync_id, generation)
+ return [item.model_dump(mode="json") for item in artifacts]
diff --git a/backend/airweave/domains/entities/canonical/text_models.py b/backend/airweave/domains/entities/canonical/text_models.py
new file mode 100644
index 0000000000..2a9afa2617
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/text_models.py
@@ -0,0 +1,91 @@
+"""Derived text descriptors; source originals remain the canonical authority."""
+
+from typing import Literal
+from uuid import UUID, uuid5
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+class TextArtifact(BaseModel):
+ """Body-free immutable descriptor, owned by one projection generation."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ id: UUID
+ part_index: int = Field(ge=0)
+ sha256: str = Field(pattern=r"^[a-f0-9]{64}$")
+ size_bytes: int = Field(ge=0)
+ characters: int = Field(ge=0)
+ content_start: int | None = Field(default=None, ge=0)
+ kind: Literal["native_text", "extracted_text", "generated_text"]
+
+ @model_validator(mode="before")
+ @classmethod
+ def legacy_kind(cls, value):
+ """Older manifests encoded only converter-vs-generated provenance."""
+ if isinstance(value, dict) and "kind" not in value:
+ return {
+ **value,
+ "kind": (
+ "extracted_text" if value.get("content_start") is not None else "generated_text"
+ ),
+ }
+ return value
+
+ @model_validator(mode="after")
+ def boundary(self):
+ """The content boundary must lie inside the retained indexing text."""
+ if (self.kind == "generated_text") != (self.content_start is None):
+ raise ValueError("Text provenance disagrees with content boundary")
+ if self.content_start is not None and self.content_start > self.characters:
+ raise ValueError("Content boundary exceeds representation length")
+ return self
+
+ def storage_key(self, sync_id: UUID, generation: UUID) -> str:
+ """No caller/provider-supplied storage path is accepted."""
+ if self.id != representation_id(generation, self.part_index):
+ raise ValueError("Representation identity does not belong to this generation")
+ return f"canonical/{sync_id}/projections/{generation}/{self.id}/{self.sha256}"
+
+
+def representation_id(generation: UUID, part_index: int) -> UUID:
+ """Opaque stable reference within this immutable publication generation."""
+ return uuid5(generation, f"text:{part_index}")
+
+
+class TextRepresentation(BaseModel):
+ """Public reference with explicit derived provenance, never a storage key."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ id: UUID
+ record_id: UUID
+ revision: int
+ generation: UUID
+ pipeline_version: int
+ part_key: str
+ kind: Literal["native_text", "extracted_text", "generated_text"]
+ media_type: Literal["text/markdown"] = "text/markdown"
+ content_characters: int | None
+ index_characters: int
+ source_anchors: Literal["unavailable"] = "unavailable"
+
+
+class TextRepresentationList(BaseModel):
+ """Unknown covers generations built before text retention, not an empty document."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ record_id: UUID
+ revision: int
+ status: Literal["available", "unavailable"]
+ representations: tuple[TextRepresentation, ...]
+
+
+class TextRead(BaseModel):
+ """Unicode character range in the selected derived representation view."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ representation: TextRepresentation
+ view: Literal["content", "index"]
+ offset: int
+ total_characters: int
+ text: str
+ next_offset: int | None
diff --git a/backend/airweave/domains/entities/canonical/text_query.py b/backend/airweave/domains/entities/canonical/text_query.py
new file mode 100644
index 0000000000..76d406ce2e
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/text_query.py
@@ -0,0 +1,225 @@
+"""Current-publication text reads, with the same visibility fence as search."""
+
+import hashlib
+from typing import Literal
+from uuid import UUID
+
+from pydantic import ValidationError
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import publication_matches
+from airweave.domains.entities.canonical.query import (
+ BlobUnavailable,
+ CanonicalQueryService,
+ RecordNotFound,
+ StaleRecordRevision,
+)
+from airweave.domains.entities.canonical.store import CanonicalStoreError
+from airweave.domains.entities.canonical.text_models import (
+ TextArtifact,
+ TextRead,
+ TextRepresentation,
+ TextRepresentationList,
+)
+from airweave.domains.storage.exceptions import StorageException
+from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES
+from airweave.domains.storage.protocols import StorageBackend
+from airweave.models.entity import Entity
+from airweave.models.projection_generation import ProjectionGeneration
+from airweave.models.sync import Sync
+
+
+class TextUnavailable(CanonicalStoreError):
+ """No retained text at this exact current publication; refresh/reproject explicitly."""
+
+ code = "text_unavailable"
+
+
+class CanonicalTextReader:
+ """Derived text has no independent authority or provider access."""
+
+ def __init__(self, records: CanonicalQueryService, storage: StorageBackend):
+ """Reuse record authorization and the configured blob backend."""
+ self.records, self.storage = records, storage
+
+ async def _snapshot(
+ self,
+ db: AsyncSession,
+ organization: UUID,
+ sync_id: UUID,
+ record_id: UUID,
+ revision: int,
+ ) -> tuple[ProjectionGeneration | None, tuple[TextArtifact, ...]]:
+ record = await self.records.read(db, organization, sync_id, record_id)
+ if record.deleted_at is not None or record.content_access != "available":
+ raise RecordNotFound("Record content is not currently available")
+ if record.revision != revision:
+ raise StaleRecordRevision("Record changed; reread before requesting text")
+ row = await db.scalar(
+ select(ProjectionGeneration)
+ .join(
+ Entity,
+ Entity.indexed_generation == ProjectionGeneration.id,
+ )
+ .where(
+ Entity.id == record_id,
+ Entity.sync_id == sync_id,
+ Entity.organization_id == organization,
+ ProjectionGeneration.organization_id == organization,
+ ProjectionGeneration.sync_id == sync_id,
+ ProjectionGeneration.record_id == record_id,
+ ProjectionGeneration.revision == revision,
+ )
+ .execution_options(populate_existing=True)
+ )
+ if row is None or row.text_representations is None:
+ return None, ()
+ try:
+ artifacts = tuple(
+ TextArtifact.model_validate(item) for item in row.text_representations
+ )
+ coverage = ExtractionCoverage.model_validate(row.extraction_coverage)
+ if len({item.id for item in artifacts}) != len(artifacts):
+ raise ValueError("Duplicate artifact")
+ if {item.part_index for item in artifacts} != {
+ part.part_index for part in coverage.parts if part.outcome == "indexed"
+ }:
+ raise ValueError("Artifact coverage mismatch")
+ for artifact in artifacts:
+ artifact.storage_key(sync_id, row.id)
+ except (ValidationError, ValueError):
+ raise TextUnavailable("Retained text metadata is unavailable") from None
+ if artifacts:
+ # Exact coverage membership above covers every descriptor. The remaining
+ # publication/ancestor predicate is common to the whole generation.
+ artifact = artifacts[0]
+ locator = ProjectionLocator(
+ record_id=record_id,
+ revision=revision,
+ pipeline_version=row.pipeline_version,
+ generation=row.id,
+ part_index=artifact.part_index,
+ )
+ visible = await db.scalar(
+ select(Entity.id)
+ .join(Sync)
+ .where(
+ Entity.organization_id == organization,
+ Entity.sync_id == sync_id,
+ Sync.organization_id == organization,
+ publication_matches(locator),
+ )
+ )
+ if visible is None:
+ return None, ()
+ return row, artifacts
+
+ @staticmethod
+ def _describe(row: ProjectionGeneration, artifact: TextArtifact) -> TextRepresentation:
+ coverage = ExtractionCoverage.model_validate(row.extraction_coverage)
+ part = next(part for part in coverage.parts if part.part_index == artifact.part_index)
+ return TextRepresentation(
+ id=artifact.id,
+ record_id=row.record_id,
+ revision=row.revision,
+ generation=row.id,
+ pipeline_version=row.pipeline_version,
+ part_key=part.key,
+ kind=artifact.kind,
+ content_characters=(
+ artifact.characters - artifact.content_start
+ if artifact.content_start is not None
+ else None
+ ),
+ index_characters=artifact.characters,
+ )
+
+ async def list(
+ self,
+ db: AsyncSession,
+ organization: UUID,
+ sync_id: UUID,
+ record_id: UUID,
+ revision: int,
+ ) -> TextRepresentationList:
+ """Describe current retained text without reading blobs or claiming legacy coverage."""
+ row, artifacts = await self._snapshot(db, organization, sync_id, record_id, revision)
+ return TextRepresentationList(
+ record_id=record_id,
+ revision=revision,
+ status="available" if row is not None and artifacts else "unavailable",
+ representations=tuple(self._describe(row, item) for item in artifacts),
+ )
+
+ async def read(
+ self,
+ db: AsyncSession,
+ organization: UUID,
+ sync_id: UUID,
+ record_id: UUID,
+ revision: int,
+ generation: UUID,
+ representation_id: UUID,
+ *,
+ offset: int = 0,
+ limit: int = 12000,
+ view: Literal["content", "index"] = "content",
+ ) -> TextRead:
+ """Default content view excludes metadata; generated-only text requires index view."""
+ if offset < 0 or not 1 <= limit <= 100000:
+ raise TextUnavailable("Text range must have nonnegative offset and limit 1–100000")
+ row, artifacts = await self._snapshot(db, organization, sync_id, record_id, revision)
+ artifact = next((item for item in artifacts if item.id == representation_id), None)
+ if row is None or row.id != generation or artifact is None:
+ raise TextUnavailable("Text publication changed or is not available; reread the record")
+ if view == "content" and artifact.content_start is None:
+ raise TextUnavailable("This is generated search text; select its index view explicitly")
+ description = self._describe(row, artifact)
+ text = await self._read_verified(artifact, sync_id, generation)
+ db.expire_all()
+ current, current_artifacts = await self._snapshot(
+ db,
+ organization,
+ sync_id,
+ record_id,
+ revision,
+ )
+ if current is None or current.id != generation or artifact not in current_artifacts:
+ raise TextUnavailable("Text publication changed during read; reread the record")
+ if view == "content":
+ text = text[artifact.content_start :]
+ if offset > len(text):
+ raise TextUnavailable("Text offset exceeds this representation")
+ end = min(len(text), offset + limit)
+ return TextRead(
+ representation=description,
+ view=view,
+ offset=offset,
+ total_characters=len(text),
+ text=text[offset:end],
+ next_offset=end if end < len(text) else None,
+ )
+
+ async def _read_verified(self, artifact: TextArtifact, sync_id: UUID, generation: UUID) -> str:
+ """Bound storage reads and verify UTF-8, digest and both recorded lengths."""
+ if artifact.size_bytes > MAX_FILE_SIZE_BYTES:
+ raise BlobUnavailable("Retained text exceeds its storage bound")
+ try:
+ raw = await self.storage.read_file(
+ artifact.storage_key(sync_id, generation),
+ max_bytes=artifact.size_bytes,
+ )
+ if (
+ len(raw) != artifact.size_bytes
+ or hashlib.sha256(raw).hexdigest() != artifact.sha256
+ ):
+ raise ValueError("Text digest mismatch")
+ text = raw.decode("utf-8")
+ if len(text) != artifact.characters:
+ raise ValueError("Text character count mismatch")
+ except (StorageException, ValueError):
+ raise BlobUnavailable("Retained text bytes are unavailable; retry later") from None
+ return text
diff --git a/backend/airweave/domains/entities/canonical/workspace_docs.py b/backend/airweave/domains/entities/canonical/workspace_docs.py
new file mode 100644
index 0000000000..b6e62b1b1e
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/workspace_docs.py
@@ -0,0 +1,165 @@
+"""Read retained native Docs content; no provider access or independent identity."""
+
+from __future__ import annotations
+
+import json
+
+from pydantic import BaseModel, ConfigDict, JsonValue
+
+from airweave.domains.entities.canonical.blob_materializer import read_blob
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.storage.protocols import StorageBackend
+from airweave.platform.sources.records.workspace_manifest import (
+ WorkspaceManifestV1,
+ parse_manifest,
+ validate_document,
+)
+
+
+class DocumentReadError(ValueError):
+ """The retained document cannot supply the requested representation."""
+
+
+class CapturedDocument(BaseModel):
+ """Original native structure with validated capture provenance, never reconstructed JSON."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ manifest: WorkspaceManifestV1
+ document: dict[str, JsonValue]
+
+ def tab(self, tab_id: str) -> dict[str, JsonValue]:
+ """Return the unchanged native subtree for one exact captured tab ID."""
+ for tab in _tabs(self.document["tabs"]):
+ properties = _object(tab.get("tabProperties"))
+ if properties.get("tabId") == tab_id:
+ return tab
+ raise DocumentReadError("Requested tab is not present in this captured document")
+
+ def text(self) -> str:
+ """Project supported authored text across all tabs and ancillary content."""
+ if self.manifest.native.status != "complete":
+ raise DocumentReadError("Native document text capture is incomplete")
+ sections = []
+ for tab in _tabs(self.document["tabs"]):
+ properties = _object(tab.get("tabProperties"))
+ title = properties.get("title")
+ if isinstance(title, str) and title:
+ sections.append(title + "\n")
+ document = _object(tab.get("documentTab"))
+ sections.append(_content(_object(document.get("body")).get("content", [])))
+ for field in ("headers", "footers", "footnotes"):
+ for value in _object(document.get(field, {})).values():
+ sections.append(_content(_object(value).get("content", [])))
+ return "\n".join(part for part in sections if part)
+
+
+def _object(value: JsonValue) -> dict[str, JsonValue]:
+ if not isinstance(value, dict):
+ raise DocumentReadError("Captured Docs structure requires an object")
+ return value
+
+
+def _array(value: JsonValue) -> list[JsonValue]:
+ if not isinstance(value, list):
+ raise DocumentReadError("Captured Docs structure requires an array")
+ return value
+
+
+def _tabs(value: JsonValue):
+ for item in _array(value):
+ tab = _object(item)
+ yield tab
+ yield from _tabs(tab.get("childTabs", []))
+
+
+def _content(value: JsonValue) -> str:
+ sections = []
+ for item in _array(value):
+ element = _object(item)
+ if "paragraph" in element:
+ sections.append(_paragraph(_object(element["paragraph"])))
+ elif "table" in element:
+ table = _object(element["table"])
+ for row in _array(table.get("tableRows", [])):
+ cells = _array(_object(row).get("tableCells", []))
+ sections.append("\t".join(_content(_object(c).get("content", [])) for c in cells))
+ elif "tableOfContents" in element:
+ sections.append(_content(_object(element["tableOfContents"]).get("content", [])))
+ elif "sectionBreak" not in element:
+ raise DocumentReadError("Unsupported captured Docs structural element")
+ return "\n".join(sections)
+
+
+def _paragraph(paragraph: dict[str, JsonValue]) -> str:
+ pieces = []
+ for item in _array(paragraph.get("elements", [])):
+ element = _object(item)
+ if "textRun" in element:
+ text = _object(element["textRun"]).get("content")
+ if not isinstance(text, str):
+ raise DocumentReadError("Captured Docs text run has no text")
+ pieces.append(text)
+ elif "person" in element:
+ properties = _object(_object(element["person"]).get("personProperties"))
+ pieces.append(_display_text(properties.get("name") or properties.get("email")))
+ elif "richLink" in element:
+ properties = _object(_object(element["richLink"]).get("richLinkProperties"))
+ pieces.append(_display_text(properties.get("title")))
+ elif "equation" in element:
+ raise DocumentReadError("Google Docs does not expose equation text for projection")
+ elif any(
+ key in element
+ for key in (
+ "inlineObjectElement",
+ "pageBreak",
+ "columnBreak",
+ "footnoteReference",
+ "horizontalRule",
+ "autoText",
+ )
+ ):
+ # Object bytes and automatic page numbers are not authored paragraph text.
+ # Footnote text is included separately from its documentTab collection.
+ continue
+ else:
+ raise DocumentReadError("Unsupported captured Docs paragraph element")
+ return "".join(pieces)
+
+
+def _display_text(value: JsonValue) -> str:
+ if not isinstance(value, str) or not value:
+ raise DocumentReadError("Captured Docs smart chip has no display text")
+ return value
+
+
+async def read_document(record: SourceRecord, storage: StorageBackend) -> CapturedDocument:
+ """Resolve only the current record's committed manifest and document bytes."""
+ if record.content_access != "available" or record.deleted_at is not None:
+ raise DocumentReadError("Captured document is unavailable")
+ if record.identity.record_type != "file" or record.payload.get("mimeType") != (
+ "application/vnd.google-apps.document"
+ ):
+ raise DocumentReadError("Captured file is not a Google document")
+ manifests = [blob for blob in record.blobs if blob.role == "representation_manifest"]
+ if len(manifests) != 1:
+ raise DocumentReadError("Captured document requires one representation manifest")
+ version = record.payload.get("version")
+ if not isinstance(version, str):
+ raise DocumentReadError("Captured document has no Drive version")
+ manifest = parse_manifest(
+ await read_blob(record, manifests[0], storage),
+ file_id=record.identity.native_id,
+ drive_version=version,
+ blobs=record.blobs,
+ )
+ digest = manifest.native.document_blob
+ if digest is None:
+ raise DocumentReadError("Native document content was not captured")
+ matches = [blob for blob in record.blobs if blob.sha256 == digest]
+ if len(matches) != 1:
+ raise DocumentReadError("Native document blob reference is ambiguous or absent")
+ document = json.loads(await read_blob(record, matches[0], storage))
+ if not isinstance(document, dict):
+ raise DocumentReadError("Native document response is not an object")
+ validate_document(document, file_id=record.identity.native_id)
+ return CapturedDocument(manifest=manifest, document=document)
diff --git a/backend/airweave/domains/entities/canonical/workspace_sheets.py b/backend/airweave/domains/entities/canonical/workspace_sheets.py
new file mode 100644
index 0000000000..8301a57f46
--- /dev/null
+++ b/backend/airweave/domains/entities/canonical/workspace_sheets.py
@@ -0,0 +1,147 @@
+"""Validated native spreadsheet reads from committed Drive-owned representations."""
+
+import json
+
+from pydantic import BaseModel, ConfigDict, JsonValue
+
+from airweave.domains.entities.canonical.blob_materializer import read_blob
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.storage.protocols import StorageBackend
+from airweave.platform.sources.records.sheets_manifest import (
+ SHEETS_MIME,
+ GridBounds,
+ GridPart,
+ WorkspaceManifestV2,
+ parse_sheet_manifest,
+)
+from airweave.platform.sources.records.sheets_models import (
+ SheetGrid,
+ SpreadsheetCell,
+ a1_range,
+ column_name,
+ metadata_shape,
+ validate_grid,
+ validate_partition,
+)
+
+
+class CapturedGrid(BaseModel):
+ """One requested native response paired with its admitted bounds."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ part: GridPart
+ data: dict[str, JsonValue]
+
+
+class CapturedSpreadsheet(BaseModel):
+ """One exact stored file revision with validated structured coverage."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ manifest: WorkspaceManifestV2
+ spreadsheet: dict[str, JsonValue]
+ grids: tuple[CapturedGrid, ...]
+
+ def cells(self, bounds: GridBounds) -> tuple[SpreadsheetCell, ...]:
+ """Return only retained cells in the selected rectangle, with coverage separately."""
+ sheets = metadata_shape(self.spreadsheet, self.manifest.file_id).sheets
+ selected = next((s for s in sheets if s.properties.sheetId == bounds.sheet_id), None)
+ if selected is None or selected.properties.gridProperties is None:
+ raise ValueError("Requested sheet is absent or non-grid")
+ dims = selected.properties.gridProperties
+ if bounds.end_row > dims.rowCount or bounds.end_column > dims.columnCount:
+ raise ValueError("Requested range exceeds sheet dimensions")
+ result = []
+ for grid in self.grids:
+ if grid.part.bounds.sheet_id == bounds.sheet_id:
+ result.extend(_selected_cells(grid.data, bounds))
+ return tuple(result)
+
+ def text(self) -> str:
+ """One searchable representation; native formula expressions remain in originals."""
+ if self.manifest.native.status != "complete" or len(self.grids) != len(
+ self.manifest.native.parts
+ ):
+ raise ValueError("Partial native grids cannot replace the full export projection")
+ lines = []
+ for sheet in metadata_shape(self.spreadsheet, self.manifest.file_id).sheets:
+ title = sheet.properties.title
+ lines.append(title)
+ for grid in self.grids:
+ if grid.part.bounds.sheet_id == sheet.properties.sheetId:
+ lines.append(a1_range(title, grid.part.bounds))
+ lines.extend(_grid_text(grid.data))
+ return "\n".join(lines)
+
+
+def _grid_text(data: dict[str, JsonValue]):
+ for block in SheetGrid.model_validate(data["sheets"][0]).data:
+ for offset, row in enumerate(block.rowData):
+ values = []
+ for cell in row.values:
+ value = cell.get("formattedValue", "")
+ note = cell.get("note", "")
+ if not isinstance(value, str) or not isinstance(note, str):
+ raise ValueError("Malformed native spreadsheet display text")
+ values.append(value + (" [Note: " + note + "]" if note else ""))
+ address = f"{column_name(block.startColumn)}{block.startRow + offset + 1}"
+ yield address + ": " + "\t".join(values)
+
+
+def _selected_cells(grid: dict[str, JsonValue], bounds: GridBounds):
+ for block in SheetGrid.model_validate(grid["sheets"][0]).data:
+ for offset, row in enumerate(block.rowData):
+ row_index = block.startRow + offset
+ if not bounds.start_row <= row_index < bounds.end_row:
+ continue
+ for offset, cell in enumerate(row.values):
+ column = block.startColumn + offset
+ if bounds.start_column <= column < bounds.end_column:
+ yield SpreadsheetCell(row=row_index, column=column, value=cell)
+
+
+def intersects(first: GridBounds, second: GridBounds) -> bool:
+ """Select retained rectangles without expanding sparse grids into dense cell arrays."""
+ return (
+ first.sheet_id == second.sheet_id
+ and max(first.start_row, second.start_row) < min(first.end_row, second.end_row)
+ and max(first.start_column, second.start_column) < min(first.end_column, second.end_column)
+ )
+
+
+async def read_spreadsheet(
+ record: SourceRecord,
+ storage: StorageBackend,
+ *,
+ bounds: GridBounds | None = None,
+ for_projection: bool = False,
+) -> CapturedSpreadsheet:
+ """Resolve only admitted blobs and validate all native range identities."""
+ if record.content_access != "available" or record.deleted_at is not None:
+ raise ValueError("Captured spreadsheet is unavailable")
+ if record.identity.record_type != "file" or record.payload.get("mimeType") != SHEETS_MIME:
+ raise ValueError("Captured file is not a spreadsheet")
+ manifests = [blob for blob in record.blobs if blob.role == "representation_manifest"]
+ if len(manifests) != 1 or not isinstance(record.payload.get("version"), str):
+ raise ValueError("Native spreadsheet has not been captured")
+ manifest = parse_sheet_manifest(
+ await read_blob(record, manifests[0], storage),
+ file_id=record.identity.native_id,
+ drive_version=record.payload["version"],
+ blobs=record.blobs,
+ )
+ by_digest = {blob.sha256: blob for blob in record.blobs}
+ metadata = json.loads(
+ await read_blob(record, by_digest[manifest.native.metadata_blob], storage)
+ )
+ validate_partition(manifest, metadata)
+ grids = []
+ for part in manifest.native.parts:
+ if not (
+ (for_projection and manifest.native.status == "complete")
+ or (bounds is not None and intersects(part.bounds, bounds))
+ ):
+ continue
+ grid = json.loads(await read_blob(record, by_digest[part.blob], storage))
+ validate_grid(grid, manifest.file_id, part.bounds)
+ grids.append(CapturedGrid(part=part, data=grid))
+ return CapturedSpreadsheet(manifest=manifest, spreadsheet=metadata, grids=tuple(grids))
diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md
new file mode 100644
index 0000000000..caf1db3b7a
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/README.md
@@ -0,0 +1,229 @@
+# Native snapshot admission
+
+This service retains versioned Almanac snapshots in the existing canonical
+store. Backend-key HTTP endpoints own source provisioning and the import
+lifecycle; they do not enumerate Almanac data. Almanac remains the authority
+for content and access.
+
+The source API binds `short_name=almanac` and
+`config_fields={owner_id, dataset}`. Datasets are `knowledge` and `sessions`.
+The service checks that binding against the authenticated organization, stored
+writer fence and every snapshot. It does not require a fabricated Composio
+account. Internal commands are not authorization credentials; HTTP callers
+cannot supply writer fences.
+
+`NativeIngestionService` owns one existing `UnitOfWork`.
+`NativeIngestionStore` composes `CanonicalRecordStore._fenced_sync` and
+`_capture_locked` inside that transaction. Those existing internal methods are
+the only locking/capture dependencies. No table, independent journal, or
+alternate revision allocator is introduced.
+
+Admission loads the batch's exact record identities and parent attestations in
+one query under that writer lock. Canonical capture still owns its per-record
+writes; batching admission does not bypass its revision or visibility checks.
+
+Retained payloads explicitly identify their authority and snapshot provenance,
+contain the unchanged original JSON, and retain the authoritative version even
+for tombstones. Knowledge versions are scalar record revisions. Session and
+message snapshots carry metadata and original-content revisions as a partial
+order; crossed components are rejected. Replay epochs are not transcript
+versions. Changed messages, including tombstones, must match their attested session
+version, including a parent submitted in the same batch. A deleted parent permits
+matching message tombstones but never message upserts. Submit parents before children.
+
+An equal version is accepted only for an identical snapshot; retries do not
+restore withdrawn visibility. Lower versions, conflicting equal versions and
+malformed retained state fail closed. Local scope withdrawal already preserves
+the canonical payload, so it also preserves native version state. Higher versions
+cannot reopen locally withdrawn records; use the explicit access observation
+operation below. Parent access gates remain canonical-owned.
+
+This is record admission, not proof of a complete session transcript or complete
+dataset. Pagination, final snapshot completion, deletion discovery, authoritative
+export validation, and source renewal remain work for the publisher boundary.
+An empty original JSON tombstone is valid but must carry an authoritative version.
+The publisher must not invent a version from capture time or an observation count.
+
+Tests use the existing disposable PostgreSQL fixture, including actual migration
+schemas and transaction rollback. They do not contact a provider or production.
+
+## Atomic snapshot pages
+
+`IngestNativePage` composes native admission with the existing canonical scan
+transaction. `CanonicalScanStore.page` accepts an internal `LockedPageCapture`
+implementation; it is code-owned behavior, never a request-supplied callback.
+The scan still owns scope validation, CAS, continuation and final-page state.
+Native `admit_locked` returns typed changed observations and unchanged record IDs.
+An unchanged record in a new sweep must pass current parent-version and active
+visibility checks before its sighting is stamped. This changes neither its
+canonical revision nor its payload and never restores withdrawn content.
+
+Ordinary exact `ingest` retries still acknowledge retained data without claiming
+fresh membership. Snapshot pages enumerate active upserts only; authoritative
+tombstones continue through ordinary ingestion. This page operation is not a
+changes-feed contract. Empty pages are valid; only `final=true` begins the existing
+reconciliation phase, and source completion remains a separate cycle barrier.
+Admission, capture, sightings and continuation either commit together or roll back.
+
+HTTP authentication, receipts, provisioning and completion/cancellation are
+implemented below, including explicit record access renewal. Authoritative
+publisher completeness and automatic source revalidation remain pending. Native jobs cannot be cancelled through the generic
+provider route: the guard checks the actual job's sync identity.
+
+## Native search projection
+
+`projection.py` supplies offline inputs to the existing canonical projector.
+It validates original IDs, revisions and timestamps against their snapshot,
+then selects native text without replacing the retained original JSON:
+
+- Knowledge: body, title, description and user notes. Path is a locator; renaming
+ it keeps record identity. Selected structured fields also enter index metadata: person contacts, roles,
+ education and relationship labels; organisation domains/industry; place address
+ and coordinates; event schedule; creative-work kind/date; profile links.
+ Reference IDs are retained but never resolved into invented searchable names.
+ Citation metadata and arbitrary extra JSON remain outside the text projection.
+- Sessions: title and description. Archived sessions remain included, matching
+ current Almanac session search; archived knowledge is excluded.
+- Messages: original string content or `type=text` blocks, with session/message
+ IDs, ordinals and block anchors. Inactive legacy imports are excluded. API
+ replay payloads, reasoning fields and arbitrary metadata are not search text.
+
+Mixed messages can publish their text while retained nontext blocks report
+`unsupported_format` and partial extraction coverage. Malformed text fails
+explicitly. Retaining an image URL in original JSON does not capture image bytes
+or make them downloadable; attachment acquisition is still separate work.
+
+Native indexed sources use the same publication and access validation as provider
+sources, without a provider registry entry or fabricated OAuth connection. This
+is an index capability, not an authentication grant. The source API binds the
+owner; a trusted authoritative publisher is still required for real ingestion.
+
+## Trusted source and import identity
+
+Backend-key `PUT /native/sources` ensures the server-derived organization/owner/
+dataset identity in one collection; `GET /native/sources/{source_id}` reads it.
+Concurrent requests serialize through the existing Sync primary key. Collection
+changes and unexplained partial state conflict; an exact ensure never reopens
+withdrawn availability. This setup is not proof of imported or indexed content.
+
+`NativeImportStore` owns the start/read transaction behind the import HTTP API. A bounded request key derives an existing SyncJob primary
+key. Exact retries recover the original state, even after cancellation or source
+withdrawal; they never activate another writer. A changed body conflicts. New
+imports require an available source and cannot replace another active writer.
+Job creation, canonical writer activation and cycle creation commit together.
+A new key after a terminal job may replace its incomplete cycle only when the
+retained cycle is associated with that preceding native import. Retained records
+remain until the subsequent declared scans reconcile them.
+
+Publisher-declared bounded snapshots use discovery-only completion policies;
+complete dataset snapshots use exhaustive policies. This declaration does not
+itself certify completeness. The HTTP lifecycle verifies final scopes and
+receipts; the publisher must establish stable export boundaries. Explicit record
+access renewal is available under an active import as described below; source-wide
+credential renewal and automatic source revalidation remain separate.
+
+## Import HTTP and page receipts
+
+Backend-key PUT/GET `/native/sources/{source_id}/imports/{request_key}` expose
+start and recovery. Keys are opaque nonsecret identifiers: application logs and
+analytics redact native paths/parameters, but infrastructure access logs may
+still record request URLs. Validation and unexpected exceptions do not log
+native request contents through the application middleware.
+
+`NativePageStore.commit` is the transaction behind PUT `/pages`.
+It resolves stored writer authority, admits
+the page, advances its cursor and retains one bounded receipt in the existing
+scan continuation. No second journal or original-content copy is created. Exact
+retries of the current page return its acknowledgement. Changed bodies, old CAS
+versions, changed reconciliation state and cancelled writers conflict rather
+than reapply writes. Cursors are bounded to32KiB, leaving receipt space inside
+the existing64KiB continuation limit. Capture acknowledgement is distinct from
+index publication or whole-import completion.
+
+## Scope HTTP workflow
+
+Under an active import, PUT `/scopes` begins or resumes an exact scope;
+POST `/scopes/read` reads its composite identity from the body without mutation;
+PUT `/pages` commits a page; POST `/scopes/reconcile` advances bounded
+reconciliation after the final page. Responses use `bounded`/`complete` coverage
+and an independent phase, hide writer credentials and internal receipt digests,
+and retain the caller cursor plus last acknowledgement for recovery. Parent
+epochs are resolved under the server's writer lock. A new import can replace a
+prior-cycle scan using its locked version; same-cycle restarts require explicit CAS.
+
+Bounded native imports use observed membership: only parents seen in their
+current completed inventory can admit child scopes or require transcripts for
+completion. Unseen retained sessions remain available but are not fresh members.
+Default provider membership and checkpoint digests are unchanged. Bounded scope
+completion never becomes evidence of exhaustive account discovery. Scope reads
+currently require an active import; durable terminal summaries and historical
+scope availability are still part of import finalization work.
+
+### Native import terminal outcomes
+
+The backend-only import API now supports `POST .../imports/{request_key}/complete`
+and `/cancel`. Completion validates all scopes required by the declared capture
+policy under the writer lock, then stores the cycle completion, job status, and a
+bounded terminal summary in one transaction. A bounded import completing does not
+mean the entire account was captured. Summary `indexing: not_verified` explicitly
+separates capture from downstream search publication.
+
+An identical completed retry returns its retained summary even after another import
+starts. Cancelling an already-terminal import preserves its outcome and never
+cancels the newer writer. A cancelled job immediately fails writer-fence checks.
+Terminal scope reads remain available only while the reusable scan row still has
+that import's cycle ID; once reused, the API reports unavailable/superseded instead
+of showing another import's progress. The import's terminal summary remains durable.
+Request keys are opaque, nonsecret identifiers: application diagnostics redact them,
+but deployment access logs require their own policy.
+
+
+## Explicit record access observations
+
+Backend-only GET/POST
+`/native/sources/{source_id}/imports/{request_key}/records/{record_id}/access`
+read or change current access. GET returns identity, retained revision, current parent
+visibility epoch, availability and removal reason, never original content. It works for a terminal import, but is
+a current-state read, not historical state at that import. Availability includes
+source authentication, active sync, record deletion and ancestor visibility.
+
+POST requires an active server-held writer and `expected_revision` from the retained
+record. Renewal of a message also requires `expected_parent_epoch` from the access
+read. Parent withdrawal/restoration invalidates old child evidence even if the child
+revision did not change. Roots use a null parent epoch. `action: withdraw` takes a
+confirmed `access_revoked` or `scope_removed` reason. It preserves original native payload/version and blob references while
+journaling withdrawal through canonical capture. A search miss or incomplete scan
+is not valid evidence. An ambiguous source404 requires provider/source interpretation.
+
+`action: renew` supplies a freshly read, owner-attested active NativeSnapshot for
+the exact record. Native owner, identity, version ordering and parent-version checks
+still apply. Equal native version requires identical content; it can now explicitly
+renew access at the expected local revision. Ordinary imports/page retries cannot.
+Newer ordinary child updates also cannot restore inherited unavailability; they require the same explicit renewal path.
+Restore parents first, then freshly attest children; restoration advances parent
+visibility so previous child attestations remain hidden until individually renewed.
+Source-wide authentication is not changed by this operation.
+
+Read retained revision *before* obtaining fresh source evidence. If another mutation
+wins meanwhile, CAS fails. After a lost response, GET current state; do not blindly
+change the expected revision and resend old evidence. This operation uses ordinary
+CAS recovery rather than a second idempotency journal. An already-active unchanged
+renewal may be a no-op. Access observations do not mark records seen in inventory;
+normal page capture still owns scan progress/completeness.
+
+No new table, migration, scheduler or source credential exists here. NativeImports
+owns the transaction; NativeAccessStore composes native admission and shared
+canonical revision/journal/visibility machinery. The trusted publisher remains
+responsible for obtaining actual source evidence and scheduling revalidation.
+
+
+### Structured-field projection rollout
+
+Changes to this mapper do not recapture or revise originals. Existing unchanged
+records need a scoped pipeline-version increment using the existing
+`backend/scripts/reproject_canonical.py` dry-run/apply workflow. Read the source's
+current version first; target exactly the next version and retry the same pair
+on uncertain delivery. Apply only to selected native sources after installing the
+new mapper in their worker. Wait for publication completion before claiming those
+fields are searchable. Fresh captures use the current mapper. The content read
+continues to return the original body; the index read includes prepared metadata.
diff --git a/backend/airweave/domains/native_ingestion/__init__.py b/backend/airweave/domains/native_ingestion/__init__.py
new file mode 100644
index 0000000000..91b7093d82
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/__init__.py
@@ -0,0 +1 @@
+"""Trusted native snapshots; Almanac remains the content and access authority."""
diff --git a/backend/airweave/domains/native_ingestion/access_models.py b/backend/airweave/domains/native_ingestion/access_models.py
new file mode 100644
index 0000000000..9b39b33d0b
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/access_models.py
@@ -0,0 +1,42 @@
+"""Explicit source-access observations, separate from native content versions."""
+
+from typing import Annotated, Literal
+from uuid import UUID
+
+from pydantic import Field
+
+from airweave.domains.entities.canonical.requests import RecordIdentity, ScopeRemovalReason
+from airweave.domains.native_ingestion.models import NativeModel, NativeSnapshot
+
+
+class WithdrawNativeRecord(NativeModel):
+ """Publisher attests exact source access loss, never an incomplete-list absence."""
+
+ action: Literal["withdraw"] = "withdraw"
+ expected_revision: int = Field(strict=True, ge=1)
+ reason: ScopeRemovalReason
+
+
+class RenewNativeRecord(NativeModel):
+ """Publisher freshly read this original after inspecting the retained revision."""
+
+ action: Literal["renew"] = "renew"
+ expected_revision: int = Field(strict=True, ge=1)
+ expected_parent_epoch: int | None = Field(default=None, strict=True, ge=1)
+ snapshot: NativeSnapshot
+
+
+NativeAccessChange = Annotated[
+ WithdrawNativeRecord | RenewNativeRecord, Field(discriminator="action")
+]
+
+
+class NativeRecordAccess(NativeModel):
+ """Current state for CAS/recovery, without retained original bodies or credentials."""
+
+ record_id: UUID
+ identity: RecordIdentity
+ revision: int
+ parent_visibility_epoch: int | None
+ available: bool
+ removal_reason: str | None
diff --git a/backend/airweave/domains/native_ingestion/access_store.py b/backend/airweave/domains/native_ingestion/access_store.py
new file mode 100644
index 0000000000..c39c09447c
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/access_store.py
@@ -0,0 +1,155 @@
+"""Compose native access CAS with the canonical journal and parent visibility rules."""
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord
+from airweave.domains.entities.canonical.store import content_is_available, source_record
+from airweave.domains.native_ingestion.access_models import (
+ NativeAccessChange,
+ NativeRecordAccess,
+ WithdrawNativeRecord,
+)
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.models import IngestNativeBatch
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.models.entity import Entity
+from airweave.models.sync import SyncStatus
+
+
+class NativeAccessStore:
+ """The caller owns the transaction; no independent access journal or retry state."""
+
+ def __init__(self, imports: NativeImportStore):
+ """Reuse import locks, native admission and canonical capture."""
+ self.imports = imports
+ self.ingestion = NativeIngestionStore(imports.canonical)
+
+ async def _record(
+ self, db: AsyncSession, organization_id: UUID, sync_id: UUID, record_id: UUID
+ ) -> Entity:
+ row = await db.scalar(
+ select(Entity)
+ .where(
+ Entity.id == record_id,
+ Entity.organization_id == organization_id,
+ Entity.sync_id == sync_id,
+ Entity.record_revision > 0,
+ )
+ .execution_options(populate_existing=True)
+ )
+ if row is None:
+ raise NativeAdmissionError("Native record is unavailable in this source")
+ return row
+
+ async def _parent_epoch(self, db: AsyncSession, row: Entity) -> int | None:
+ """Native messages have one session parent; no arbitrary ancestor token scheme."""
+ if row.parent_record_type is None:
+ return None
+ record = source_record(row)
+ assert record.parent is not None
+ return await db.scalar(
+ select(Entity.visibility_epoch).where(
+ Entity.organization_id == row.organization_id,
+ Entity.sync_id == row.sync_id,
+ Entity.entity_definition_short_name == record.parent.record_type,
+ Entity.entity_id == record.parent.entity_key,
+ Entity.record_revision > 0,
+ )
+ )
+
+ async def read(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ record_id: UUID,
+ ) -> NativeRecordAccess:
+ """Read current CAS state even when the original is inaccessible or import terminal."""
+ bound, _, _ = await self.imports.load(db, organization_id, source_id, request_key)
+ row = await self._record(db, organization_id, bound.sync.id, record_id)
+ record = source_record(row)
+ available = await db.scalar(select(content_is_available()).where(Entity.id == row.id))
+ return NativeRecordAccess(
+ record_id=row.id,
+ identity=record.identity,
+ revision=row.record_revision,
+ parent_visibility_epoch=await self._parent_epoch(db, row),
+ available=bool(available)
+ and row.deleted_at is None
+ and bound.source.is_authenticated
+ and bound.sync.status == SyncStatus.ACTIVE,
+ removal_reason=row.removal_reason,
+ )
+
+ async def change(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ record_id: UUID,
+ request: NativeAccessChange,
+ ) -> NativeRecordAccess:
+ """Fresh source evidence and retained revision must agree under the active writer."""
+ imported = await self.imports.active(db, organization_id, source_id, request_key)
+ sync = await self.imports.canonical._fenced_sync(db, imported.fence)
+ row = await self._record(db, organization_id, sync.id, record_id)
+ if row.record_revision != request.expected_revision:
+ raise NativeAdmissionError("Native record changed; read current access before retrying")
+ record = source_record(row)
+ now = datetime.now(timezone.utc)
+ if isinstance(request, WithdrawNativeRecord):
+ # Access is a local observation. Preserve original content/version and blobs.
+ observations = (
+ CaptureRecord(
+ identity=record.identity,
+ parent=record.parent,
+ payload=record.payload,
+ payload_schema_version=record.payload_schema_version,
+ content_hash=record.content_hash,
+ completeness=record.completeness,
+ blobs=record.blobs,
+ kind="delete",
+ removal_reason=request.reason,
+ observed_at=now,
+ source_created_at=record.source_created_at,
+ source_updated_at=record.source_updated_at,
+ ),
+ )
+ else:
+ parent_epoch = await self._parent_epoch(db, row)
+ if request.expected_parent_epoch != parent_epoch or (
+ record.parent is not None and parent_epoch is None
+ ):
+ raise NativeAdmissionError(
+ "Native parent access changed; obtain fresh source evidence"
+ )
+ if (
+ request.snapshot.identity != record.identity
+ or request.snapshot.operation != "upsert"
+ ):
+ raise NativeAdmissionError("Renewal requires this exact active native original")
+ admitted = await self.ingestion.admit_locked(
+ db,
+ sync,
+ IngestNativeBatch(
+ fence=imported.fence,
+ observed_at=now,
+ snapshots=(request.snapshot,),
+ ),
+ revalidated_ids=(row.id,),
+ )
+ observations = admitted.records
+ await self.imports.canonical._capture_locked(
+ db,
+ sync,
+ CaptureBatch(fence=imported.fence, records=observations),
+ mark_seen=False,
+ )
+ return await self.read(db, organization_id, source_id, request_key, record_id)
diff --git a/backend/airweave/domains/native_ingestion/errors.py b/backend/airweave/domains/native_ingestion/errors.py
new file mode 100644
index 0000000000..a3b0f91e09
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/errors.py
@@ -0,0 +1,15 @@
+"""Stable native admission failures without source content in diagnostics."""
+
+from airweave.domains.entities.canonical.store import CanonicalStoreError
+
+
+class NativeAdmissionError(CanonicalStoreError):
+ """A native source, version or retained state cannot safely admit this write."""
+
+ code = "native_admission_failed"
+
+
+class NativeImportNotFound(NativeAdmissionError):
+ """No import exists at the authenticated source/request identity."""
+
+ code = "native_import_not_found"
diff --git a/backend/airweave/domains/native_ingestion/import_models.py b/backend/airweave/domains/native_ingestion/import_models.py
new file mode 100644
index 0000000000..22575a7650
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/import_models.py
@@ -0,0 +1,81 @@
+"""Native import intent and durable identity, without client-controlled writer fences."""
+
+from typing import Literal
+from uuid import UUID
+
+from pydantic import AwareDatetime, Field, model_validator
+
+from airweave.core.shared_models import SyncJobStatus
+from airweave.domains.entities.canonical.coverage_models import CompletionPolicy
+from airweave.domains.entities.canonical.requests import WriterFence
+from airweave.domains.native_ingestion.models import NativeModel
+
+
+class StartNativeImport(NativeModel):
+ """Publisher attestation; bounded roots do not authorize root absence deletion."""
+
+ snapshot_id: str = Field(min_length=1, max_length=256)
+ coverage: Literal["bounded", "complete"]
+ transcript_coverage: Literal["bounded", "complete"] = Field(
+ default="bounded",
+ description="Complete message enumeration for each selected session; sessions only",
+ )
+
+ def completion_policy(self, record_type: str) -> CompletionPolicy:
+ """Complete transcripts own only children of the selected session roots."""
+ if self.coverage == "complete" or (
+ record_type == "message" and self.transcript_coverage == "complete"
+ ):
+ return "exhaustive"
+ return "discovery_only"
+
+
+class NativeImportSummary(NativeModel):
+ """Durable capture outcome; indexing completion is separately observed."""
+
+ outcome: Literal["completed", "cancelled"]
+ coverage: Literal["bounded", "complete"]
+ finished_at: AwareDatetime
+ sequence: int | None = Field(default=None, ge=0)
+ completed_scopes: int = Field(ge=0)
+ capture_complete: bool
+ indexing: Literal["not_verified"] = "not_verified"
+
+ @model_validator(mode="after")
+ def consistent_outcome(self) -> "NativeImportSummary":
+ """Malformed retained outcomes must not become contradictory public claims."""
+ completed = self.outcome == "completed"
+ if self.capture_complete != completed or (completed and self.sequence is None):
+ raise ValueError("Native terminal summary contradicts its capture outcome")
+ return self
+
+
+class NativeImportReceipt(NativeModel):
+ """Server-only job metadata; an identical request never reactivates a writer."""
+
+ schema_version: Literal[1] = 1
+ source_id: UUID
+ request_key: str = Field(min_length=1, max_length=128)
+ request: StartNativeImport
+ fence: WriterFence
+ cycle_id: UUID
+ summary: NativeImportSummary | None = None
+
+ @model_validator(mode="after")
+ def consistent_coverage(self) -> "NativeImportReceipt":
+ """A terminal result cannot expand the publisher's original scope claim."""
+ if self.summary is not None and self.summary.coverage != self.request.coverage:
+ raise ValueError("Native terminal coverage contradicts import intent")
+ return self
+
+
+class NativeImportState(NativeModel):
+ """Public identity and execution state; running is not evidence of full capture."""
+
+ source_id: UUID
+ import_id: UUID
+ request_key: str
+ request: StartNativeImport
+ status: SyncJobStatus
+ cycle_id: UUID
+ summary: NativeImportSummary | None = None
diff --git a/backend/airweave/domains/native_ingestion/import_service.py b/backend/airweave/domains/native_ingestion/import_service.py
new file mode 100644
index 0000000000..a060ecc8da
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/import_service.py
@@ -0,0 +1,141 @@
+"""Atomic native import start/read application boundary."""
+
+from uuid import UUID
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.native_ingestion.access_models import NativeAccessChange, NativeRecordAccess
+from airweave.domains.native_ingestion.access_store import NativeAccessStore
+from airweave.domains.native_ingestion.import_models import NativeImportState, StartNativeImport
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.page_models import CommitNativePage, NativePageAck
+from airweave.domains.native_ingestion.page_store import NativePageStore
+from airweave.domains.native_ingestion.scope_models import (
+ BeginNativeScope,
+ NativeScopeRef,
+ NativeScopeState,
+ ReconcileNativeScope,
+)
+from airweave.domains.native_ingestion.scope_store import NativeScopeStore
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+
+
+class NativeImports:
+ """Own the transaction; never enqueue provider execution or expose writer fences."""
+
+ def __init__(self, store: NativeImportStore):
+ """Inject the native import persistence boundary."""
+ self.store = store
+ self.access = NativeAccessStore(store)
+ self.scopes = NativeScopeStore(store)
+ self.pages = NativePageStore(store, NativeIngestionStore(store.canonical))
+
+ async def start(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: StartNativeImport,
+ ) -> NativeImportState:
+ """Create or recover a durable import and commit its writer/cycle atomically."""
+ async with UnitOfWork(db):
+ return await self.store.start(db, organization_id, source_id, request_key, request)
+
+ async def read(
+ self, db: AsyncSession, organization_id: UUID, source_id: UUID, request_key: str
+ ) -> NativeImportState:
+ """Read durable progress under source authorization without reactivation."""
+ async with UnitOfWork(db):
+ return await self.store.read(db, organization_id, source_id, request_key)
+
+ async def begin_scope(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: BeginNativeScope,
+ ) -> NativeScopeState:
+ """Authorize and execute begin scope in one transaction."""
+ async with UnitOfWork(db):
+ return await self.scopes.begin(db, organization_id, source_id, request_key, request)
+
+ async def read_scope(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: NativeScopeRef,
+ ) -> NativeScopeState:
+ """Authorize and execute read scope in one transaction."""
+ async with UnitOfWork(db):
+ return await self.scopes.read(db, organization_id, source_id, request_key, request)
+
+ async def reconcile_scope(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: ReconcileNativeScope,
+ ) -> NativeScopeState:
+ """Authorize and execute reconcile scope in one transaction."""
+ async with UnitOfWork(db):
+ return await self.scopes.reconcile(db, organization_id, source_id, request_key, request)
+
+ async def page(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: CommitNativePage,
+ ) -> NativePageAck:
+ """Authorize and execute page in one transaction."""
+ async with UnitOfWork(db):
+ return await self.pages.commit(db, organization_id, source_id, request_key, request)
+
+ async def finish(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ *,
+ cancel: bool = False,
+ ) -> NativeImportState:
+ """Finalize capture and its durable outcome in one transaction."""
+ async with UnitOfWork(db):
+ return await self.store.finish(
+ db, organization_id, source_id, request_key, cancel=cancel
+ )
+
+ async def read_access(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ record_id: UUID,
+ ) -> NativeRecordAccess:
+ """Read retained revision and current availability without disclosing original content."""
+ async with UnitOfWork(db):
+ return await self.access.read(db, organization_id, source_id, request_key, record_id)
+
+ async def change_access(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ record_id: UUID,
+ request: NativeAccessChange,
+ ) -> NativeRecordAccess:
+ """Commit source access evidence and publication invalidation atomically."""
+ async with UnitOfWork(db):
+ return await self.access.change(
+ db, organization_id, source_id, request_key, record_id, request
+ )
diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py
new file mode 100644
index 0000000000..54c656248c
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/import_store.py
@@ -0,0 +1,290 @@
+"""Native import start/recovery in the existing source writer transaction."""
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from uuid import UUID, uuid4, uuid5
+
+from pydantic import ValidationError
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.core.shared_models import SyncStatus
+from airweave.domains.entities.canonical.cycle_models import (
+ BeginCycle,
+ CompleteCycle,
+ CycleConfiguration,
+ RestartCycle,
+)
+from airweave.domains.entities.canonical.cycle_store import (
+ begin_cycle,
+ complete_cycle,
+ cursor_row,
+ cycle_state,
+ restart_cycle,
+)
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.errors import NativeAdmissionError, NativeImportNotFound
+from airweave.domains.native_ingestion.import_models import (
+ NativeImportReceipt,
+ NativeImportState,
+ NativeImportSummary,
+ StartNativeImport,
+)
+from airweave.domains.native_ingestion.source_store import LockedNativeSource, NativeSourceStore
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.sync_job import SyncJob
+
+# Persisted v1 identity: changing this namespace would break retry recovery.
+IMPORT_NAMESPACE = UUID("8aab29e2-c0a6-40cd-a540-5316f1b6c4b3")
+
+
+def native_import_id(source_id: UUID, request_key: str) -> UUID:
+ """Unambiguous, source-scoped identity backed by the existing job primary key."""
+ if not 1 <= len(request_key) <= 128:
+ raise NativeAdmissionError("Import request key must contain 1 to 128 characters")
+ return uuid5(IMPORT_NAMESPACE, json.dumps([str(source_id), request_key], separators=(",", ":")))
+
+
+def receipt(job: SyncJob, source_id: UUID, request_key: str) -> NativeImportReceipt:
+ """Reject unrelated jobs and malformed state instead of adopting or restarting them."""
+ try:
+ value = NativeImportReceipt.model_validate(job.sync_metadata)
+ except ValidationError as error:
+ raise NativeAdmissionError("Native import receipt is malformed") from error
+ if (
+ (
+ value.summary is not None
+ and (
+ value.summary.outcome != job.status
+ or value.summary.coverage != value.request.coverage
+ or value.summary.capture_complete != (job.status == "completed")
+ )
+ )
+ or value.source_id != source_id
+ or value.request_key != request_key
+ or value.fence.job_id != job.id
+ or value.fence.sync_id != job.sync_id
+ or value.fence.organization_id != job.organization_id
+ ):
+ raise NativeAdmissionError("Native import receipt does not match its job")
+ return value
+
+
+def import_state(job: SyncJob, value: NativeImportReceipt) -> NativeImportState:
+ """Never return the internal writer fence to the publisher."""
+ return NativeImportState(
+ source_id=value.source_id,
+ import_id=job.id,
+ request_key=value.request_key,
+ request=value.request,
+ status=job.status,
+ cycle_id=value.cycle_id,
+ summary=value.summary,
+ )
+
+
+class NativeImportStore:
+ """No commits: caller owns one UnitOfWork for job, writer and cycle creation."""
+
+ def __init__(self, sources: NativeSourceStore, canonical: CanonicalRecordStore):
+ """Reuse native source ownership and canonical writer/cycle machinery."""
+ self.sources = sources
+ self.canonical = canonical
+
+ async def start(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: StartNativeImport,
+ ) -> NativeImportState:
+ """Recover exact requests, otherwise create one fenced snapshot import."""
+ bound = await self.sources.require(db, organization_id, source_id)
+ if request.transcript_coverage == "complete" and bound.binding.dataset != "sessions":
+ raise NativeAdmissionError("Complete transcript coverage requires a sessions source")
+ job_id = native_import_id(source_id, request_key)
+ job = await db.scalar(select(SyncJob).where(SyncJob.id == job_id).with_for_update())
+ if job is not None:
+ if job.sync_id != bound.sync.id or job.organization_id != organization_id:
+ raise NativeAdmissionError("Import job belongs to another source")
+ saved = receipt(job, source_id, request_key)
+ if saved.request != request:
+ raise NativeAdmissionError("Import request key has conflicting intent")
+ return import_state(job, saved)
+
+ if not bound.source.is_authenticated or bound.sync.status != SyncStatus.ACTIVE:
+ raise NativeAdmissionError("Native source is unavailable for import")
+ previous_job_id = bound.sync.writer_job_id
+ job = SyncJob(
+ id=job_id,
+ sync_id=bound.sync.id,
+ organization_id=organization_id,
+ provisioning_generation=bound.sync.provisioning_generation,
+ status="running",
+ started_at=datetime.now(timezone.utc).replace(tzinfo=None),
+ )
+ db.add(job)
+ await db.flush()
+ fence = await self.canonical.activate_writer(
+ db,
+ organization_id,
+ bound.sync.id,
+ job_id,
+ attempt_id=uuid4(),
+ attempt_number=1,
+ )
+ plan = {"native_import_id": str(job_id), **request.model_dump(mode="json")}
+ fingerprint = hashlib.sha256(json.dumps(plan, sort_keys=True).encode()).hexdigest()
+ parents = (
+ {"knowledge": (None,)}
+ if bound.binding.dataset == "knowledge"
+ else {"session": (None,), "message": ("session",)}
+ )
+ configuration = CycleConfiguration(
+ fingerprint=fingerprint,
+ parents=parents,
+ completion_policies={kind: request.completion_policy(kind) for kind in parents},
+ membership="observed" if request.coverage == "bounded" else "retained",
+ )
+ previous = cycle_state(await cursor_row(db, fence))
+ if previous is not None and previous.source_plan.get("native_import_id") != str(
+ previous_job_id
+ ):
+ raise NativeAdmissionError("Existing capture cycle belongs to another import lifecycle")
+ if previous is not None and previous.phase == "active":
+ # Activation already rejects a different live writer. A new import after
+ # a terminal job explicitly replaces its incomplete enumeration only.
+ cycle = await restart_cycle(
+ db,
+ RestartCycle(
+ fence=fence,
+ expected=previous.version,
+ configuration=configuration,
+ source_plan=plan,
+ ),
+ )
+ else:
+ cycle = await begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence,
+ expected=previous.version if previous else None,
+ configuration=configuration,
+ source_plan=plan,
+ ),
+ )
+ saved = NativeImportReceipt(
+ source_id=source_id,
+ request_key=request_key,
+ request=request,
+ fence=fence,
+ cycle_id=cycle.version.cycle_id,
+ )
+ job.sync_metadata = saved.model_dump(mode="json")
+ await db.flush()
+ return import_state(job, saved)
+
+ async def read(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ ) -> NativeImportState:
+ """Authorize actual source/job association, including terminal import retries."""
+ _, job, saved = await self.load(db, organization_id, source_id, request_key)
+ return import_state(job, saved)
+
+ async def active(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ ) -> NativeImportReceipt:
+ """Resolve stored authority for a mutation, never a publisher-supplied fence."""
+ bound, _, saved = await self.load(db, organization_id, source_id, request_key)
+ if not bound.source.is_authenticated or bound.sync.status != SyncStatus.ACTIVE:
+ raise NativeAdmissionError("Native source is unavailable for import")
+ await self.canonical._fenced_sync(db, saved.fence)
+ return saved
+
+ async def load(
+ self, db: AsyncSession, organization_id: UUID, source_id: UUID, request_key: str
+ ) -> tuple[LockedNativeSource, SyncJob, NativeImportReceipt]:
+ """Authorize and lock the actual job, even after its writer became terminal."""
+ bound = await self.sources.require(db, organization_id, source_id)
+ job = await db.scalar(
+ select(SyncJob)
+ .where(
+ SyncJob.id == native_import_id(source_id, request_key),
+ SyncJob.sync_id == bound.sync.id,
+ SyncJob.organization_id == organization_id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if job is None:
+ raise NativeImportNotFound("Native import does not exist")
+ return bound, job, receipt(job, source_id, request_key)
+
+ async def finish(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ *,
+ cancel: bool = False,
+ ) -> NativeImportState:
+ """Persist capture completion or cancellation without changing another writer."""
+ bound, job, saved = await self.load(db, organization_id, source_id, request_key)
+ if job.status in ("completed", "cancelled", "failed"):
+ if cancel or (job.status == "completed" and saved.summary is not None):
+ return import_state(job, saved)
+ raise NativeAdmissionError("Import is terminal without successful capture completion")
+ if cancel:
+ # Even an older job may be cancelled; only its status fences its writes.
+ outcome = "cancelled"
+ else:
+ if not bound.source.is_authenticated or bound.sync.status != SyncStatus.ACTIVE:
+ raise NativeAdmissionError("Native source is unavailable for import")
+ await self.canonical._fenced_sync(db, saved.fence)
+ current = cycle_state(await cursor_row(db, saved.fence))
+ if current is None or current.version.cycle_id != saved.cycle_id:
+ raise NativeAdmissionError("Import no longer owns its capture cycle")
+ await complete_cycle(
+ db, bound.sync, CompleteCycle(fence=saved.fence, expected=current.version)
+ )
+ outcome = "completed"
+ finished_at = datetime.now(timezone.utc)
+ count = await db.scalar(
+ select(func.count())
+ .select_from(CaptureScan)
+ .where(
+ CaptureScan.organization_id == organization_id,
+ CaptureScan.sync_id == bound.sync.id,
+ CaptureScan.cycle_id == saved.cycle_id,
+ CaptureScan.phase == "complete",
+ )
+ )
+ # A superseded cancellation must not claim the newer writer's sequence.
+ sequence = (
+ bound.sync.observed_change_sequence if bound.sync.writer_job_id == job.id else None
+ )
+ summary = NativeImportSummary(
+ outcome=outcome,
+ coverage=saved.request.coverage,
+ finished_at=finished_at,
+ sequence=sequence,
+ completed_scopes=count,
+ capture_complete=not cancel,
+ )
+ saved = saved.model_copy(update={"summary": summary})
+ job.sync_metadata = saved.model_dump(mode="json")
+ job.status = outcome
+ job.completed_at = finished_at.replace(tzinfo=None)
+ await db.flush()
+ return import_state(job, saved)
diff --git a/backend/airweave/domains/native_ingestion/knowledge_fields.py b/backend/airweave/domains/native_ingestion/knowledge_fields.py
new file mode 100644
index 0000000000..a3aa6f78b2
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/knowledge_fields.py
@@ -0,0 +1,123 @@
+"""Searchable native wiki fields, independent of body text and reference resolution."""
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictStr
+
+
+class _Fields(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+
+
+class _Link(_Fields):
+ label: StrictStr
+ url: StrictStr
+ handle: StrictStr | None = None
+
+
+class _Role(_Fields):
+ title: StrictStr
+ start: StrictStr | None = None
+ end: StrictStr | None = None
+
+
+class _Education(_Fields):
+ degree: StrictStr | None = None
+ subject: StrictStr | None = None
+ start: StrictStr | None = None
+ end: StrictStr | None = None
+
+
+class _RelatedPerson(_Fields):
+ relationship: StrictStr
+
+
+class _Person(_Fields):
+ related_people: tuple[_RelatedPerson, ...] = ()
+ emails: tuple[StrictStr, ...] = ()
+ phones: tuple[StrictStr, ...] = ()
+ timezone: StrictStr | None = None
+ roles: tuple[_Role, ...] = ()
+ experience: tuple[_Role, ...] = ()
+ education: tuple[_Education, ...] = ()
+ links: tuple[_Link, ...] = ()
+
+
+class _Organisation(_Fields):
+ related_people: tuple[_RelatedPerson, ...] = ()
+ domains: tuple[StrictStr, ...] = ()
+ website: StrictStr | None = None
+ industry: tuple[StrictStr, ...] = ()
+ founded: StrictStr | None = None
+ links: tuple[_Link, ...] = ()
+
+
+class _Address(_Fields):
+ street: StrictStr | None = None
+ locality: StrictStr | None = None
+ region: StrictStr | None = None
+ postal_code: StrictStr | None = None
+ country: StrictStr | None = None
+
+
+class _Coordinates(_Fields):
+ latitude: float = Field(ge=-90, le=90, allow_inf_nan=False)
+ longitude: float = Field(ge=-180, le=180, allow_inf_nan=False)
+
+
+class _Place(_Fields):
+ coordinates: _Coordinates | None = None
+ place_kind: StrictStr | None = None
+ address: _Address | None = None
+
+
+class _Timed(_Fields):
+ kind: Literal["timed"]
+ start: StrictStr
+ end: StrictStr
+ timezone: StrictStr
+
+
+class _AllDay(_Fields):
+ kind: Literal["all_day"]
+ start_on: StrictStr
+ end_on_exclusive: StrictStr
+
+
+class _Event(_Fields):
+ schedule: Annotated[_Timed | _AllDay, Field(discriminator="kind")] | None = None
+
+
+class _CreativeWork(_Fields):
+ work_kind: StrictStr | None = None
+ published_on: StrictStr | None = None
+
+
+def _text(value: _Fields) -> str:
+ """Flatten only the validated allowlist, never arbitrary native payload metadata."""
+ lines = []
+ for key, item in value.model_dump(exclude_none=True).items():
+ label = key.replace("_", " ")
+ if isinstance(item, str):
+ lines.append(f"{label}: {item}")
+ elif isinstance(item, list | tuple):
+ for entry in item:
+ if isinstance(entry, str):
+ lines.append(f"{label}: {entry}")
+ elif entry:
+ lines.append(f"{label}: " + "; ".join(f"{k}: {v}" for k, v in entry.items()))
+ elif isinstance(item, dict) and item:
+ lines.append(f"{label}: " + "; ".join(f"{k}: {v}" for k, v in item.items()))
+ return "\n".join(lines)
+
+
+def knowledge_details(kind: str, original: dict[str, JsonValue]) -> str | None:
+ """Do not fetch referenced records or invent names from their opaque IDs."""
+ model = {
+ "person": _Person,
+ "organisation": _Organisation,
+ "place": _Place,
+ "creative_work": _CreativeWork,
+ "event": _Event,
+ }.get(kind)
+ return _text(model.model_validate(original)) or None if model else None
diff --git a/backend/airweave/domains/native_ingestion/models.py b/backend/airweave/domains/native_ingestion/models.py
new file mode 100644
index 0000000000..3d7a0b4198
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/models.py
@@ -0,0 +1,120 @@
+"""Versioned native snapshots, distinct from provider observations and originals."""
+
+from typing import Annotated, Literal
+from uuid import UUID
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator
+
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+ ScanVersion,
+ WriterFence,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+
+
+class NativeModel(BaseModel):
+ """Strict immutable boundary for publisher-owned data."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False)
+
+
+class RecordVersion(NativeModel):
+ """Almanac knowledge record revision, independent of capture revision."""
+
+ kind: Literal["record"] = "record"
+ revision: int = Field(strict=True, ge=1)
+
+
+class SessionVersion(NativeModel):
+ """Metadata and original-content revisions; replay epochs are not history."""
+
+ kind: Literal["session"] = "session"
+ revision: int = Field(strict=True, ge=1)
+ content_revision: int = Field(strict=True, ge=0)
+
+
+NativeVersion = Annotated[RecordVersion | SessionVersion, Field(discriminator="kind")]
+
+
+class NativeSourceBinding(NativeModel):
+ """Server-provisioned owner and dataset; never supplied by an untrusted reader."""
+
+ owner_id: str = Field(min_length=1)
+ dataset: Literal["knowledge", "sessions"]
+
+
+class NativeSnapshot(NativeModel):
+ """An attested copy of native JSON, not a second editable authority."""
+
+ authority: Literal["almanac"] = "almanac"
+ representation: Literal["snapshot"] = "snapshot"
+ schema_version: Literal[1] = 1
+ owner_id: str = Field(min_length=1)
+ identity: RecordIdentity
+ parent: RecordIdentity | None = None
+ version: NativeVersion
+ operation: Literal["upsert", "delete"] = "upsert"
+ original: dict[str, JsonValue]
+ source_created_at: AwareDatetime | None = None
+ source_updated_at: AwareDatetime | None = None
+
+ @model_validator(mode="after")
+ def identity_contract(self) -> "NativeSnapshot":
+ """Keep version domains and session child identity explicit."""
+ if self.identity.record_type == "knowledge":
+ if self.version.kind != "record" or self.parent is not None:
+ raise ValueError("Knowledge snapshots require record versions and no parent")
+ elif self.identity.record_type == "session":
+ if self.version.kind != "session" or self.parent is not None:
+ raise ValueError("Session snapshots require session versions and no parent")
+ elif self.identity.record_type == "message":
+ if (
+ self.version.kind != "session"
+ or self.parent is None
+ or self.parent.record_type != "session"
+ or self.parent.container_id is not None
+ or self.identity.container_id != self.parent.native_id
+ ):
+ raise ValueError("Original messages require their session parent and version")
+ else:
+ raise ValueError("Unsupported native snapshot record type")
+ if self.identity.record_type != "message" and self.identity.container_id is not None:
+ raise ValueError("Native roots have no container")
+ return self
+
+
+class IngestNativeBatch(NativeModel):
+ """Internal service command; the authenticated edge must issue the writer fence."""
+
+ fence: WriterFence
+ observed_at: AwareDatetime
+ snapshots: tuple[NativeSnapshot, ...] = Field(min_length=1, max_length=500)
+
+ @model_validator(mode="after")
+ def unique_records(self) -> "IngestNativeBatch":
+ """One version per identity makes atomic admission unambiguous."""
+ keys = [(item.identity.record_type, item.identity.entity_key) for item in self.snapshots]
+ if len(set(keys)) != len(keys):
+ raise ValueError("A native batch must contain unique record identities")
+ return self
+
+
+class NativeAdmission(NativeModel):
+ """Admitted changes and exact retained retries under the current writer lock."""
+
+ records: tuple[CaptureRecord, ...]
+ unchanged_ids: tuple[UUID, ...]
+
+
+class IngestNativePage(IngestNativeBatch):
+ """Native originals plus the existing scan CAS; empty final pages are valid."""
+
+ snapshots: tuple[NativeSnapshot, ...] = Field(max_length=500)
+ scope: CompletedScope
+ cycle_id: UUID
+ expected: ScanVersion
+ continuation: ScanContinuation
+ final: bool = False
diff --git a/backend/airweave/domains/native_ingestion/page_models.py b/backend/airweave/domains/native_ingestion/page_models.py
new file mode 100644
index 0000000000..3cc4f3870f
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/page_models.py
@@ -0,0 +1,56 @@
+"""Bounded publisher page commands and durable acknowledgements."""
+
+import json
+from typing import Literal
+from uuid import UUID
+
+from pydantic import Field, JsonValue, field_validator, model_validator
+
+from airweave.domains.entities.canonical.requests import CompletedScope, ScanVersion
+from airweave.domains.native_ingestion.models import NativeModel, NativeSnapshot
+
+
+class CommitNativePage(NativeModel):
+ """One in-flight page per scope; preserve this exact body on an uncertain retry."""
+
+ page_id: UUID
+ scope: CompletedScope
+ expected: ScanVersion
+ snapshots: tuple[NativeSnapshot, ...] = Field(max_length=500)
+ cursor: dict[str, JsonValue] = Field(default_factory=dict)
+ final: bool = False
+
+ @field_validator("cursor")
+ @classmethod
+ def bounded_cursor(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]:
+ """Leave space for the receipt inside the existing 64KiB scan continuation."""
+ if len(json.dumps(value, allow_nan=False).encode()) > 32768:
+ raise ValueError("Native page cursor exceeds 32KiB")
+ return value
+
+ @model_validator(mode="after")
+ def unique_snapshots(self) -> "CommitNativePage":
+ """Reject duplicate originals before attempting any page writes."""
+ keys = [(item.identity.record_type, item.identity.entity_key) for item in self.snapshots]
+ if len(keys) != len(set(keys)):
+ raise ValueError("Native page contains duplicate originals")
+ return self
+
+
+class NativePageAck(NativeModel):
+ """Durable capture acknowledgement, not search publication or import completion."""
+
+ page_id: UUID
+ version: ScanVersion
+ phase: Literal["collecting", "reconciling", "complete"]
+ sequence: int
+ changed: int
+ unchanged: int
+
+
+class NativePageReceipt(NativeModel):
+ """Only the most recently committed page is retryable with identical success."""
+
+ schema_version: Literal[1] = 1
+ digest: str = Field(pattern=r"^[a-f0-9]{64}$")
+ acknowledgement: NativePageAck
diff --git a/backend/airweave/domains/native_ingestion/page_store.py b/backend/airweave/domains/native_ingestion/page_store.py
new file mode 100644
index 0000000000..375f8c5892
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/page_store.py
@@ -0,0 +1,100 @@
+"""Atomic page admission and bounded lost-response recovery on existing scan state."""
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from uuid import UUID
+
+from pydantic import ValidationError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.entities.canonical.scan_store import CanonicalScanStore, ScanConflict
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.models import IngestNativePage
+from airweave.domains.native_ingestion.page_models import (
+ CommitNativePage,
+ NativePageAck,
+ NativePageReceipt,
+)
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+
+_RECEIPT = "native_page_receipt"
+
+
+class NativePageStore:
+ """Caller owns one transaction; no side journal and no retained response payload copies."""
+
+ def __init__(self, imports: NativeImportStore, ingestion: NativeIngestionStore):
+ """Share the existing canonical writer and scan boundary."""
+ self.imports = imports
+ self.ingestion = ingestion
+ self.scans = CanonicalScanStore(ingestion.canonical)
+
+ async def commit(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: CommitNativePage,
+ ) -> NativePageAck:
+ """A committed current page retry returns its receipt without another capture."""
+ imported = await self.imports.active(db, organization_id, source_id, request_key)
+ current = await self.scans.read(db, imported.fence, request.scope)
+ if current is None:
+ raise NativeAdmissionError("Native scope has not been started")
+ if current.cycle_id != imported.cycle_id:
+ raise NativeAdmissionError("Scope belongs to another native import")
+ digest = hashlib.sha256(
+ json.dumps(
+ request.model_dump(mode="json"),
+ sort_keys=True,
+ separators=(",", ":"),
+ ensure_ascii=False,
+ allow_nan=False,
+ ).encode()
+ ).hexdigest()
+ stored = current.continuation.value.get(_RECEIPT)
+ if stored is not None:
+ try:
+ previous = NativePageReceipt.model_validate(stored)
+ except ValidationError as error:
+ raise NativeAdmissionError("Native page receipt is malformed") from error
+ if previous.acknowledgement.page_id == request.page_id:
+ if previous.digest != digest or previous.acknowledgement.version != current.version:
+ raise ScanConflict("Page retry conflicts with committed progress")
+ return previous.acknowledgement
+ if current.version != request.expected:
+ raise ScanConflict("Page changed; read committed scope progress before retrying")
+ result = await self.ingestion.page(
+ db,
+ IngestNativePage(
+ fence=imported.fence,
+ observed_at=datetime.now(timezone.utc),
+ snapshots=request.snapshots,
+ scope=request.scope,
+ cycle_id=imported.cycle_id,
+ expected=request.expected,
+ continuation=ScanContinuation(value={"cursor": request.cursor}),
+ final=request.final,
+ ),
+ )
+ ack = NativePageAck(
+ page_id=request.page_id,
+ version=result.state.version,
+ phase=result.state.phase,
+ sequence=result.capture.sequence,
+ changed=len(result.capture.changes),
+ unchanged=result.capture.unchanged,
+ )
+ # The scan owns this row and its lock; receipt and page share the outer commit.
+ row = await self.scans._row(db, imported.fence, request.scope)
+ assert row is not None # Existing scan remains locked through this transaction.
+ row.continuation = {
+ "cursor": request.cursor,
+ _RECEIPT: NativePageReceipt(digest=digest, acknowledgement=ack).model_dump(mode="json"),
+ }
+ await db.flush()
+ return ack
diff --git a/backend/airweave/domains/native_ingestion/projection.py b/backend/airweave/domains/native_ingestion/projection.py
new file mode 100644
index 0000000000..55ad69c04d
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/projection.py
@@ -0,0 +1,221 @@
+"""Offline views of original Almanac records, sessions and protocol messages."""
+
+from typing import Literal
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool
+
+from airweave.domains.entities.canonical.extraction_models import ExtractionPart
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.projection_inputs import ProjectionInput, ProjectionInputs
+from airweave.domains.native_ingestion.knowledge_fields import knowledge_details
+from airweave.domains.native_ingestion.models import NativeSnapshot
+from airweave.domains.sync_pipeline.pipeline.text_models import NativeTextBody
+from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields
+from airweave.platform.entities._airweave_field import AirweaveField
+from airweave.platform.entities._base import BaseEntity
+
+
+class _Original(BaseModel):
+ """Read only declared projection fields; all other original JSON stays retained."""
+
+ model_config = ConfigDict(extra="ignore")
+ id: str = Field(min_length=1)
+ created_at: AwareDatetime
+
+
+class _Knowledge(_Original):
+ type: Literal["person", "organisation", "place", "event", "creative_work", "topic", "page"]
+ revision: int = Field(strict=True, ge=1)
+ title: str = Field(min_length=1)
+ description: str = Field(min_length=1)
+ path: str = Field(min_length=1)
+ body: str
+ updated_at: AwareDatetime
+ archived_at: AwareDatetime | None
+ user_notes: str | None = None
+
+
+class _Session(_Original):
+ title: str = Field(min_length=1)
+ description: str = Field(min_length=1)
+ revision: int = Field(strict=True, ge=1)
+ content_revision: int = Field(strict=True, ge=0)
+ archived: StrictBool
+ updated_at: AwareDatetime
+
+
+class _ImportProvenance(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ legacy_row_id: int = Field(strict=True, gt=0)
+ active: StrictBool
+ compacted: StrictBool
+
+
+class _TextBlock(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ type: Literal["text"]
+ text: str
+
+
+class _Block(BaseModel):
+ model_config = ConfigDict(extra="allow")
+ type: str = Field(min_length=1)
+
+
+class _Payload(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ role: Literal["user", "assistant", "tool", "system"]
+ content: str | list[dict[str, JsonValue]] | None
+
+
+class _Message(_Original):
+ ordinal: int = Field(strict=True, ge=0)
+ import_provenance: _ImportProvenance | None = None
+ payload: _Payload
+
+
+class NativeContentEntity(BaseEntity):
+ """Explicit searchable metadata; original payloads never enter entity serialization."""
+
+ native_key: str = AirweaveField(..., is_entity_id=True, embeddable=False)
+ title: str = AirweaveField(..., is_name=True, embeddable=True)
+ description: str | None = AirweaveField(None, embeddable=True)
+ user_notes: str | None = AirweaveField(None, embeddable=True)
+ details: str | None = AirweaveField(None, embeddable=True)
+ native_type: str = AirweaveField(..., embeddable=False)
+ path: str | None = AirweaveField(None, embeddable=False)
+ session_id: str | None = AirweaveField(None, embeddable=False)
+ message_id: str | None = AirweaveField(None, embeddable=False)
+ message_ordinal: int | None = AirweaveField(None, embeddable=False)
+ block_index: int | None = AirweaveField(None, embeddable=False)
+ role: str | None = AirweaveField(None, embeddable=False)
+ archived: bool | None = AirweaveField(None, embeddable=False)
+
+
+def _read(record: SourceRecord) -> tuple[NativeSnapshot, _Knowledge | _Session | _Message]:
+ snapshot = NativeSnapshot.model_validate(record.payload)
+ if snapshot.identity != record.identity or snapshot.parent != record.parent:
+ raise ValueError("Native snapshot identity differs from retained record")
+ if snapshot.operation != "upsert":
+ raise ValueError("Native tombstone cannot project content")
+ match snapshot.identity.record_type:
+ case "knowledge":
+ original = _Knowledge.model_validate(snapshot.original)
+ if snapshot.version.kind != "record" or original.revision != snapshot.version.revision:
+ raise ValueError("Knowledge version differs from original")
+ case "session":
+ original = _Session.model_validate(snapshot.original)
+ if (
+ snapshot.version.kind != "session"
+ or original.revision != snapshot.version.revision
+ or original.content_revision != snapshot.version.content_revision
+ ):
+ raise ValueError("Session version differs from original")
+ case "message":
+ original = _Message.model_validate(snapshot.original)
+ case _:
+ raise ValueError("Unsupported native projection type")
+ _validate_identity(record, snapshot, original)
+ return snapshot, original
+
+
+def _validate_identity(
+ record: SourceRecord, snapshot: NativeSnapshot, original: _Knowledge | _Session | _Message
+) -> None:
+ if original.id != snapshot.identity.native_id:
+ raise ValueError("Native original ID differs from snapshot")
+ if (
+ original.created_at != record.source_created_at
+ or original.created_at != snapshot.source_created_at
+ ):
+ raise ValueError("Native original creation timestamp differs from snapshot")
+ if isinstance(original, (_Knowledge, _Session)) and (
+ original.updated_at != record.source_updated_at
+ or original.updated_at != snapshot.source_updated_at
+ ):
+ raise ValueError("Native original update timestamp differs from snapshot")
+
+
+def excluded_native(record: SourceRecord) -> bool:
+ """Follow current knowledge and session search visibility policy."""
+ _, original = _read(record)
+ if isinstance(original, _Knowledge):
+ return original.archived_at is not None
+ return (
+ isinstance(original, _Message)
+ and original.import_provenance is not None
+ and not original.import_provenance.active
+ )
+
+
+def _part(index: int, key: str, entity: NativeContentEntity, text: str) -> ProjectionInput:
+ populate_base_fields(entity)
+ return ProjectionInput(
+ part=ExtractionPart(part_index=index, key=key, kind="body", media_type="text/plain"),
+ entity=entity,
+ native_body=NativeTextBody(text=text),
+ )
+
+
+def map_native(record: SourceRecord) -> ProjectionInputs:
+ """Select verbatim bodies and individual text blocks without flattening protocol JSON."""
+ snapshot, original = _read(record)
+ key = snapshot.identity.native_id
+ common = {
+ "created_at": record.source_created_at,
+ "updated_at": record.source_updated_at,
+ "breadcrumbs": [],
+ }
+ if isinstance(original, _Knowledge):
+ entity = NativeContentEntity(
+ native_key=key,
+ title=original.title,
+ description=original.description,
+ user_notes=original.user_notes,
+ details=knowledge_details(original.type, snapshot.original),
+ path=original.path,
+ native_type=original.type,
+ **common,
+ )
+ return ProjectionInputs(parts=(_part(0, key, entity, original.body),))
+ if isinstance(original, _Session):
+ entity = NativeContentEntity(
+ native_key=key,
+ title=original.title,
+ native_type="session",
+ session_id=key,
+ archived=original.archived,
+ **common,
+ )
+ return ProjectionInputs(parts=(_part(0, key, entity, original.description),))
+ content = original.payload.content
+ blocks = [{"type": "text", "text": content}] if isinstance(content, str) else content
+ # Null or an empty list is retained nontext protocol content, not invented text.
+ blocks = blocks or [{"type": "nontext"}]
+ parts = []
+ for index, block in enumerate(blocks):
+ kind = _Block.model_validate(block).type
+ part_key = f"{snapshot.identity.entity_key}:block:{index}"
+ if kind != "text":
+ parts.append(
+ ProjectionInput(
+ part=ExtractionPart(part_index=index, key=part_key, kind="record"),
+ entity=None,
+ omission="unsupported_format",
+ )
+ )
+ continue
+ text = _TextBlock.model_validate(block).text
+ entity = NativeContentEntity(
+ native_key=part_key,
+ title=f"{original.payload.role} message",
+ native_type="message",
+ session_id=snapshot.identity.container_id,
+ message_id=original.id,
+ message_ordinal=original.ordinal,
+ block_index=index,
+ role=original.payload.role,
+ **common,
+ )
+ parts.append(_part(index, part_key, entity, text))
+ return ProjectionInputs(parts=tuple(parts))
diff --git a/backend/airweave/domains/native_ingestion/scope_models.py b/backend/airweave/domains/native_ingestion/scope_models.py
new file mode 100644
index 0000000000..87d431415b
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/scope_models.py
@@ -0,0 +1,42 @@
+"""Publisher-visible scope progress, without storage receipts or writer credentials."""
+
+from typing import Literal
+from uuid import UUID
+
+from pydantic import Field, JsonValue
+
+from airweave.domains.entities.canonical.requests import CompletedScope, ScanVersion
+from airweave.domains.native_ingestion.models import NativeModel
+from airweave.domains.native_ingestion.page_models import NativePageAck
+
+
+class NativeScopeRef(NativeModel):
+ """Exact native container identity; a read does not start or restart work."""
+
+ scope: CompletedScope
+
+
+class BeginNativeScope(NativeScopeRef):
+ """Resume an unchanged scope; explicit same-import restarts require its version."""
+
+ expected: ScanVersion | None = None
+ restart: bool = False
+
+
+class ReconcileNativeScope(NativeScopeRef):
+ """Advance bounded reconciliation only after the final page is committed."""
+
+ expected: ScanVersion
+ limit: int = Field(default=250, ge=1, le=500)
+
+
+class NativeScopeState(NativeModel):
+ """Capture progress; complete refers to this scope's declared coverage policy."""
+
+ scope: CompletedScope
+ cycle_id: UUID
+ version: ScanVersion
+ phase: Literal["collecting", "reconciling", "complete"]
+ coverage: Literal["bounded", "complete"]
+ cursor: dict[str, JsonValue]
+ last_page: NativePageAck | None = None
diff --git a/backend/airweave/domains/native_ingestion/scope_store.py b/backend/airweave/domains/native_ingestion/scope_store.py
new file mode 100644
index 0000000000..95b713b447
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/scope_store.py
@@ -0,0 +1,121 @@
+"""Native scope lifecycle composed with canonical scope admission and reconciliation."""
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+from pydantic import ValidationError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.cycle_store import attest_cycle, scope_owner
+from airweave.domains.entities.canonical.scan_models import BeginScan, ReconcileScan, ScanState
+from airweave.domains.entities.canonical.scan_store import CanonicalScanStore
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.page_models import NativePageReceipt
+from airweave.domains.native_ingestion.scope_models import (
+ BeginNativeScope,
+ NativeScopeRef,
+ NativeScopeState,
+ ReconcileNativeScope,
+)
+
+
+def scope_state(state: ScanState) -> NativeScopeState:
+ """Hide the internal receipt envelope, retaining its bounded last-page acknowledgement."""
+ try:
+ saved = state.continuation.value.get("native_page_receipt")
+ ack = NativePageReceipt.model_validate(saved).acknowledgement if saved is not None else None
+ return NativeScopeState(
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ version=state.version,
+ phase=state.phase,
+ coverage="complete" if state.completion_policy == "exhaustive" else "bounded",
+ cursor=state.continuation.value.get("cursor", {}),
+ last_page=ack,
+ )
+ except ValidationError as error:
+ raise NativeAdmissionError("Native scope progress is malformed") from error
+
+
+class NativeScopeStore:
+ """No commits; import authorization and scope mutation share the outer transaction."""
+
+ def __init__(self, imports: NativeImportStore):
+ """Use one canonical scan implementation for native and provider correctness."""
+ self.imports = imports
+ self.scans = CanonicalScanStore(imports.canonical)
+
+ async def begin(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: BeginNativeScope,
+ ) -> NativeScopeState:
+ """Derive current parent authority server-side; do not reset same-import scans silently."""
+ imported = await self.imports.active(db, organization_id, source_id, request_key)
+ _, cycle = await attest_cycle(db, imported.fence, imported.cycle_id)
+ parent = await scope_owner(db, imported.fence, cycle, request.scope)
+ previous = await self.scans.read(db, imported.fence, request.scope)
+ expected = request.expected
+ if expected is None and previous is not None and previous.cycle_id != imported.cycle_id:
+ # A new import explicitly replaces a preceding cycle, under the same writer lock.
+ expected = previous.version
+ state = await self.scans.begin(
+ db,
+ BeginScan(
+ fence=imported.fence,
+ scope=request.scope,
+ cycle_id=imported.cycle_id,
+ fingerprint=cycle.configuration.fingerprint,
+ expected=expected,
+ restart=request.restart,
+ expected_parent_epoch=parent.visibility_epoch if parent else None,
+ ),
+ )
+ return scope_state(state)
+
+ async def read(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: NativeScopeRef,
+ ) -> NativeScopeState:
+ """Read retained scope progress only while its row still belongs to this import."""
+ _, _, imported = await self.imports.load(db, organization_id, source_id, request_key)
+ row = await self.scans._row(db, imported.fence, request.scope)
+ if row is None or row.cycle_id != imported.cycle_id:
+ raise NativeAdmissionError("Scope is unavailable or superseded for this import")
+ state = (await self.scans._state(db, row)).model_copy(
+ update={
+ "completion_policy": imported.request.completion_policy(request.scope.record_type)
+ }
+ )
+ return scope_state(state)
+
+ async def reconcile(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ request_key: str,
+ request: ReconcileNativeScope,
+ ) -> NativeScopeState:
+ """The canonical completion policy controls whether absence can withdraw records."""
+ imported = await self.imports.active(db, organization_id, source_id, request_key)
+ result = await self.scans.reconcile(
+ db,
+ ReconcileScan(
+ fence=imported.fence,
+ scope=request.scope,
+ cycle_id=imported.cycle_id,
+ expected=request.expected,
+ observed_at=datetime.now(timezone.utc),
+ limit=request.limit,
+ ),
+ )
+ return scope_state(result.state)
diff --git a/backend/airweave/domains/native_ingestion/service.py b/backend/airweave/domains/native_ingestion/service.py
new file mode 100644
index 0000000000..ae618839c8
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/service.py
@@ -0,0 +1,27 @@
+"""Atomic internal native ingestion; no provider I/O or public credential handling."""
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.models import CaptureResult
+from airweave.domains.entities.canonical.scan_models import ScanResult
+from airweave.domains.native_ingestion.models import IngestNativeBatch, IngestNativePage
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+
+
+class NativeIngestionService:
+ """Commit admission and canonical capture together, or neither."""
+
+ def __init__(self, store: NativeIngestionStore):
+ """Inject persistence; callers supply a fresh session without pending writes."""
+ self.store = store
+
+ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureResult:
+ """Admit one bounded snapshot batch under its current writer fence."""
+ async with UnitOfWork(db):
+ return await self.store.ingest(db, request)
+
+ async def page(self, db: AsyncSession, request: IngestNativePage) -> ScanResult:
+ """Commit admitted originals, sweep sightings and continuation together."""
+ async with UnitOfWork(db):
+ return await self.store.page(db, request)
diff --git a/backend/airweave/domains/native_ingestion/source_models.py b/backend/airweave/domains/native_ingestion/source_models.py
new file mode 100644
index 0000000000..0216e14eec
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/source_models.py
@@ -0,0 +1,48 @@
+"""Immutable native source identity and backend-only ensure/read contracts."""
+
+import json
+from uuid import UUID, uuid5
+
+from pydantic import Field
+
+from airweave.domains.native_ingestion.models import NativeModel, NativeSourceBinding
+
+# Versioned protocol constants: changing these would create duplicate source identities.
+_SOURCE_NAMESPACE_V1 = UUID("b13e2720-64ab-4ac5-9a0f-6264e2fb783f")
+_SYNC_NAMESPACE_V1 = UUID("01b71224-aad4-4a85-914b-a42fc1e8a554")
+
+
+def native_source_id(organization_id: UUID, binding: NativeSourceBinding) -> UUID:
+ """Organization, owner and dataset define identity; collection is immutable configuration."""
+ name = json.dumps(
+ [str(organization_id), binding.owner_id, binding.dataset],
+ ensure_ascii=False,
+ separators=(",", ":"),
+ )
+ return uuid5(_SOURCE_NAMESPACE_V1, name)
+
+
+def native_sync_id(source_id: UUID) -> UUID:
+ """One sync owns each native source's retained records and writer fence."""
+ return uuid5(_SYNC_NAMESPACE_V1, str(source_id))
+
+
+class EnsureNativeSource(NativeSourceBinding):
+ """The backend attests the owner; no provider or credential selectors are accepted."""
+
+ collection: str = Field(min_length=1, max_length=255)
+
+ def binding(self) -> NativeSourceBinding:
+ """Exclude configuration from stable source identity."""
+ return NativeSourceBinding(owner_id=self.owner_id, dataset=self.dataset)
+
+
+class NativeSource(NativeModel):
+ """Provisioned identifiers, never proof of completed import or indexed coverage."""
+
+ source_connection_id: UUID
+ sync_id: UUID
+ organization_id: UUID
+ binding: NativeSourceBinding
+ collection: str
+ available: bool
diff --git a/backend/airweave/domains/native_ingestion/source_service.py b/backend/airweave/domains/native_ingestion/source_service.py
new file mode 100644
index 0000000000..9d4e35c3b7
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/source_service.py
@@ -0,0 +1,29 @@
+"""Own atomic native binding transactions independently of provider provisioning."""
+
+from uuid import UUID
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.native_ingestion.source_models import EnsureNativeSource, NativeSource
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+
+
+class NativeSources:
+ """Ensure/read only; source availability is not proof of imported content."""
+
+ def __init__(self, store: NativeSourceStore):
+ """Inject the single persistence boundary."""
+ self.store = store
+
+ async def ensure(
+ self, db: AsyncSession, organization_id: UUID, request: EnsureNativeSource
+ ) -> NativeSource:
+ """Publish the complete source/sync binding or roll back both."""
+ async with UnitOfWork(db):
+ return (await self.store.ensure(db, organization_id, request)).response()
+
+ async def get(self, db: AsyncSession, organization_id: UUID, source_id: UUID) -> NativeSource:
+ """Read a currently valid binding without opening a provider connection."""
+ async with UnitOfWork(db):
+ return (await self.store.require(db, organization_id, source_id)).response()
diff --git a/backend/airweave/domains/native_ingestion/source_store.py b/backend/airweave/domains/native_ingestion/source_store.py
new file mode 100644
index 0000000000..baf18b3de4
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/source_store.py
@@ -0,0 +1,160 @@
+"""Primary-key serialized native binding; no JSON discovery or external side effects."""
+
+from uuid import UUID
+
+from fastapi import HTTPException
+from pydantic import BaseModel, ConfigDict, ValidationError
+from sqlalchemy import select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.native_ingestion.models import NativeSourceBinding
+from airweave.domains.native_ingestion.source_models import (
+ EnsureNativeSource,
+ NativeSource,
+ native_source_id,
+ native_sync_id,
+)
+from airweave.models.collection import Collection
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+
+
+class LockedNativeSource(BaseModel):
+ """Internal ORM state held under Sync then SourceConnection row locks."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", frozen=True)
+ sync: Sync
+ source: SourceConnection
+ binding: NativeSourceBinding
+
+ def response(self) -> NativeSource:
+ """Expose only stable identifiers and typed binding, not ORM/auth data."""
+ return NativeSource(
+ source_connection_id=self.source.id,
+ sync_id=self.sync.id,
+ organization_id=self.sync.organization_id,
+ binding=self.binding,
+ collection=self.source.readable_collection_id,
+ available=self.source.is_authenticated,
+ )
+
+
+class NativeSourceStore:
+ """All methods run within their caller's transaction; locks live through commit."""
+
+ async def ensure(
+ self, db: AsyncSession, organization_id: UUID, request: EnsureNativeSource
+ ) -> LockedNativeSource:
+ """Concurrent exact ensures serialize on the deterministic Sync primary key."""
+ collection = await db.scalar(
+ select(Collection.id).where(
+ Collection.organization_id == organization_id,
+ Collection.readable_id == request.collection,
+ )
+ )
+ if collection is None:
+ raise HTTPException(404, "Native collection not found")
+ binding = request.binding()
+ source_id = native_source_id(organization_id, binding)
+ sync_id = native_sync_id(source_id)
+ inserted = await db.scalar(
+ insert(Sync)
+ .values(
+ id=sync_id,
+ organization_id=organization_id,
+ name=f"Almanac {binding.dataset}",
+ status="ACTIVE",
+ sync_type="full",
+ )
+ .on_conflict_do_nothing(index_elements=[Sync.id])
+ .returning(Sync.id)
+ )
+ # The unique insert waits for any concurrent ensure before this lock/read.
+ sync = await db.scalar(
+ select(Sync)
+ .where(
+ Sync.id == sync_id,
+ Sync.organization_id == organization_id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if sync is None:
+ raise HTTPException(409, "Native source identity conflicts with existing state")
+ if inserted is not None:
+ existing = await db.scalar(
+ select(SourceConnection.id).where(SourceConnection.id == source_id)
+ )
+ if existing is not None:
+ raise HTTPException(409, "Native source identity conflicts with existing state")
+ db.add(
+ SourceConnection(
+ id=source_id,
+ organization_id=organization_id,
+ sync_id=sync_id,
+ name=f"Almanac {binding.dataset}",
+ short_name="almanac",
+ config_fields=binding.model_dump(mode="json"),
+ readable_collection_id=request.collection,
+ is_authenticated=True,
+ )
+ )
+ await db.flush()
+ locked = await self.require(db, organization_id, source_id, missing_status=409)
+ if locked.binding != binding or locked.source.readable_collection_id != request.collection:
+ raise HTTPException(409, "Native source binding is immutable")
+ return locked
+
+ async def require(
+ self,
+ db: AsyncSession,
+ organization_id: UUID,
+ source_id: UUID,
+ *,
+ missing_status: int = 404,
+ ) -> LockedNativeSource:
+ """Lock Sync then source and fail closed on foreign or unexplained partial state."""
+ sync = await db.scalar(
+ select(Sync)
+ .where(
+ Sync.id == native_sync_id(source_id),
+ Sync.organization_id == organization_id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ source = await db.scalar(
+ select(SourceConnection)
+ .where(
+ SourceConnection.id == source_id,
+ SourceConnection.organization_id == organization_id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if sync is None or source is None:
+ raise HTTPException(missing_status, "Native source not found or incomplete")
+ try:
+ binding = NativeSourceBinding.model_validate(source.config_fields)
+ except ValidationError as error:
+ raise HTTPException(409, "Native source binding is malformed") from error
+ if (
+ source.short_name != "almanac"
+ or source.sync_id != sync.id
+ or native_source_id(organization_id, binding) != source_id
+ or source.readable_auth_provider_id is not None
+ or source.auth_provider_config is not None
+ or source.connection_id is not None
+ or source.connection_init_session_id is not None
+ ):
+ raise HTTPException(409, "Native source identity conflicts with existing state")
+ collection = await db.scalar(
+ select(Collection.id).where(
+ Collection.readable_id == source.readable_collection_id,
+ Collection.organization_id == organization_id,
+ )
+ )
+ if collection is None:
+ raise HTTPException(409, "Native source collection is unavailable")
+ return LockedNativeSource(sync=sync, source=source, binding=binding)
diff --git a/backend/airweave/domains/native_ingestion/store.py b/backend/airweave/domains/native_ingestion/store.py
new file mode 100644
index 0000000000..02062a1476
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/store.py
@@ -0,0 +1,291 @@
+"""Native admission composed with the existing canonical writer transaction."""
+
+from uuid import UUID
+
+from pydantic import ValidationError
+from sqlalchemy import select, tuple_, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.models import CaptureResult
+from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord
+from airweave.domains.entities.canonical.scan_models import CommitScanPage, ScanResult
+from airweave.domains.entities.canonical.scan_store import CanonicalScanStore
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, content_is_available
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.models import (
+ IngestNativeBatch,
+ IngestNativePage,
+ NativeAdmission,
+ NativeSnapshot,
+ NativeSourceBinding,
+ NativeVersion,
+)
+from airweave.models.entity import Entity
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+
+
+def compare_versions(previous: NativeVersion, incoming: NativeVersion) -> int:
+ """Partial ordering: reject crossed session revisions instead of guessing."""
+ if previous.kind != incoming.kind:
+ raise NativeAdmissionError("Native version domain changed")
+ old = (previous.revision,)
+ new = (incoming.revision,)
+ if previous.kind == "session" and incoming.kind == "session":
+ old += (previous.content_revision,)
+ new += (incoming.content_revision,)
+ lower = any(left < right for left, right in zip(new, old, strict=True))
+ higher = any(left > right for left, right in zip(new, old, strict=True))
+ if lower and higher:
+ raise NativeAdmissionError("Native version components are incomparable")
+ return -1 if lower else 1 if higher else 0
+
+
+def _admit(entity: Entity, snapshot: NativeSnapshot, *, renew: bool = False) -> bool:
+ """Validate retained version state; False acknowledges an exact retained retry."""
+ try:
+ previous = NativeSnapshot.model_validate(entity.source_payload)
+ except ValidationError as error:
+ raise NativeAdmissionError("Retained native version state is malformed") from error
+ if previous.identity != snapshot.identity or previous.owner_id != snapshot.owner_id:
+ raise NativeAdmissionError("Retained native identity does not match its source")
+ comparison = compare_versions(previous.version, snapshot.version)
+ if comparison < 0:
+ raise NativeAdmissionError("Native snapshot version is stale")
+ if comparison == 0:
+ if previous != snapshot:
+ raise NativeAdmissionError("Native snapshot version has conflicting content")
+ # Ordinary retries acknowledge retention. Only explicit access CAS renews.
+ return renew
+ if not renew and entity.removal_reason in ("access_revoked", "scope_removed", "absent"):
+ raise NativeAdmissionError("Native source visibility requires explicit renewal")
+ return True
+
+
+def _validate_message_parent(
+ snapshot: NativeSnapshot,
+ submitted: dict[tuple[str, str], NativeSnapshot],
+ retained: dict[tuple[str, str], Entity],
+) -> None:
+ """A late transcript page cannot join a newer session snapshot."""
+ if snapshot.identity.record_type != "message":
+ return
+ parent = snapshot.parent
+ assert parent is not None # NativeSnapshot validates message identity before admission.
+ key = (parent.record_type, parent.entity_key)
+ attested = submitted.get(key)
+ if attested is None:
+ entity = retained.get(key)
+ try:
+ attested = NativeSnapshot.model_validate(entity.source_payload if entity else None)
+ except ValidationError as error:
+ raise NativeAdmissionError("Message requires an attested native session") from error
+ if (
+ attested.identity != parent
+ or attested.owner_id != snapshot.owner_id
+ or (snapshot.operation == "upsert" and attested.operation != "upsert")
+ or attested.version != snapshot.version
+ ):
+ raise NativeAdmissionError("Message version must match its attested native session")
+
+
+def _observation(snapshot: NativeSnapshot, request: IngestNativeBatch) -> CaptureRecord:
+ return CaptureRecord(
+ identity=snapshot.identity,
+ parent=snapshot.parent,
+ payload=snapshot.model_dump(mode="json"),
+ kind=snapshot.operation,
+ removal_reason="provider_deleted" if snapshot.operation == "delete" else None,
+ observed_at=request.observed_at,
+ source_created_at=snapshot.source_created_at,
+ source_updated_at=snapshot.source_updated_at,
+ )
+
+
+def _admit_snapshots(
+ snapshots: tuple[NativeSnapshot, ...],
+ binding: NativeSourceBinding,
+ submitted: dict[tuple[str, str], NativeSnapshot],
+ retained: dict[tuple[str, str], Entity],
+ request: IngestNativeBatch,
+ *,
+ sightings: bool,
+ revalidated_ids: tuple[UUID, ...] = (),
+) -> NativeAdmission:
+ """Apply version and membership rules to one preloaded bounded batch."""
+ accepted = []
+ unchanged = []
+ for snapshot in snapshots:
+ if sightings and snapshot.operation != "upsert":
+ raise NativeAdmissionError("Snapshot pages enumerate active originals, not tombstones")
+ dataset = "knowledge" if snapshot.identity.record_type == "knowledge" else "sessions"
+ if snapshot.owner_id != binding.owner_id or dataset != binding.dataset:
+ raise NativeAdmissionError("Snapshot does not belong to the bound native source")
+ entity = retained.get((snapshot.identity.record_type, snapshot.identity.entity_key))
+ if entity is not None and not _admit(entity, snapshot, renew=entity.id in revalidated_ids):
+ if sightings:
+ _validate_message_parent(snapshot, submitted, retained)
+ if entity.deleted_at is not None:
+ raise NativeAdmissionError("A new sweep requires active original sightings")
+ unchanged.append(entity.id)
+ continue
+ _validate_message_parent(snapshot, submitted, retained)
+ accepted.append(_observation(snapshot, request))
+ return NativeAdmission(records=tuple(accepted), unchanged_ids=tuple(unchanged))
+
+
+class NativeIngestionStore:
+ """Compose canonical internals inside the service's single UnitOfWork.
+
+ `_fenced_sync` acquires the existing writer lock and validates the active run.
+ `_capture_locked` persists admitted snapshots under that same lock. Neither
+ method commits; duplicating either would split version checks from capture.
+ """
+
+ def __init__(self, canonical: CanonicalRecordStore):
+ """Reuse canonical locking, visibility and revision publication."""
+ self.canonical = canonical
+
+ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureResult:
+ """Run inside the caller's UnitOfWork; never commit independently."""
+ sync = await self.canonical._fenced_sync(db, request.fence)
+ admitted = await self.admit_locked(db, sync, request)
+ result = await self.canonical._capture_locked(
+ db, sync, CaptureBatch(fence=request.fence, records=admitted.records)
+ )
+ return result.model_copy(
+ update={"unchanged": result.unchanged + len(admitted.unchanged_ids)}
+ )
+
+ async def admit_locked(
+ self,
+ db: AsyncSession,
+ sync: Sync,
+ request: IngestNativeBatch,
+ *,
+ sightings: bool = False,
+ revalidated_ids: tuple[UUID, ...] = (),
+ ) -> NativeAdmission:
+ """Caller must hold the canonical writer lock through capture and commit."""
+ if (
+ sync.id != request.fence.sync_id
+ or sync.organization_id != request.fence.organization_id
+ ):
+ raise NativeAdmissionError("Native admission lock belongs to another source")
+ source = (
+ await db.execute(
+ select(SourceConnection)
+ .where(
+ SourceConnection.organization_id == request.fence.organization_id,
+ SourceConnection.sync_id == sync.id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ ).scalar_one_or_none()
+ if source is None or source.short_name != "almanac":
+ raise NativeAdmissionError("Writer is not bound to an Almanac native source")
+ try:
+ binding = NativeSourceBinding.model_validate(source.config_fields)
+ except ValidationError as error:
+ raise NativeAdmissionError("Native source binding is malformed") from error
+ submitted = {
+ (item.identity.record_type, item.identity.entity_key): item
+ for item in request.snapshots
+ }
+ keys = set(submitted)
+ keys.update(
+ (item.parent.record_type, item.parent.entity_key)
+ for item in request.snapshots
+ if item.parent is not None
+ )
+ # The sync writer lock already serializes mutations. Fetch exact identities
+ # and parent attestations once rather than adding a round trip per record.
+ rows = await db.scalars(
+ select(Entity)
+ .where(
+ Entity.organization_id == request.fence.organization_id,
+ Entity.sync_id == sync.id,
+ tuple_(Entity.entity_definition_short_name, Entity.entity_id).in_(sorted(keys)),
+ )
+ .execution_options(populate_existing=True)
+ )
+ retained = {(row.entity_definition_short_name, row.entity_id): row for row in rows}
+ admitted = _admit_snapshots(
+ request.snapshots,
+ binding,
+ submitted,
+ retained,
+ request,
+ sightings=sightings,
+ revalidated_ids=revalidated_ids,
+ )
+ # A newer original cannot silently rebind a retained child after ancestor
+ # withdrawal. Explicit renewal attests both its local revision and parent epoch.
+ protected = tuple(
+ row.id
+ for item in admitted.records
+ if item.kind == "upsert"
+ and (row := retained.get((item.identity.record_type, item.identity.entity_key)))
+ is not None
+ and row.id not in revalidated_ids
+ )
+ if protected:
+ visible = set(
+ await db.scalars(
+ select(Entity.id).where(Entity.id.in_(protected), content_is_available())
+ )
+ )
+ if visible != set(protected):
+ raise NativeAdmissionError("Native source visibility requires explicit renewal")
+ unchanged = admitted.unchanged_ids
+ if sightings and unchanged:
+ available = set(
+ await db.scalars(
+ select(Entity.id).where(Entity.id.in_(unchanged), content_is_available())
+ )
+ )
+ if available != set(unchanged):
+ raise NativeAdmissionError("Withdrawn native originals cannot be marked seen")
+ return admitted
+
+ async def page(self, db: AsyncSession, request: IngestNativePage) -> ScanResult:
+ """Use canonical scan CAS and one capture callback within the caller's transaction."""
+ observations = tuple(_observation(item, request) for item in request.snapshots)
+
+ async def capture_page(
+ session: AsyncSession, sync: Sync, batch: CaptureBatch, *, seen_id: UUID
+ ) -> CaptureResult:
+ # CanonicalScanStore has already validated these exact observations/scope.
+ if batch.records != observations or batch.fence != request.fence:
+ raise NativeAdmissionError("Native scan capture command changed")
+ admitted = await self.admit_locked(session, sync, request, sightings=True)
+ result = await self.canonical._capture_locked(
+ session,
+ sync,
+ CaptureBatch(fence=batch.fence, records=admitted.records),
+ seen_id=seen_id,
+ )
+ if admitted.unchanged_ids:
+ await session.execute(
+ update(Entity)
+ .where(Entity.id.in_(admitted.unchanged_ids))
+ .values(last_seen_run_id=seen_id)
+ )
+ return result.model_copy(
+ update={"unchanged": result.unchanged + len(admitted.unchanged_ids)}
+ )
+
+ return await CanonicalScanStore(self.canonical).page(
+ db,
+ CommitScanPage(
+ fence=request.fence,
+ scope=request.scope,
+ cycle_id=request.cycle_id,
+ expected=request.expected,
+ records=observations,
+ continuation=request.continuation,
+ final=request.final,
+ ),
+ capture_page=capture_page,
+ )
diff --git a/backend/airweave/domains/native_ingestion/tests/__init__.py b/backend/airweave/domains/native_ingestion/tests/__init__.py
new file mode 100644
index 0000000000..f633ca572c
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/__init__.py
@@ -0,0 +1 @@
+"""Native ingestion tests use package-qualified names alongside canonical tests."""
diff --git a/backend/airweave/domains/native_ingestion/tests/conftest.py b/backend/airweave/domains/native_ingestion/tests/conftest.py
new file mode 100644
index 0000000000..e7b95331ff
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/conftest.py
@@ -0,0 +1,5 @@
+"""Reuse isolated real canonical migrations, never an application database."""
+
+from airweave.domains.entities.canonical.tests.conftest import database, source
+
+__all__ = ["database", "source"]
diff --git a/backend/airweave/domains/native_ingestion/tests/test_access_api.py b/backend/airweave/domains/native_ingestion/tests/test_access_api.py
new file mode 100644
index 0000000000..7ac8bacfa6
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_access_api.py
@@ -0,0 +1,209 @@
+"""Explicit access lifecycle using HTTP and isolated PostgreSQL, not provider mutations."""
+
+# Imported fixture is resolved through pytest parameters.
+# ruff: noqa: F811
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.models import IngestNativeBatch, SessionVersion
+from airweave.domains.native_ingestion.service import NativeIngestionService
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.domains.native_ingestion.tests.test_import_api import native_api # noqa: F401
+from airweave.domains.native_ingestion.tests.test_ingestion import NOW, snapshot
+from airweave.models.entity import Entity
+from airweave.models.source_connection import SourceConnection
+
+
+async def capture(client, base, original):
+ scope = {"record_type": original.identity.record_type}
+ if original.parent:
+ scope.update(
+ container_id=original.identity.container_id,
+ parent=original.parent.model_dump(mode="json"),
+ )
+ opened = await client.put(base + "/scopes", json={"scope": scope})
+ assert opened.status_code == 200, opened.text
+ result = await client.put(
+ base + "/pages",
+ json={
+ "page_id": str(uuid4()),
+ "scope": scope,
+ "expected": opened.json()["version"],
+ "snapshots": [original.model_dump(mode="json")],
+ "final": True,
+ },
+ )
+ assert result.status_code == 200, result.text
+ completed = await client.post(
+ base + "/scopes/reconcile",
+ json={
+ "scope": scope,
+ "expected": result.json()["version"],
+ },
+ )
+ assert completed.status_code == 200, completed.text
+
+
+async def row_id(database, source, kind):
+ async with database() as db:
+ return await db.scalar(
+ select(Entity.id).where(
+ Entity.sync_id == source.sync_id,
+ Entity.entity_definition_short_name == kind,
+ )
+ )
+
+
+async def test_access_cas_preserves_original_and_rejects_stale_restore(database, native_api):
+ client, ctx, source, logs = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/access"
+ await client.put(base, json={"snapshot_id": "one", "coverage": "bounded"})
+ original = snapshot(owner_id="owner", revision=3)
+ await capture(client, base, original)
+ record_id = await row_id(database, source, "knowledge")
+ url = base + f"/records/{record_id}/access"
+ initial = await client.get(url)
+ assert initial.json()["available"] and initial.json()["revision"] == 1
+ assert "original" not in initial.text and "payload" not in initial.text
+ withdraw = {"action": "withdraw", "expected_revision": 1, "reason": "access_revoked"}
+ revoked = await client.post(url, json=withdraw)
+ assert revoked.status_code == 200, revoked.text
+ assert revoked.json()["revision"] == 2 and not revoked.json()["available"]
+ # Lost response recovery reads state; old CAS never replays across newer observations.
+ assert (await client.post(url, json=withdraw)).status_code == 409
+ assert (await client.get(url)).json() == revoked.json()
+ async with database() as db:
+ retained = await db.get(Entity, record_id)
+ assert retained.source_payload == original.model_dump(mode="json")
+ renew = {
+ "action": "renew",
+ "expected_revision": 2,
+ "snapshot": original.model_dump(mode="json"),
+ }
+ for changes in (
+ {"owner_id": "wrong"},
+ {"original": {"body": "same-version-conflict"}},
+ {"version": {"kind": "record", "revision": 2}},
+ ):
+ rejected = await client.post(url, json=renew | {"snapshot": renew["snapshot"] | changes})
+ assert rejected.status_code == 409, rejected.text
+ assert (await client.get(url)).json() == revoked.json()
+ restored = await client.post(url, json=renew)
+ assert restored.status_code == 200, restored.text
+ assert restored.json()["available"] and restored.json()["revision"] == 3
+ revoked_again = await client.post(url, json=withdraw | {"expected_revision": 3})
+ assert revoked_again.json()["revision"] == 4
+ assert (await client.post(url, json=renew)).status_code == 409
+ assert not (await client.get(url)).json()["available"]
+ ctx.is_api_key_auth = False
+ assert (await client.get(url)).status_code == 403
+ ctx.is_api_key_auth = True
+ ctx.organization.id = uuid4()
+ assert (await client.get(url)).status_code == 404
+ assert (await client.post(url, json=renew)).status_code == 404
+ assert "same-version-conflict" not in str(logs)
+
+
+@pytest.mark.parametrize("native_api", ["sessions"], indirect=True)
+async def test_parent_restore_requires_fresh_child_attestation_and_active_import(
+ database, native_api
+):
+ client, _, source, _ = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/access"
+ await client.put(base, json={"snapshot_id": "one", "coverage": "bounded"})
+ parent = snapshot(
+ owner_id="owner",
+ identity=RecordIdentity(record_type="session", native_id="s"),
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ child = snapshot(
+ owner_id="owner",
+ identity=RecordIdentity(record_type="message", native_id="m", container_id="s"),
+ parent=parent.identity,
+ version=parent.version,
+ )
+ await capture(client, base, parent)
+ await capture(client, base, child)
+ parent_url = base + f"/records/{await row_id(database, source, 'session')}/access"
+ child_url = base + f"/records/{await row_id(database, source, 'message')}/access"
+ withdrawn = await client.post(
+ parent_url,
+ json={
+ "action": "withdraw",
+ "reason": "scope_removed",
+ "expected_revision": 1,
+ },
+ )
+ assert withdrawn.status_code == 200, withdrawn.text
+ assert not (await client.get(child_url)).json()["available"]
+ child_renew = {
+ "expected_parent_epoch": 1,
+ "action": "renew",
+ "expected_revision": 1,
+ "snapshot": child.model_dump(mode="json"),
+ }
+ assert (await client.post(child_url, json=child_renew)).status_code == 409
+ restored = await client.post(
+ parent_url,
+ json={
+ "action": "renew",
+ "expected_revision": 2,
+ "snapshot": parent.model_dump(mode="json"),
+ },
+ )
+ assert restored.status_code == 200, restored.text
+ assert not (await client.get(child_url)).json()["available"]
+ # The request prepared before revocation must not cross the parent epoch change.
+ assert (await client.post(child_url, json=child_renew)).status_code == 409
+ parent_new = parent.model_copy(
+ update={"version": SessionVersion(revision=2, content_revision=1)}
+ )
+ parent_update = await client.post(
+ parent_url,
+ json={
+ "action": "renew",
+ "expected_revision": 3,
+ "snapshot": parent_new.model_dump(mode="json"),
+ },
+ )
+ assert parent_update.status_code == 200, parent_update.text
+ child_new = child.model_copy(update={"version": parent_new.version})
+ canonical = CanonicalRecordStore()
+ imports = NativeImportStore(NativeSourceStore(), canonical)
+ async with database() as db, UnitOfWork(db):
+ saved = await imports.active(
+ db, source.organization_id, source.source_connection_id, "access"
+ )
+ async with database() as db:
+ with pytest.raises(NativeAdmissionError, match="renewal"):
+ await NativeIngestionService(NativeIngestionStore(canonical)).ingest(
+ db,
+ IngestNativeBatch(fence=saved.fence, observed_at=NOW, snapshots=(child_new,)),
+ )
+ current = (await client.get(child_url)).json()
+ assert current["revision"] == 1 and current["parent_visibility_epoch"] == 2
+ child_renew = child_renew | {
+ "expected_parent_epoch": current["parent_visibility_epoch"],
+ "snapshot": child_new.model_dump(mode="json"),
+ }
+ refreshed = await client.post(child_url, json=child_renew)
+ assert refreshed.status_code == 200, refreshed.text
+ assert refreshed.json()["available"] and refreshed.json()["revision"] == 2
+ await client.post(base + "/cancel")
+ assert (await client.get(child_url)).status_code == 200
+ assert (
+ await client.post(child_url, json=child_renew | {"expected_revision": 2})
+ ).status_code == 409
+ async with database() as db:
+ bound = await db.get(SourceConnection, source.source_connection_id)
+ bound.is_authenticated = False
+ await db.commit()
+ assert not (await client.get(child_url)).json()["available"]
diff --git a/backend/airweave/domains/native_ingestion/tests/test_import_api.py b/backend/airweave/domains/native_ingestion/tests/test_import_api.py
new file mode 100644
index 0000000000..0ebc3cae24
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_import_api.py
@@ -0,0 +1,143 @@
+"""Native HTTP contract with actual PostgreSQL transactions and safe diagnostics."""
+
+from types import SimpleNamespace
+from uuid import uuid4
+
+import httpx
+import pytest
+from fastapi import FastAPI
+from fastapi.exceptions import RequestValidationError
+from sqlalchemy import func, select
+
+from airweave.api import deps, middleware
+from airweave.api.v1.endpoints.native_imports import router
+from airweave.db.session import get_db
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.import_service import NativeImports
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.source_models import EnsureNativeSource
+from airweave.domains.native_ingestion.source_service import NativeSources
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.models import Collection, Organization, SyncJob
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+@pytest.fixture(params=["knowledge"])
+async def native_api(database, monkeypatch, request):
+ org = uuid4()
+ async with database() as db:
+ db.add(Organization(id=org, name="Native API"))
+ metadata = VectorDbDeploymentMetadata(
+ dense_embedder="test", sparse_embedder="test", embedding_dimensions=3
+ )
+ db.add(metadata)
+ await db.flush()
+ db.add(
+ Collection(
+ name="API",
+ readable_id="native-api",
+ organization_id=org,
+ vector_db_deployment_metadata_id=metadata.id,
+ )
+ )
+ await db.commit()
+ sources = NativeSourceStore()
+ async with database() as db:
+ source = await NativeSources(sources).ensure(
+ db,
+ org,
+ EnsureNativeSource(owner_id="owner", dataset=request.param, collection="native-api"),
+ )
+ service = NativeImports(NativeImportStore(sources, CanonicalRecordStore()))
+ ctx = SimpleNamespace(is_api_key_auth=True, organization=SimpleNamespace(id=org))
+ app = FastAPI()
+ app.include_router(router, prefix="/native/sources")
+ app.exception_handler(RequestValidationError)(middleware.validation_exception_handler)
+ app.middleware("http")(middleware.log_requests)
+ logs = []
+ monkeypatch.setattr(
+ middleware,
+ "logger",
+ SimpleNamespace(info=logs.append, error=logs.append, warning=logs.append),
+ )
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: ctx
+ app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(native_imports=service)
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app), base_url="http://test"
+ ) as client:
+ yield client, ctx, source, logs
+
+
+async def test_import_http_retry_conflict_and_key_validation(database, native_api):
+ client, _, source, logs = native_api
+ path = f"/native/sources/{source.source_connection_id}/imports/private-request-key"
+ body = {"snapshot_id": "private-snapshot-value", "coverage": "bounded"}
+ first = await client.put(path, json=body)
+ assert first.status_code == 200, first.text
+ assert (await client.put(path, json=body)).json() == first.json()
+ assert (await client.get(path)).json() == first.json()
+ assert "fence" not in first.json()
+ conflict = await client.put(path, json={**body, "coverage": "complete"})
+ assert conflict.status_code == 409
+ assert (await client.put(path + "-other", json=body)).status_code == 409
+ assert (await client.get(path + "-missing")).status_code == 404
+ assert (await client.put(path + "x" * 129, json=body)).status_code == 422
+ assert (
+ await client.put(path, json={**body, "coverage": "private-invalid-value"})
+ ).status_code == 422
+ assert "private-request-key" not in str(logs)
+ assert "private-snapshot-value" not in str(logs)
+ assert "private-invalid-value" not in str(logs)
+ assert "{request_key}" in str(logs)
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+
+
+async def test_import_http_rejects_session_and_foreign_organization(native_api):
+ client, ctx, source, _ = native_api
+ path = f"/native/sources/{source.source_connection_id}/imports/request-key"
+ body = {"snapshot_id": "snapshot", "coverage": "bounded"}
+ ctx.is_api_key_auth = False
+ assert (await client.put(path, json=body)).status_code == 403
+ assert (await client.get(path)).status_code == 403
+ ctx.is_api_key_auth = True
+ ctx.organization.id = uuid4()
+ assert (await client.put(path, json=body)).status_code == 404
+ assert (await client.get(path)).status_code == 404
+
+
+async def test_native_analytics_and_exception_logs_keep_only_safe_route(monkeypatch):
+ from starlette.requests import Request
+ from starlette.responses import Response
+
+ request = Request(
+ {
+ "type": "http",
+ "method": "PUT",
+ "scheme": "http",
+ "headers": [],
+ "path": "/native/sources/source/imports/private-request-key",
+ "query_string": b"payload=private-query",
+ "path_params": {"source_id": "source", "request_key": "private-request-key"},
+ "route": SimpleNamespace(path="/native/sources/{source_id}/imports/{request_key}"),
+ }
+ )
+ captured = []
+ monkeypatch.setattr(middleware.analytics, "track_event", lambda *a, **k: captured.append(a))
+ monkeypatch.setattr(middleware, "logger", SimpleNamespace(error=captured.append))
+ await middleware._track_api_call_async(None, Response(status_code=200), 12.5, request)
+
+ async def failed(_):
+ raise ValueError("private-payload-content")
+
+ response = await middleware.exception_logging_middleware(request, failed)
+ assert response.status_code == 500
+ assert "private-" not in str(captured)
+ assert "{request_key}" in str(captured)
+ assert "private-" not in response.body.decode()
diff --git a/backend/airweave/domains/native_ingestion/tests/test_imports.py b/backend/airweave/domains/native_ingestion/tests/test_imports.py
new file mode 100644
index 0000000000..d7d73c7b7e
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_imports.py
@@ -0,0 +1,215 @@
+"""Real PostgreSQL import request recovery, concurrent creation and rollback."""
+
+import asyncio
+import hashlib
+import json
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import func, select
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, WriterBusy
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_models import StartNativeImport
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.source_models import EnsureNativeSource
+from airweave.domains.native_ingestion.source_service import NativeSources
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.models.collection import Collection
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+from airweave.models.sync_cursor import SyncCursor
+from airweave.models.sync_job import SyncJob
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+@pytest.fixture(params=["knowledge", "sessions"])
+async def native(database, source, request):
+ _, fence = source
+ async with database() as db:
+ deployment = VectorDbDeploymentMetadata(
+ dense_embedder="fake", embedding_dimensions=3, sparse_embedder="fake"
+ )
+ db.add(deployment)
+ await db.flush()
+ db.add(
+ Collection(
+ organization_id=fence.organization_id,
+ name="Native import",
+ readable_id="native",
+ vector_db_deployment_metadata_id=deployment.id,
+ )
+ )
+ await db.commit()
+ async with database() as db:
+ return await NativeSources(NativeSourceStore()).ensure(
+ db,
+ fence.organization_id,
+ EnsureNativeSource(owner_id="owner", dataset=request.param, collection="native"),
+ )
+
+
+async def start(database, native, key="request-one", request=None):
+ store = NativeImportStore(NativeSourceStore(), CanonicalRecordStore())
+ async with database() as db, UnitOfWork(db):
+ return await store.start(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ key,
+ request or StartNativeImport(snapshot_id="snapshot-one", coverage="bounded"),
+ )
+
+
+async def test_concurrent_start_recovers_one_writer_and_rejects_changed_key_body(database, native):
+ first, duplicate = await asyncio.gather(start(database, native), start(database, native))
+ assert first == duplicate
+ async with database() as db:
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == native.sync_id)
+ )
+ == 1
+ )
+ sync = await db.get(Sync, native.sync_id)
+ assert sync.writer_epoch == 1
+ with pytest.raises(NativeAdmissionError, match="conflicting intent"):
+ await start(
+ database, native, request=StartNativeImport(snapshot_id="other", coverage="complete")
+ )
+ with pytest.raises(WriterBusy):
+ await start(database, native, key="other-key")
+ async with database() as db:
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == native.sync_id)
+ )
+ == 1
+ )
+
+
+async def test_terminal_retry_never_reactivates_and_new_request_replaces_partial_cycle(
+ database, native
+):
+ first = await start(database, native)
+ async with database() as db:
+ job = await db.get(SyncJob, first.import_id)
+ job.status = "cancelled"
+ await db.commit()
+ duplicate = await start(database, native)
+ assert duplicate.status == "cancelled" and duplicate.cycle_id == first.cycle_id
+ async with database() as db:
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 1
+ next_import = await start(database, native, key="new-request")
+ assert next_import.cycle_id != first.cycle_id
+ async with database() as db:
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 2
+ assert (await start(database, native)).status == "cancelled"
+ async with database() as db:
+ source = await db.get(SourceConnection, native.source_connection_id)
+ source.is_authenticated = False
+ await db.commit()
+ assert (await start(database, native)).status == "cancelled"
+ with pytest.raises(NativeAdmissionError, match="unavailable"):
+ await start(database, native, key="after-withdrawal")
+ async with database() as db:
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 2
+
+
+async def test_cycle_failure_rolls_back_job_and_fence(database, native, monkeypatch):
+ import airweave.domains.native_ingestion.import_store as module
+
+ async def fail(*args, **kwargs):
+ raise RuntimeError("cycle publication failed")
+
+ monkeypatch.setattr(module, "begin_cycle", fail)
+ with pytest.raises(RuntimeError, match="cycle publication"):
+ await start(database, native)
+ async with database() as db:
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 0
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == native.sync_id)
+ )
+ == 0
+ )
+ from fastapi import HTTPException
+
+ async with database() as db, UnitOfWork(db):
+ with pytest.raises(HTTPException) as error:
+ await NativeImportStore(NativeSourceStore(), CanonicalRecordStore()).read(
+ db, uuid4(), native.source_connection_id, "request-one"
+ )
+ assert error.value.status_code == 404
+
+
+async def test_transcript_coverage_is_sessions_only_and_part_of_retry_identity(database, native):
+ request = StartNativeImport(
+ snapshot_id="snapshot-one", coverage="bounded", transcript_coverage="complete"
+ )
+ if native.binding.dataset == "knowledge":
+ with pytest.raises(NativeAdmissionError, match="requires a sessions source"):
+ await start(database, native, request=request)
+ async with database() as db:
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 0
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(SyncJob)
+ .where(SyncJob.sync_id == native.sync_id)
+ )
+ == 0
+ )
+ return
+
+ initial = await start(database, native)
+ assert initial.request.transcript_coverage == "bounded"
+ with pytest.raises(NativeAdmissionError, match="conflicting intent"):
+ await start(database, native, request=request)
+ assert await start(database, native) == initial
+
+
+@pytest.mark.parametrize("native", ["sessions"], indirect=True)
+async def test_legacy_import_without_transcript_coverage_resumes_without_changing_plan(
+ database, native
+):
+ initial = await start(database, native)
+ async with database() as db:
+ job = await db.get(SyncJob, initial.import_id)
+ request = dict(job.sync_metadata["request"])
+ request.pop("transcript_coverage")
+ job.sync_metadata = {**job.sync_metadata, "request": request}
+ cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == native.sync_id))
+ cycle = cursor.cursor_data["canonical_cycle"]
+ plan = dict(cycle["source_plan"])
+ plan.pop("transcript_coverage")
+ fingerprint = hashlib.sha256(json.dumps(plan, sort_keys=True).encode()).hexdigest()
+ legacy_cycle = {
+ **cycle,
+ "source_plan": plan,
+ "configuration": {**cycle["configuration"], "fingerprint": fingerprint},
+ }
+ cursor.cursor_data = {**cursor.cursor_data, "canonical_cycle": legacy_cycle}
+ await db.commit()
+
+ assert await start(database, native) == initial
+ async with database() as db:
+ cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == native.sync_id))
+ assert cursor.cursor_data["canonical_cycle"] == legacy_cycle
+ assert (await db.get(Sync, native.sync_id)).writer_epoch == 1
+ job = await db.get(SyncJob, initial.import_id)
+ assert "transcript_coverage" not in job.sync_metadata["request"]
+
+
+async def test_legacy_cancellation_uses_actual_native_sync_identity(database, native, source):
+ from fastapi import HTTPException
+
+ from airweave.domains.owned_provisioning.guard import require_provider_sync
+
+ async with database() as db:
+ with pytest.raises(HTTPException) as error:
+ await require_provider_sync(db, native.sync_id, native.organization_id)
+ assert error.value.status_code == 409
+ _, provider = source
+ await require_provider_sync(db, provider.sync_id, provider.organization_id)
diff --git a/backend/airweave/domains/native_ingestion/tests/test_ingestion.py b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py
new file mode 100644
index 0000000000..cdf9ef88a2
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py
@@ -0,0 +1,296 @@
+"""Real transactions prove version admission, rollback and retained access state."""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.requests import (
+ CaptureBatch,
+ CaptureRecord,
+ RecordIdentity,
+ RemovedScope,
+)
+from airweave.domains.entities.canonical.store import (
+ CanonicalRecordStore,
+ CanonicalStoreError,
+ StaleWriter,
+ content_is_available,
+)
+from airweave.domains.entities.canonical.tests.helpers import bind_projection
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.models import (
+ IngestNativeBatch,
+ NativeSnapshot,
+ RecordVersion,
+ SessionVersion,
+)
+from airweave.domains.native_ingestion.service import NativeIngestionService
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.models.entity import Entity
+from airweave.models.entity_change import EntityChange
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+
+NOW = datetime(2026, 10, 1, tzinfo=timezone.utc)
+
+
+def snapshot(native_id="one", revision=1, **changes):
+ """Keep original Unicode JSON distinct from index text."""
+ fields = {
+ "owner_id": "owner-one",
+ "identity": RecordIdentity(record_type="knowledge", native_id=native_id),
+ "version": RecordVersion(revision=revision),
+ "original": {"body": "नमस्ते — original", "references": ["people/sam"]},
+ }
+ return NativeSnapshot(**(fields | changes))
+
+
+async def bind(database, fence, *, dataset="knowledge", provider="almanac"):
+ binding = await bind_projection(database, fence, provider)
+ async with database() as db:
+ source = await db.get(SourceConnection, binding.source_connection_id)
+ source.config_fields = {"owner_id": "owner-one", "dataset": dataset}
+ await db.commit()
+
+
+async def ingest(database, fence, *snapshots):
+ service = NativeIngestionService(NativeIngestionStore(CanonicalRecordStore()))
+ async with database() as db:
+ return await service.ingest(
+ db, IngestNativeBatch(fence=fence, observed_at=NOW, snapshots=snapshots)
+ )
+
+
+async def retained(database, fence):
+ async with database() as db:
+ return list((await db.scalars(select(Entity).where(Entity.sync_id == fence.sync_id))).all())
+
+
+async def test_atomic_versions_retry_tombstone_and_stale_resurrection(database, source):
+ _, fence = source
+ await bind(database, fence)
+ first = snapshot(revision=3)
+ result = await ingest(database, fence, first)
+ assert result.sequence == 1
+ assert (await ingest(database, fence, first)).unchanged == 1
+ for invalid in (snapshot(revision=2), snapshot(revision=3, original={"body": "conflict"})):
+ with pytest.raises(NativeAdmissionError):
+ await ingest(database, fence, snapshot("second"), invalid)
+ rows = await retained(database, fence)
+ assert len(rows) == 1 and rows[0].record_revision == 1
+ deleted = snapshot(revision=4, operation="delete", original={})
+ result = await ingest(database, fence, deleted)
+ assert result.sequence == 2 and result.changes[0].kind == "delete"
+ with pytest.raises(NativeAdmissionError, match="stale"):
+ await ingest(database, fence, first)
+ assert (await ingest(database, fence, deleted)).unchanged == 1
+ result = await ingest(database, fence, snapshot(revision=5))
+ assert result.sequence == 3 and result.changes[0].record.deleted_at is None
+
+
+async def test_session_versions_are_partial_order_and_original_payload_retained(database, source):
+ _, fence = source
+ await bind(database, fence, dataset="sessions")
+ item = snapshot(
+ identity=RecordIdentity(record_type="session", native_id="session-one"),
+ version=SessionVersion(revision=3, content_revision=7),
+ )
+ await ingest(database, fence, item)
+ with pytest.raises(NativeAdmissionError, match="incomparable"):
+ await ingest(
+ database,
+ fence,
+ item.model_copy(update={"version": SessionVersion(revision=4, content_revision=6)}),
+ )
+ higher = item.model_copy(update={"version": SessionVersion(revision=3, content_revision=8)})
+ result = await ingest(database, fence, higher)
+ assert result.changes[0].record.payload["original"] == item.original
+ assert result.changes[0].record.payload["representation"] == "snapshot"
+
+
+async def test_visibility_withdrawal_preserves_version_and_retry_does_not_restore(database, source):
+ canonical, fence = source
+ await bind(database, fence)
+ item = snapshot(revision=3)
+ await ingest(database, fence, item)
+ async with database() as db:
+ await canonical.remove_scope(
+ db,
+ fence,
+ RemovedScope(record_type="knowledge", removal_reason="access_revoked", observed_at=NOW),
+ )
+ assert (await ingest(database, fence, item)).unchanged == 1
+ row = (await retained(database, fence))[0]
+ assert row.removal_reason == "access_revoked"
+ assert row.source_payload["version"]["revision"] == 3
+ with pytest.raises(NativeAdmissionError, match="renewal"):
+ await ingest(database, fence, snapshot(revision=4))
+
+
+async def test_bound_source_owner_malformed_state_and_superseded_writer(database, source):
+ canonical, fence = source
+ with pytest.raises(NativeAdmissionError, match="not bound"):
+ await ingest(database, fence, snapshot())
+ await bind(database, fence)
+ with pytest.raises(NativeAdmissionError, match="bound native source"):
+ await ingest(database, fence, snapshot(owner_id="other-owner"))
+ async with database() as db:
+ await canonical.capture(
+ db,
+ CaptureBatch(
+ fence=fence,
+ records=(
+ CaptureRecord(
+ identity=snapshot().identity, payload={"original": {}}, observed_at=NOW
+ ),
+ ),
+ ),
+ )
+ with pytest.raises(NativeAdmissionError, match="malformed"):
+ await ingest(database, fence, snapshot())
+ async with database() as db:
+ await canonical.activate_writer(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ fence.job_id,
+ attempt_id=uuid4(),
+ attempt_number=2,
+ )
+ with pytest.raises(StaleWriter):
+ await ingest(database, fence, snapshot("second"))
+ assert len(await retained(database, fence)) == 1
+
+
+async def test_provider_binding_cannot_be_used_for_native_ingestion(database, source):
+ _, fence = source
+ await bind(database, fence, provider="gmail")
+ with pytest.raises(NativeAdmissionError, match="not bound"):
+ await ingest(database, fence, snapshot())
+
+
+async def test_capture_failure_rolls_back_batch_and_parent_withdrawal_hides_children(
+ database, source
+):
+ canonical, fence = source
+ await bind(database, fence, dataset="sessions")
+ parent = snapshot(
+ identity=RecordIdentity(record_type="session", native_id="session-one"),
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ missing = RecordIdentity(record_type="session", native_id="missing")
+ late_parent = parent.model_copy(update={"identity": missing})
+ child = snapshot(
+ identity=RecordIdentity(
+ record_type="message", native_id="message-one", container_id="missing"
+ ),
+ parent=missing,
+ version=parent.version,
+ original={"role": "user", "content": "original message"},
+ )
+ with pytest.raises(CanonicalStoreError):
+ # All pass admission; the child arrives before its parent in capture order.
+ await ingest(database, fence, parent, child, late_parent)
+ assert await retained(database, fence) == []
+ async with database() as db:
+ assert (
+ await db.scalar(select(Sync.observed_change_sequence).where(Sync.id == fence.sync_id))
+ == 0
+ )
+ assert (
+ await db.scalar(select(EntityChange.id).where(EntityChange.sync_id == fence.sync_id))
+ is None
+ )
+ child = child.model_copy(
+ update={
+ "parent": parent.identity,
+ "identity": RecordIdentity(
+ record_type="message", native_id="message-one", container_id="session-one"
+ ),
+ }
+ )
+ result = await ingest(database, fence, parent, child)
+ assert result.sequence == 2 and [change.sequence for change in result.changes] == [1, 2]
+ async with database() as db:
+ await canonical.remove_scope(
+ db,
+ fence,
+ RemovedScope(record_type="session", removal_reason="access_revoked", observed_at=NOW),
+ )
+ assert (await ingest(database, fence, child)).unchanged == 1
+ async with database() as db:
+ visible = await db.scalar(
+ select(content_is_available()).where(
+ Entity.sync_id == fence.sync_id, Entity.native_id == "message-one"
+ )
+ )
+ assert visible is False
+ higher = child.model_copy(update={"version": SessionVersion(revision=1, content_revision=2)})
+ with pytest.raises(CanonicalStoreError):
+ await ingest(database, fence, higher)
+ rows = await retained(database, fence)
+ message = next(row for row in rows if row.native_id == "message-one")
+ assert message.record_revision == 1
+
+
+async def test_late_unseen_message_cannot_join_newer_session(database, source):
+ _, fence = source
+ await bind(database, fence, dataset="sessions")
+ parent = snapshot(
+ identity=RecordIdentity(record_type="session", native_id="session-one"),
+ version=SessionVersion(revision=3, content_revision=10),
+ )
+ await ingest(database, fence, parent)
+ child = snapshot(
+ identity=RecordIdentity(
+ record_type="message", native_id="late", container_id="session-one"
+ ),
+ parent=parent.identity,
+ version=SessionVersion(revision=2, content_revision=8),
+ )
+ with pytest.raises(NativeAdmissionError, match="attested native session"):
+ await ingest(database, fence, child)
+ assert len(await retained(database, fence)) == 1
+ current = child.model_copy(update={"version": parent.version})
+ assert (await ingest(database, fence, current)).sequence == 2
+ advanced = parent.model_copy(
+ update={"version": SessionVersion(revision=3, content_revision=11)}
+ )
+ assert (await ingest(database, fence, advanced)).sequence == 3
+ assert (await ingest(database, fence, current)).unchanged == 1
+ unseen = current.model_copy(
+ update={
+ "identity": RecordIdentity(
+ record_type="message", native_id="unseen", container_id="session-one"
+ )
+ }
+ )
+ with pytest.raises(NativeAdmissionError, match="attested native session"):
+ await ingest(database, fence, unseen)
+ stale_delete = current.model_copy(update={"operation": "delete", "original": {}})
+ with pytest.raises(NativeAdmissionError, match="conflicting content"):
+ await ingest(database, fence, stale_delete)
+ # A later child version can still be stale relative to the parent.
+ newest = advanced.model_copy(
+ update={"version": SessionVersion(revision=3, content_revision=12)}
+ )
+ await ingest(database, fence, newest)
+ stale_delete = stale_delete.model_copy(update={"version": advanced.version})
+ with pytest.raises(NativeAdmissionError, match="attested native session"):
+ await ingest(database, fence, stale_delete)
+ current_delete = stale_delete.model_copy(update={"version": newest.version})
+ result = await ingest(database, fence, current_delete)
+ assert result.changes[0].kind == "delete"
+ # Parent tombstones permit matching child tombstones, but never child upserts.
+ parent_delete = newest.model_copy(
+ update={
+ "version": SessionVersion(revision=3, content_revision=13),
+ "operation": "delete",
+ "original": {},
+ }
+ )
+ child_delete = current_delete.model_copy(update={"version": parent_delete.version})
+ result = await ingest(database, fence, parent_delete, child_delete)
+ assert [change.kind for change in result.changes] == ["delete", "delete"]
diff --git a/backend/airweave/domains/native_ingestion/tests/test_maintenance_ownership.py b/backend/airweave/domains/native_ingestion/tests/test_maintenance_ownership.py
new file mode 100644
index 0000000000..a53b13b981
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_maintenance_ownership.py
@@ -0,0 +1,27 @@
+"""Provider timeout maintenance cannot terminate resumable native imports."""
+
+from datetime import timedelta
+
+import pytest
+from sqlalchemy import update
+
+from airweave.core.datetime_utils import utc_now_naive
+from airweave.crud.crud_sync_job import sync_job
+from airweave.domains.native_ingestion.tests.test_imports import native, start # noqa: F401
+from airweave.models.sync_job import SyncJob
+
+
+@pytest.mark.parametrize("native", ["knowledge"], indirect=True)
+async def test_stuck_provider_query_excludes_native_import(database, source, native): # noqa: F811
+ imported = await start(database, native)
+ _, provider = source
+ cutoff = utc_now_naive() - timedelta(minutes=15)
+ async with database() as db:
+ await db.execute(update(SyncJob).values(started_at=cutoff - timedelta(minutes=1)))
+ await db.commit()
+ jobs = await sync_job.get_stuck_jobs_by_status(
+ db, status=["running"], started_before=cutoff
+ )
+ assert {job.id for job in jobs} == {provider.job_id}
+ # Exclusion does not terminate or replace the durable native import.
+ assert await start(database, native) == imported
diff --git a/backend/airweave/domains/native_ingestion/tests/test_observed_membership.py b/backend/airweave/domains/native_ingestion/tests/test_observed_membership.py
new file mode 100644
index 0000000000..ed1f8833d9
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_observed_membership.py
@@ -0,0 +1,177 @@
+"""Bounded native imports own only freshly observed parent/ancestor membership."""
+
+import hashlib
+import json
+
+import pytest
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.cycle_models import (
+ BeginCycle,
+ CompleteCycle,
+ CycleConfiguration,
+)
+from airweave.domains.entities.canonical.cycle_store import (
+ CycleConflict,
+ begin_cycle,
+ complete_cycle,
+ cursor_row,
+ cycle_state,
+ next_scope_work,
+)
+from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ ReconcileScan,
+ ScanContinuation,
+)
+from airweave.domains.entities.canonical.scan_store import CanonicalScanStore
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.models import IngestNativePage, SessionVersion
+from airweave.domains.native_ingestion.service import NativeIngestionService
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.domains.native_ingestion.tests.test_ingestion import (
+ NOW,
+ bind,
+ ingest,
+ retained,
+ snapshot,
+)
+
+
+async def start_cycle(database, canonical, fence, *, expected=None):
+ async with database() as db, UnitOfWork(db):
+ await canonical._fenced_sync(db, fence)
+ return await begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence,
+ expected=expected,
+ configuration=CycleConfiguration(
+ fingerprint="a" * 64,
+ parents={"session": (None,), "message": ("session",)},
+ completion_policies={"session": "discovery_only", "message": "discovery_only"},
+ membership="observed",
+ ),
+ ),
+ )
+
+
+async def finish_scope(database, canonical, fence, cycle, scope, items=(), *, epoch=None):
+ scans = CanonicalScanStore(canonical)
+ async with database() as db, UnitOfWork(db):
+ previous = await scans.read(db, fence, scope)
+ state = await scans.begin(
+ db,
+ BeginScan(
+ fence=fence,
+ scope=scope,
+ cycle_id=cycle.version.cycle_id,
+ fingerprint=cycle.configuration.fingerprint,
+ expected=previous.version if previous else None,
+ expected_parent_epoch=epoch,
+ ),
+ )
+ async with database() as db:
+ result = await NativeIngestionService(NativeIngestionStore(canonical)).page(
+ db,
+ IngestNativePage(
+ fence=fence,
+ observed_at=NOW,
+ snapshots=items,
+ scope=scope,
+ cycle_id=cycle.version.cycle_id,
+ expected=state.version,
+ final=True,
+ continuation=ScanContinuation(),
+ ),
+ )
+ async with database() as db, UnitOfWork(db):
+ await scans.reconcile(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=scope,
+ cycle_id=cycle.version.cycle_id,
+ expected=result.state.version,
+ observed_at=NOW,
+ ),
+ )
+ return result
+
+
+async def finish_cycle(database, canonical, fence):
+ async with database() as db, UnitOfWork(db):
+ sync = await canonical._fenced_sync(db, fence)
+ cycle = cycle_state(await cursor_row(db, fence))
+ return await complete_cycle(db, sync, CompleteCycle(fence=fence, expected=cycle.version))
+
+
+async def test_only_observed_unchanged_parent_requires_transcript_and_empty_import_keeps_all(
+ database, source
+):
+ _, fence = source
+ canonical = CanonicalRecordStore()
+ await bind(database, fence, dataset="sessions")
+ first, second = (
+ snapshot(
+ name,
+ identity=RecordIdentity(record_type="session", native_id=name),
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ for name in ("first", "second")
+ )
+ await ingest(database, fence, first, second)
+ cycle = await start_cycle(database, canonical, fence)
+ root = CompletedScope(record_type="session")
+ result = await finish_scope(database, canonical, fence, cycle, root, (first,))
+ assert result.capture.unchanged == 1 and result.capture.changes == ()
+ with pytest.raises(CycleConflict, match="incomplete child"):
+ await finish_cycle(database, canonical, fence)
+ async with database() as db, UnitOfWork(db):
+ work = await next_scope_work(db, fence, cycle.version.cycle_id)
+ assert work.parent.identity == first.identity
+ with pytest.raises(CycleConflict, match="membership"):
+ await finish_scope(
+ database,
+ canonical,
+ fence,
+ cycle,
+ CompletedScope(record_type="message", container_id="second", parent=second.identity),
+ epoch=1,
+ )
+ await finish_scope(
+ database,
+ canonical,
+ fence,
+ cycle,
+ CompletedScope(record_type="message", container_id="first", parent=first.identity),
+ epoch=work.parent_visibility_epoch,
+ )
+ completed = await finish_cycle(database, canonical, fence)
+ assert completed.last_full_capture.discovery == "incomplete"
+ empty = await start_cycle(database, canonical, fence, expected=completed.version)
+ await finish_scope(database, canonical, fence, empty, root)
+ async with database() as db, UnitOfWork(db):
+ assert await next_scope_work(db, fence, empty.version.cycle_id) is None
+ completed_empty = await finish_cycle(database, canonical, fence)
+ assert completed_empty.last_full_capture.discovery == "incomplete"
+ rows = await retained(database, fence)
+ assert len(rows) == 2 and all(row.deleted_at is None for row in rows)
+
+
+def test_default_configuration_digest_stays_compatible():
+ configuration = CycleConfiguration(fingerprint="a" * 64, parents={"session": (None,)})
+ legacy = {
+ "fingerprint": "a" * 64,
+ "parents": {"session": [None]},
+ "completion_policies": {"session": "exhaustive"},
+ }
+ assert (
+ configuration.digest()
+ == hashlib.sha256(json.dumps(legacy, sort_keys=True).encode()).hexdigest()
+ )
+ assert (
+ configuration.model_copy(update={"membership": "observed"}).digest()
+ != configuration.digest()
+ )
diff --git a/backend/airweave/domains/native_ingestion/tests/test_page_receipts.py b/backend/airweave/domains/native_ingestion/tests/test_page_receipts.py
new file mode 100644
index 0000000000..d2c9e1bab3
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_page_receipts.py
@@ -0,0 +1,147 @@
+"""Exact lost-response recovery uses the same real transaction as native page writes."""
+
+
+# Imported pytest fixture names are resolved through function parameters.
+# ruff: noqa: F811
+
+import asyncio
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import func, select
+
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.cycle_store import cursor_row, cycle_state
+from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity
+from airweave.domains.entities.canonical.scan_models import BeginScan
+from airweave.domains.entities.canonical.scan_store import CanonicalScanStore, ScanConflict
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, StaleWriter
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.models import SessionVersion
+from airweave.domains.native_ingestion.page_models import CommitNativePage
+from airweave.domains.native_ingestion.page_store import NativePageStore
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.domains.native_ingestion.tests.test_imports import native, start # noqa: F401
+from airweave.domains.native_ingestion.tests.test_ingestion import snapshot
+from airweave.models.entity import Entity
+from airweave.models.sync_job import SyncJob
+
+
+@pytest.fixture
+async def scope(database, native):
+ await start(database, native)
+ canonical = CanonicalRecordStore()
+ imports = NativeImportStore(NativeSourceStore(), canonical)
+ async with database() as db, UnitOfWork(db):
+ saved = await imports.active(
+ db, native.organization_id, native.source_connection_id, "request-one"
+ )
+ cycle = cycle_state(await cursor_row(db, saved.fence))
+ kind = "knowledge" if native.binding.dataset == "knowledge" else "session"
+ state = await CanonicalScanStore(canonical).begin(
+ db,
+ BeginScan(
+ fence=saved.fence,
+ scope=CompletedScope(record_type=kind),
+ cycle_id=saved.cycle_id,
+ fingerprint=cycle.configuration.fingerprint,
+ ),
+ )
+ return state
+
+
+def page_request(native, state):
+ item = snapshot(owner_id=native.binding.owner_id)
+ if native.binding.dataset == "sessions":
+ item = item.model_copy(
+ update={
+ "identity": RecordIdentity(record_type="session", native_id="one"),
+ "version": SessionVersion(revision=1, content_revision=1),
+ }
+ )
+ return CommitNativePage(
+ page_id=uuid4(),
+ scope=state.scope,
+ expected=state.version,
+ snapshots=(item,),
+ cursor={"next": "two"},
+ )
+
+
+async def commit(database, native, request):
+ canonical = CanonicalRecordStore()
+ pages = NativePageStore(
+ NativeImportStore(NativeSourceStore(), canonical), NativeIngestionStore(canonical)
+ )
+ async with database() as db, UnitOfWork(db):
+ return await pages.commit(
+ db, native.organization_id, native.source_connection_id, "request-one", request
+ )
+
+
+async def test_concurrent_page_retry_conflict_and_superseded_receipt(database, native, scope):
+ request = page_request(native, scope)
+ first, duplicate = await asyncio.gather(
+ commit(database, native, request), commit(database, native, request)
+ )
+ assert first == duplicate and first.changed == 1 and first.sequence == 1
+ async with database() as db:
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(Entity).where(Entity.sync_id == native.sync_id)
+ )
+ == 1
+ )
+ with pytest.raises(ScanConflict, match="retry conflicts"):
+ await commit(database, native, request.model_copy(update={"final": True}))
+ last = request.model_copy(
+ update={
+ "page_id": uuid4(),
+ "expected": first.version,
+ "snapshots": (),
+ "final": True,
+ }
+ )
+ final = await commit(database, native, last)
+ assert final.phase == "reconciling"
+ assert await commit(database, native, last) == final
+ with pytest.raises(ScanConflict, match="Page changed"):
+ await commit(database, native, request)
+
+
+async def test_cancelled_import_cannot_replay_a_page_receipt(database, native, scope):
+ request = page_request(native, scope)
+ await commit(database, native, request)
+ imported = await start(database, native)
+ async with database() as db:
+ job = await db.get(SyncJob, imported.import_id)
+ job.status = "cancelled"
+ await db.commit()
+ with pytest.raises(StaleWriter):
+ await commit(database, native, request)
+
+
+async def test_receipt_failure_rolls_back_capture_and_cursor(database, native, scope, monkeypatch):
+ import airweave.domains.native_ingestion.page_store as module
+
+ def fail(**kwargs):
+ raise RuntimeError("receipt persistence failed")
+
+ monkeypatch.setattr(module, "NativePageReceipt", fail)
+ with pytest.raises(RuntimeError, match="receipt persistence"):
+ await commit(database, native, page_request(native, scope))
+ canonical = CanonicalRecordStore()
+ imports = NativeImportStore(NativeSourceStore(), canonical)
+ async with database() as db, UnitOfWork(db):
+ saved = await imports.active(
+ db, native.organization_id, native.source_connection_id, "request-one"
+ )
+ current = await CanonicalScanStore(canonical).read(db, saved.fence, scope.scope)
+ assert current == scope
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(Entity).where(Entity.sync_id == native.sync_id)
+ )
+ == 0
+ )
diff --git a/backend/airweave/domains/native_ingestion/tests/test_pages.py b/backend/airweave/domains/native_ingestion/tests/test_pages.py
new file mode 100644
index 0000000000..8384285eba
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_pages.py
@@ -0,0 +1,230 @@
+"""Native admission, capture and sweep continuation share one real transaction."""
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration
+from airweave.domains.entities.canonical.requests import (
+ CompletedScope,
+ RecordIdentity,
+ RemovedScope,
+)
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ ReconcileScan,
+ ScanContinuation,
+)
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.models import IngestNativePage, SessionVersion
+from airweave.domains.native_ingestion.service import NativeIngestionService
+from airweave.domains.native_ingestion.store import NativeIngestionStore
+from airweave.domains.native_ingestion.tests.test_ingestion import (
+ NOW,
+ bind,
+ ingest,
+ retained,
+ snapshot,
+)
+from airweave.models.entity import Entity
+
+
+async def begin(database, canonical, fence, *, sessions=False):
+ async with database() as db:
+ return await canonical.begin_cycle(
+ db,
+ BeginCycle(
+ fence=fence,
+ configuration=CycleConfiguration(
+ fingerprint="a" * 64,
+ parents={"session": (None,), "message": ("session",)}
+ if sessions
+ else {"knowledge": (None,)},
+ ),
+ ),
+ )
+
+
+async def scan(database, canonical, fence, cycle, scope, epoch=None):
+ async with database() as db:
+ return await canonical.begin_scan(
+ db,
+ BeginScan(
+ fence=fence,
+ scope=scope,
+ cycle_id=cycle.version.cycle_id,
+ fingerprint="a" * 64,
+ expected_parent_epoch=epoch,
+ ),
+ )
+
+
+async def page(database, fence, state, *items, final=False):
+ service = NativeIngestionService(NativeIngestionStore(CanonicalRecordStore()))
+ async with database() as db:
+ return await service.page(
+ db,
+ IngestNativePage(
+ fence=fence,
+ observed_at=NOW,
+ snapshots=items,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ continuation=ScanContinuation(value={"next": "two"}),
+ final=final,
+ ),
+ )
+
+
+async def test_unchanged_new_sweep_survives_reconciliation_without_new_revision(database, source):
+ canonical, fence = source
+ await bind(database, fence)
+ original = snapshot()
+ await ingest(database, fence, original)
+ cycle = await begin(database, canonical, fence)
+ state = await scan(database, canonical, fence, cycle, CompletedScope(record_type="knowledge"))
+ result = await page(database, fence, state, original, final=True)
+ assert result.capture.unchanged == 1 and result.capture.changes == ()
+ async with database() as db:
+ reconciled = await canonical.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=result.state.version,
+ observed_at=NOW,
+ ),
+ )
+ row = (await retained(database, fence))[0]
+ assert row.last_seen_run_id == state.version.sweep_id
+ assert row.record_revision == 1 and row.deleted_at is None
+ assert reconciled.capture.changes == ()
+
+
+async def test_withdrawn_or_tombstoned_snapshot_cannot_be_seen_again(database, source):
+ canonical, fence = source
+ await bind(database, fence)
+ item = snapshot()
+ await ingest(database, fence, item)
+ cycle = await begin(database, canonical, fence)
+ state = await scan(database, canonical, fence, cycle, CompletedScope(record_type="knowledge"))
+ async with database() as db:
+ await canonical.remove_scope(
+ db,
+ fence,
+ RemovedScope(record_type="knowledge", removal_reason="access_revoked", observed_at=NOW),
+ )
+ with pytest.raises(NativeAdmissionError, match="active original sightings"):
+ await page(database, fence, state, item)
+ with pytest.raises(NativeAdmissionError, match="not tombstones"):
+ await page(database, fence, state, item.model_copy(update={"operation": "delete"}))
+ async with database() as db:
+ current = await canonical.read_scan(db, fence, state.scope)
+ assert current == state
+
+
+async def test_stale_unchanged_child_does_not_attest_membership_of_newer_parent(database, source):
+ canonical, fence = source
+ await bind(database, fence, dataset="sessions")
+ parent = snapshot(
+ identity=RecordIdentity(record_type="session", native_id="s"),
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ child = snapshot(
+ identity=RecordIdentity(record_type="message", native_id="m", container_id="s"),
+ parent=parent.identity,
+ version=parent.version,
+ )
+ await ingest(database, fence, parent, child)
+ newer = parent.model_copy(update={"version": SessionVersion(revision=1, content_revision=2)})
+ await ingest(database, fence, newer)
+ cycle = await begin(database, canonical, fence, sessions=True)
+ root = await scan(database, canonical, fence, cycle, CompletedScope(record_type="session"))
+ root_result = await page(database, fence, root, newer, final=True)
+ async with database() as db:
+ await canonical.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=root.scope,
+ cycle_id=root.cycle_id,
+ expected=root_result.state.version,
+ observed_at=NOW,
+ ),
+ )
+ async with database() as db:
+ epoch = await db.scalar(
+ select(Entity.visibility_epoch).where(
+ Entity.sync_id == fence.sync_id, Entity.native_id == "s"
+ )
+ )
+ state = await scan(
+ database,
+ canonical,
+ fence,
+ cycle,
+ CompletedScope(record_type="message", container_id="s", parent=parent.identity),
+ epoch,
+ )
+ with pytest.raises(NativeAdmissionError, match="attested native session"):
+ await page(database, fence, state, child)
+ assert (
+ await ingest(database, fence, child)
+ ).unchanged == 1 # Ordinary retry is not a new sighting.
+ updated = child.model_copy(update={"version": newer.version})
+ result = await page(database, fence, state, updated)
+ assert result.capture.changes[0].record.revision == 2
+
+
+async def test_failure_after_capture_rolls_back_records_sightings_and_cursor(
+ database, source, monkeypatch
+):
+ import airweave.domains.entities.canonical.scan_store as scan_module
+
+ canonical, fence = source
+ await bind(database, fence)
+ original = snapshot()
+ await ingest(database, fence, original)
+ before = (await retained(database, fence))[0].last_seen_run_id
+ cycle = await begin(database, canonical, fence)
+ state = await scan(database, canonical, fence, cycle, CompletedScope(record_type="knowledge"))
+
+ def fail(*args, **kwargs):
+ raise RuntimeError("synthetic failure after page writes")
+
+ monkeypatch.setattr(scan_module, "publish_scope", fail)
+ with pytest.raises(RuntimeError, match="after page writes"):
+ await page(database, fence, state, original, snapshot("new"), final=True)
+ rows = await retained(database, fence)
+ assert len(rows) == 1 and rows[0].last_seen_run_id == before
+ async with database() as db:
+ assert await canonical.read_scan(db, fence, state.scope) == state
+
+
+async def test_empty_final_page_is_explicit_completion_not_implicit_absence(database, source):
+ canonical, fence = source
+ await bind(database, fence)
+ await ingest(database, fence, snapshot())
+ cycle = await begin(database, canonical, fence)
+ state = await scan(database, canonical, fence, cycle, CompletedScope(record_type="knowledge"))
+ partial = await page(database, fence, state)
+ assert partial.state.phase == "collecting"
+ assert (await retained(database, fence))[0].deleted_at is None
+ final = await page(database, fence, partial.state, final=True)
+ assert final.state.phase == "reconciling"
+ assert (await retained(database, fence))[0].deleted_at is None
+ async with database() as db:
+ await canonical.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=final.state.version,
+ observed_at=NOW,
+ ),
+ )
+ row = (await retained(database, fence))[0]
+ assert row.removal_reason == "absent" and row.source_payload["version"]["revision"] == 1
diff --git a/backend/airweave/domains/native_ingestion/tests/test_projection.py b/backend/airweave/domains/native_ingestion/tests/test_projection.py
new file mode 100644
index 0000000000..36e6849547
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_projection.py
@@ -0,0 +1,349 @@
+"""Native model-shaped fixtures; real SQL publication, synthetic embedding/destination."""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.domains.entities.canonical.projection_mappers import map_record
+from airweave.domains.entities.canonical.projection_policy import excluded_from_search
+from airweave.domains.entities.canonical.projection_store import (
+ CanonicalProjectionStore,
+ current_extraction,
+)
+from airweave.domains.entities.canonical.query import CanonicalQueryService
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.entities.canonical.tests.test_extraction_coverage import (
+ destination,
+ projector,
+)
+from airweave.domains.entities.canonical.text_query import CanonicalTextReader
+from airweave.domains.native_ingestion.models import RecordVersion, SessionVersion
+from airweave.domains.native_ingestion.tests.test_ingestion import NOW, bind, ingest, snapshot
+
+
+def knowledge(**changes):
+ original = {
+ "id": "one",
+ "type": "page",
+ "revision": 1,
+ "title": "Fundraising",
+ "description": "A useful original page",
+ "path": "notes/fundraising",
+ "body": "नमस्ते 👋\n# Original body",
+ "created_at": NOW.isoformat(),
+ "updated_at": NOW.isoformat(),
+ "archived_at": None,
+ "citations": {},
+ "field_citations": {},
+ "aliases": [],
+ "image": None,
+ "about": [],
+ "user_notes": "A note",
+ }
+ return snapshot(original=original | changes, source_created_at=NOW, source_updated_at=NOW)
+
+
+def session():
+ return snapshot(
+ identity=RecordIdentity(record_type="session", native_id="s1"),
+ version=SessionVersion(revision=1, content_revision=1),
+ source_created_at=NOW,
+ source_updated_at=NOW,
+ original={
+ "id": "s1",
+ "title": "Session title",
+ "title_source": "user",
+ "description": "Original session description",
+ "source": "chat",
+ "metadata": {},
+ "archived": True,
+ "pinned": False,
+ "revision": 1,
+ "content_revision": 1,
+ "replay_epoch": 4,
+ "message_count": 1,
+ "earlier_message_count": 0,
+ "project": None,
+ "effective_project": None,
+ "work": None,
+ "created_at": NOW.isoformat(),
+ "updated_at": NOW.isoformat(),
+ },
+ )
+
+
+def message(content, **changes):
+ parent = session()
+ original = {
+ "id": "m1",
+ "input_id": None,
+ "ordinal": 0,
+ "import_provenance": None,
+ "payload": {
+ "role": "user",
+ "content": content,
+ "api_content": "DO NOT INDEX",
+ "reasoning_content": "PRIVATE REPLAY TEXT",
+ },
+ "created_at": NOW.isoformat(),
+ }
+ return snapshot(
+ identity=RecordIdentity(record_type="message", native_id="m1", container_id="s1"),
+ parent=parent.identity,
+ version=parent.version,
+ original=original | changes,
+ source_created_at=NOW,
+ )
+
+
+async def test_knowledge_identity_rename_archive_and_malformed_original(database, source):
+ _, fence = source
+ await bind(database, fence)
+ result = await ingest(database, fence, knowledge())
+ record = result.changes[0].record
+ async with map_record(record, "almanac", AsyncMock()) as mapped:
+ assert mapped.parts[0].entity.entity_id == "one"
+ assert mapped.parts[0].native_body.text == knowledge().original["body"]
+ renamed = knowledge(path="notes/renamed", revision=2).model_copy(
+ update={"version": RecordVersion(revision=2)}
+ )
+ record = (await ingest(database, fence, renamed)).changes[0].record
+ async with map_record(record, "almanac", AsyncMock()) as mapped:
+ assert mapped.parts[0].entity.entity_id == "one"
+ assert mapped.parts[0].entity.path == "notes/renamed"
+ bad = record.model_copy(
+ update={
+ "payload": renamed.model_copy(
+ update={"original": renamed.original | {"id": "wrong"}}
+ ).model_dump(mode="json")
+ }
+ )
+ with pytest.raises(ValueError, match="ID differs"):
+ async with map_record(bad, "almanac", AsyncMock()):
+ pass
+ archived = renamed.model_copy(
+ update={"original": renamed.original | {"archived_at": NOW.isoformat()}}
+ )
+ assert excluded_from_search(
+ record.model_copy(update={"payload": archived.model_dump(mode="json")}), "almanac"
+ )
+
+
+async def test_mixed_message_publishes_partial_coverage_and_exact_native_text(
+ database, source, tmp_path
+):
+ _, fence = source
+ await bind(database, fence, dataset="sessions")
+ text = "नमस्ते 👋\n# Content\nOriginal user words"
+ item = message(
+ [
+ {"type": "text", "text": text},
+ {"type": "image_url", "image_url": {"url": "https://example.org/image"}},
+ ]
+ )
+ result = await ingest(database, fence, session(), item)
+ assert not excluded_from_search(
+ result.changes[0].record, "almanac"
+ ) # Archived sessions included today.
+ storage = FilesystemBackend(tmp_path)
+ projection = projector(database, storage)
+ indexed = []
+
+ async def fixed_chunks(entities, context, runtime):
+ indexed.extend(entity.textual_representation for entity in entities)
+ return projection._processor._multiply_entities(
+ entities, [[{"text": entity.textual_representation}] for entity in entities], context
+ )
+
+ projection._processor._chunk_entities = fixed_chunks
+ async with database() as db:
+ pending = await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id)
+ work = next(work for work in pending if work.record.identity.record_type == "message")
+ async with map_record(work.record, "almanac", storage) as mapped:
+ assert mapped.parts[0].entity.message_id == "m1"
+ assert mapped.parts[0].entity.session_id == "s1"
+ assert mapped.parts[0].entity.block_index == 0
+ assert mapped.parts[0].native_body.text == text
+ assert mapped.parts[1].omission == "unsupported_format"
+ assert await projection.project_one(
+ work, "almanac", destination(work.binding.collection_id), MagicMock()
+ )
+ assert indexed[0].endswith(text)
+ assert "DO NOT INDEX" not in indexed[0] and "PRIVATE REPLAY TEXT" not in indexed[0]
+ async with database() as db:
+ coverage = await current_extraction(
+ db, fence.organization_id, fence.sync_id, work.record.id, work.record.revision
+ )
+ assert coverage.status == "partial"
+ assert [part.outcome for part in coverage.parts] == ["indexed", "unsupported"]
+ reader = CanonicalTextReader(
+ CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-key"), storage
+ )
+ async with database() as db:
+ representations = await reader.list(
+ db, fence.organization_id, fence.sync_id, work.record.id, work.record.revision
+ )
+ ref = representations.representations[0]
+ native = await reader.read(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ work.record.id,
+ work.record.revision,
+ ref.generation,
+ ref.id,
+ )
+ assert ref.kind == "native_text" and native.text == text
+
+
+async def test_inactive_import_and_invalid_text_are_not_successful_empty_indexing(database, source):
+ _, fence = source
+ await bind(database, fence, dataset="sessions")
+ inactive = message(
+ "old", import_provenance={"legacy_row_id": 1, "active": False, "compacted": True}
+ )
+ result = await ingest(database, fence, session(), inactive)
+ assert excluded_from_search(result.changes[1].record, "almanac")
+ malformed = message([{"type": "text", "text": {"unexpected": "object"}}])
+ record = result.changes[1].record.model_copy(
+ update={"payload": malformed.model_dump(mode="json")}
+ )
+ with pytest.raises(ValueError):
+ async with map_record(record, "almanac", AsyncMock()):
+ pass
+
+
+async def test_person_fields_enter_index_without_rewriting_native_body(database, source, tmp_path):
+ _, fence = source
+ await bind(database, fence)
+ item = knowledge(
+ type="person",
+ emails=["sam@example.org"],
+ roles=[
+ {
+ "organisation": {"type": "organisation", "id": "opaque-id"},
+ "title": "Research engineer",
+ "start": "2024",
+ "end": None,
+ }
+ ],
+ links=[{"label": "Profile", "url": "https://example.org/sam", "handle": "samdev"}],
+ metadata={"secret": "UNSELECTED_METADATA"},
+ )
+ captured = await ingest(database, fence, item)
+ storage = FilesystemBackend(tmp_path)
+ projection = projector(database, storage)
+ indexed = []
+
+ async def fixed_chunks(entities, context, runtime):
+ indexed.extend(entity.textual_representation for entity in entities)
+ return projection._processor._multiply_entities(
+ entities, [[{"text": entity.textual_representation}] for entity in entities], context
+ )
+
+ projection._processor._chunk_entities = fixed_chunks
+ async with database() as db:
+ work = (await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id))[
+ 0
+ ]
+ assert await projection.project_one(
+ work, "almanac", destination(work.binding.collection_id), MagicMock()
+ )
+ assert "sam@example.org" in indexed[0] and "Research engineer" in indexed[0]
+ assert "samdev" in indexed[0] and "2024" in indexed[0]
+ assert "UNSELECTED_METADATA" not in indexed[0] and "opaque-id" not in indexed[0]
+ assert captured.changes[0].record.payload["original"] == item.original
+ reader = CanonicalTextReader(
+ CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-key"), storage
+ )
+ async with database() as db:
+ representations = await reader.list(
+ db, fence.organization_id, fence.sync_id, work.record.id, 1
+ )
+ ref = representations.representations[0]
+ content = await reader.read(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ work.record.id,
+ 1,
+ ref.generation,
+ ref.id,
+ view="content",
+ )
+ prepared = await reader.read(
+ db,
+ fence.organization_id,
+ fence.sync_id,
+ work.record.id,
+ 1,
+ ref.generation,
+ ref.id,
+ view="index",
+ )
+ assert content.text == item.original["body"]
+ assert "sam@example.org" in prepared.text
+
+
+@pytest.mark.parametrize(
+ "kind,fields,expected",
+ [
+ (
+ "organisation",
+ {"domains": ["example.org"], "industry": ["Robotics"], "founded": "2020"},
+ "Robotics",
+ ),
+ (
+ "place",
+ {"place_kind": "Office", "address": {"locality": "मुंबई", "country": "India"}},
+ "मुंबई",
+ ),
+ ("creative_work", {"work_kind": "Book", "published_on": "2025-01-01"}, "2025-01-01"),
+ ],
+)
+def test_selected_knowledge_fields_preserve_native_text(kind, fields, expected):
+ from airweave.domains.native_ingestion.knowledge_fields import knowledge_details
+
+ text = knowledge_details(kind, fields | {"metadata": {"private": "UNSELECTED"}})
+ assert expected in text and "UNSELECTED" not in text
+
+
+def test_invalid_selected_field_fails_instead_of_stringifying_json():
+ from pydantic import ValidationError
+
+ from airweave.domains.native_ingestion.knowledge_fields import knowledge_details
+
+ with pytest.raises(ValidationError):
+ knowledge_details("person", {"emails": [{"unexpected": "value"}]})
+
+
+def test_event_and_relationship_values_are_searchable_without_resolving_references():
+ from airweave.domains.native_ingestion.knowledge_fields import knowledge_details
+
+ assert "2026-10-01" in knowledge_details(
+ "event",
+ {
+ "schedule": {
+ "kind": "all_day",
+ "start_on": "2026-10-01",
+ "end_on_exclusive": "2026-10-02",
+ }
+ },
+ )
+ details = knowledge_details(
+ "person",
+ {
+ "related_people": [
+ {"person": {"type": "person", "id": "hidden-id"}, "relationship": "Co-founder"}
+ ],
+ "education": [{"institution": {"type": "organisation", "id": "hidden-school"}}],
+ },
+ )
+ assert details == "related people: relationship: Co-founder"
+ coordinates = knowledge_details(
+ "place", {"coordinates": {"latitude": 19.07, "longitude": 72.87}}
+ )
+ assert "19.07" in coordinates and "72.87" in coordinates
diff --git a/backend/airweave/domains/native_ingestion/tests/test_publish_cli.py b/backend/airweave/domains/native_ingestion/tests/test_publish_cli.py
new file mode 100644
index 0000000000..34b6c14c8d
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_publish_cli.py
@@ -0,0 +1,171 @@
+"""CLI validation and outcome claims, without duplicating publisher lifecycle tests."""
+
+import json
+from uuid import UUID
+
+import httpx
+import pytest
+from evaluation import native_import_cli as cli
+
+from airweave.domains.native_ingestion.tests.test_ingestion import snapshot
+
+
+@pytest.fixture
+def invocation(tmp_path, monkeypatch):
+ monkeypatch.setenv("AIRWEAVE_API_KEY", "private-test-key")
+ path = tmp_path / "staged.json"
+ path.write_text(json.dumps([snapshot().model_dump(mode="json")]))
+ return path, [
+ "--url",
+ "http://127.0.0.1:18086/api/v1/",
+ "--input",
+ str(path),
+ "--owner",
+ "owner-one",
+ "--dataset",
+ "knowledge",
+ "--collection",
+ "native-test",
+ "--request-key",
+ "repeat-this",
+ ]
+
+
+def transport(monkeypatch, handler):
+ client = httpx.AsyncClient
+
+ def configured(**kw):
+ assert kw["follow_redirects"] is False and kw["trust_env"] is False
+ return client(transport=httpx.MockTransport(handler), **kw)
+
+ monkeypatch.setattr(cli.httpx, "AsyncClient", configured)
+
+
+@pytest.mark.parametrize("invalid", ["json", "oversize", "owner", "credentials"])
+def test_invalid_input_makes_no_http(invocation, monkeypatch, capsys, invalid):
+ path, args = invocation
+ calls = []
+
+ def no_http(request):
+ calls.append(request)
+ raise AssertionError("Invalid input reached HTTP")
+
+ transport(monkeypatch, no_http)
+ if invalid == "json":
+ path.write_text('{"secret-content":')
+ elif invalid == "oversize":
+ monkeypatch.setattr(cli, "MAX_INPUT_BYTES", 1)
+ elif invalid == "owner":
+ args[5] = "wrong-owner"
+ else:
+ monkeypatch.delenv("AIRWEAVE_API_KEY")
+ assert cli.main(args) in (2, 4)
+ assert not calls
+ output = capsys.readouterr()
+ assert not output.out
+ assert "secret-content" not in output.err and "private-test-key" not in output.err
+
+
+def test_lost_response_is_unknown_not_success(invocation, monkeypatch, capsys):
+ _, args = invocation
+ calls = []
+
+ def lose(request):
+ calls.append(request)
+ raise httpx.ReadTimeout("private-test-key and private body", request=request)
+
+ transport(monkeypatch, lose)
+ assert cli.main(args) == 3
+ assert len(calls) == 1
+ output = capsys.readouterr()
+ assert not output.out
+ assert "outcome unknown" in output.err and "private" not in output.err
+
+
+def test_completed_retry_outputs_capture_summary(invocation, monkeypatch, capsys):
+ _, args = invocation
+ requests = []
+ identifier = str(UUID(int=1))
+
+ def destination(request):
+ requests.append(request)
+ assert request.headers["X-API-Key"] == "private-test-key"
+ if request.url.path.endswith("/native/sources"):
+ return httpx.Response(
+ 200,
+ json={
+ "source_connection_id": identifier,
+ "sync_id": identifier,
+ "organization_id": identifier,
+ "binding": {"owner_id": "owner-one", "dataset": "knowledge"},
+ "collection": "native-test",
+ "available": True,
+ },
+ )
+ assert request.url.path.endswith("/imports/repeat-this")
+ return httpx.Response(
+ 200,
+ json={
+ "source_id": identifier,
+ "import_id": identifier,
+ "cycle_id": identifier,
+ "request_key": "repeat-this",
+ "request": json.loads(request.content),
+ "status": "completed",
+ "summary": {
+ "outcome": "completed",
+ "coverage": "bounded",
+ "finished_at": "2026-10-01T00:00:00Z",
+ "sequence": 1,
+ "completed_scopes": 1,
+ "capture_complete": True,
+ "indexing": "not_verified",
+ },
+ },
+ )
+
+ transport(monkeypatch, destination)
+ assert cli.main(args) == 0
+ assert len(requests) == 2
+ output = capsys.readouterr()
+ assert not output.err
+ summary = json.loads(output.out)
+ assert summary["capture_complete"] and summary["indexing"] == "not_verified"
+ assert "private-test-key" not in output.out
+
+
+@pytest.mark.parametrize(
+ "url",
+ [
+ "https://user:secret@example.com/api/",
+ "https://example.com/api/?",
+ "https://example.com/api/#",
+ "http://localhost.example.com/api/",
+ "file:///tmp/api/",
+ ],
+)
+def test_unsafe_url_rejected(invocation, url, capsys):
+ _, args = invocation
+ args[1] = url
+ assert cli.main(args) == 2
+ assert url not in capsys.readouterr().err
+
+
+def test_help_needs_no_server_settings(tmp_path):
+ import os
+ import subprocess
+ import sys
+ from pathlib import Path
+
+ backend = Path(__file__).resolve().parents[4]
+ result = subprocess.run(
+ [sys.executable, "-m", "evaluation.native_import_cli", "--help"],
+ cwd=tmp_path,
+ env={"PATH": os.defpath, "PYTHONPATH": f"{backend}:{backend.parent}"},
+ capture_output=True,
+ text=True,
+ timeout=15,
+ )
+ assert result.returncode == 0, result.stderr
+ assert "--request-key" in result.stdout
+ assert not result.stderr
diff --git a/backend/airweave/domains/native_ingestion/tests/test_publisher.py b/backend/airweave/domains/native_ingestion/tests/test_publisher.py
new file mode 100644
index 0000000000..2f50fbb8e7
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_publisher.py
@@ -0,0 +1,148 @@
+"""Staged publisher resumes durable HTTP progress after lost responses."""
+
+# ruff: noqa: F811
+import httpx
+import pytest
+from evaluation.native_import import NativePublishError, publish_native
+from sqlalchemy import func, select
+
+from airweave.api import deps
+from airweave.api.v1.endpoints.native_sources import router as source_router
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.native_ingestion.models import SessionVersion
+from airweave.domains.native_ingestion.source_service import NativeSources
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.domains.native_ingestion.tests.test_import_api import native_api # noqa: F401
+from airweave.domains.native_ingestion.tests.test_ingestion import snapshot
+from airweave.models.entity import Entity
+from airweave.models.sync_job import SyncJob
+
+
+class LostResponse(httpx.AsyncBaseTransport):
+ def __init__(self, inner, suffix):
+ self.inner = inner
+ self.suffix = suffix
+ self.dropped = False
+
+ async def handle_async_request(self, request):
+ response = await self.inner.handle_async_request(request)
+ if not self.dropped and request.url.path.endswith(self.suffix):
+ self.dropped = True
+ await response.aclose()
+ raise httpx.ReadError("synthetic lost response", request=request)
+ return response
+
+ async def aclose(self):
+ await self.inner.aclose()
+
+
+def wire_sources(client):
+ app = client._transport.app
+ app.include_router(source_router, prefix="/native/sources")
+ container = app.dependency_overrides[deps.get_container]()
+ container.native_sources = NativeSources(NativeSourceStore())
+ app.dependency_overrides[deps.get_container] = lambda: container
+
+
+@pytest.mark.parametrize("lost", ["/pages", "/scopes/reconcile", "/complete"])
+async def test_lost_response_resume_and_changed_input_conflict(
+ native_api, database, lost, monkeypatch
+):
+ client, _, source, _ = native_api
+ wire_sources(client)
+ client._transport = LostResponse(client._transport, lost)
+ import evaluation.native_import as publisher
+
+ monkeypatch.setattr(publisher, "_PAGE_SIZE", 2)
+ originals = tuple(snapshot(str(index), owner_id="owner") for index in range(3))
+ kwargs = {
+ "owner_id": "owner",
+ "dataset": "knowledge",
+ "collection": "native-api",
+ "request_key": "stable-publisher",
+ "snapshots": originals,
+ }
+ with pytest.raises(httpx.ReadError):
+ await publish_native(client, **kwargs)
+ result = await publish_native(client, **kwargs)
+ assert result.status == "completed" and result.summary.capture_complete
+ assert await publish_native(client, **kwargs) == result
+ with pytest.raises(NativePublishError, match="HTTP 409"):
+ await publish_native(
+ client, **{**kwargs, "snapshots": (snapshot("different", owner_id="owner"),)}
+ )
+ async with database() as db:
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(Entity).where(Entity.sync_id == source.sync_id)
+ )
+ == 3
+ )
+ records = (await db.scalars(select(Entity).where(Entity.sync_id == source.sync_id))).all()
+ assert all(record.record_revision == 1 for record in records)
+ assert result.summary.sequence == 3
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == source.sync_id)
+ )
+ == 1
+ )
+
+
+@pytest.mark.parametrize("native_api", ["knowledge", "sessions"], indirect=True)
+async def test_empty_dataset_and_empty_session(native_api):
+ client, _, source, _ = native_api
+ wire_sources(client)
+ snapshots = ()
+ if source.binding.dataset == "sessions":
+ snapshots = (
+ snapshot(
+ "session",
+ owner_id="owner",
+ identity=RecordIdentity(record_type="session", native_id="session"),
+ version=SessionVersion(revision=1, content_revision=0),
+ ),
+ )
+ result = await publish_native(
+ client,
+ owner_id="owner",
+ dataset=source.binding.dataset,
+ collection="native-api",
+ request_key="empty",
+ snapshots=snapshots,
+ )
+ assert result.status == "completed"
+ assert result.summary.completed_scopes == (2 if snapshots else 1)
+
+
+async def test_preflight_rejects_mismatched_message_before_http():
+ calls = []
+
+ def response(request):
+ calls.append(request)
+ return httpx.Response(500)
+
+ parent = snapshot(
+ "session",
+ identity=RecordIdentity(record_type="session", native_id="session"),
+ version=SessionVersion(revision=1, content_revision=2),
+ )
+ child = snapshot(
+ "message",
+ identity=RecordIdentity(record_type="message", native_id="message", container_id="session"),
+ parent=parent.identity,
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ async with httpx.AsyncClient(
+ base_url="http://test/", transport=httpx.MockTransport(response)
+ ) as client:
+ with pytest.raises(NativePublishError, match="matching session"):
+ await publish_native(
+ client,
+ owner_id=parent.owner_id,
+ dataset="sessions",
+ collection="test",
+ request_key="invalid",
+ snapshots=(parent, child),
+ )
+ assert calls == []
diff --git a/backend/airweave/domains/native_ingestion/tests/test_scope_api.py b/backend/airweave/domains/native_ingestion/tests/test_scope_api.py
new file mode 100644
index 0000000000..d103261baf
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_scope_api.py
@@ -0,0 +1,168 @@
+"""HTTP scope/page/reconcile flow using actual retained records and scan transactions."""
+
+# Imported fixtures are resolved by pytest parameter name.
+# ruff: noqa: F811
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.native_ingestion.models import SessionVersion
+from airweave.domains.native_ingestion.tests.test_import_api import native_api # noqa: F401
+from airweave.domains.native_ingestion.tests.test_ingestion import snapshot
+from airweave.models.entity import Entity
+from airweave.models.sync_job import SyncJob
+
+
+@pytest.mark.parametrize("coverage", ["bounded", "complete"])
+async def test_http_scope_pages_retry_read_and_reconcile(database, native_api, coverage):
+ client, _, source, logs = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/scope-key"
+ started = await client.put(base, json={"snapshot_id": "snapshot", "coverage": coverage})
+ assert started.status_code == 200
+ scope = {"record_type": "knowledge"}
+ assert (await client.post(base + "/scopes/read", json={"scope": scope})).status_code == 409
+ opened = await client.put(base + "/scopes", json={"scope": scope})
+ assert opened.status_code == 200, opened.text
+ assert (await client.put(base + "/scopes", json={"scope": scope})).json() == opened.json()
+ assert (
+ await client.post(
+ base + "/scopes/reconcile",
+ json={
+ "scope": scope,
+ "expected": opened.json()["version"],
+ },
+ )
+ ).status_code == 409
+ original = snapshot(owner_id="owner", original={"body": "private-native-content नमस्ते"})
+ body = {
+ "page_id": str(uuid4()),
+ "scope": scope,
+ "expected": opened.json()["version"],
+ "snapshots": [original.model_dump(mode="json")],
+ "cursor": {"position": 1},
+ "final": True,
+ }
+ committed = await client.put(base + "/pages", json=body)
+ assert committed.status_code == 200, committed.text
+ assert (await client.put(base + "/pages", json=body)).json() == committed.json()
+ read = await client.post(base + "/scopes/read", json={"scope": scope})
+ assert read.json()["last_page"] == committed.json()
+ assert read.json()["cursor"] == {"position": 1}
+ assert "digest" not in read.text and "fence" not in read.text
+ reconciled = await client.post(
+ base + "/scopes/reconcile",
+ json={
+ "scope": scope,
+ "expected": committed.json()["version"],
+ },
+ )
+ assert reconciled.status_code == 200, reconciled.text
+ assert reconciled.json()["phase"] == "complete"
+ assert reconciled.json()["coverage"] == coverage
+ assert (await client.put(base + "/pages", json=body)).status_code == 409
+ assert (
+ await client.post(base + "/scopes/read", json={"scope": scope})
+ ).json() == reconciled.json()
+ async with database() as db:
+ row = await db.scalar(select(Entity).where(Entity.sync_id == source.sync_id))
+ assert row.record_revision == 1
+ assert row.source_payload["original"] == original.original
+ assert "scope-key" not in str(logs) and "private-native-content" not in str(logs)
+
+
+async def test_scope_and_page_reject_session_foreign_scope_and_client_fence(native_api):
+ client, ctx, source, _ = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/scope-key"
+ await client.put(base, json={"snapshot_id": "snapshot", "coverage": "bounded"})
+ ctx.is_api_key_auth = False
+ assert (
+ await client.put(base + "/scopes", json={"scope": {"record_type": "knowledge"}})
+ ).status_code == 403
+ ctx.is_api_key_auth = True
+ assert (
+ await client.put(base + "/scopes", json={"scope": {"record_type": "email"}})
+ ).status_code == 409
+ scope = await client.put(base + "/scopes", json={"scope": {"record_type": "knowledge"}})
+ body = {
+ "scope": {"record_type": "knowledge"},
+ "page_id": str(uuid4()),
+ "expected": scope.json()["version"],
+ "snapshots": [],
+ "final": True,
+ "fence": {"job_id": str(uuid4())},
+ }
+ assert (await client.put(base + "/pages", json=body)).status_code == 422
+ body.pop("fence")
+ assert (await client.put(base + "/pages", json=body)).status_code == 200
+ ctx.organization.id = uuid4()
+ assert (
+ await client.post(base + "/scopes/read", json={"scope": {"record_type": "knowledge"}})
+ ).status_code == 404
+
+
+@pytest.mark.parametrize("native_api", ["sessions"], indirect=True)
+async def test_sessions_http_parent_and_child_scopes_reused_by_new_import(database, native_api):
+ client, _, source, _ = native_api
+ parent = snapshot(
+ owner_id="owner",
+ identity=RecordIdentity(record_type="session", native_id="s"),
+ version=SessionVersion(revision=1, content_revision=1),
+ )
+ child = snapshot(
+ owner_id="owner",
+ identity=RecordIdentity(record_type="message", native_id="m", container_id="s"),
+ parent=parent.identity,
+ version=parent.version,
+ )
+ root_scope = {"record_type": "session"}
+ child_scope = {
+ "record_type": "message",
+ "container_id": "s",
+ "parent": parent.identity.model_dump(mode="json"),
+ }
+
+ async def upload(base, scope, item):
+ opened = await client.put(base + "/scopes", json={"scope": scope})
+ assert opened.status_code == 200, opened.text
+ captured = await client.put(
+ base + "/pages",
+ json={
+ "page_id": str(uuid4()),
+ "scope": scope,
+ "expected": opened.json()["version"],
+ "snapshots": [item.model_dump(mode="json")],
+ "final": True,
+ },
+ )
+ assert captured.status_code == 200, captured.text
+ reconciled = await client.post(
+ base + "/scopes/reconcile",
+ json={
+ "scope": scope,
+ "expected": captured.json()["version"],
+ },
+ )
+ assert reconciled.status_code == 200, reconciled.text
+ return captured.json()
+
+ base = f"/native/sources/{source.source_connection_id}/imports/first"
+ started = await client.put(base, json={"snapshot_id": "one", "coverage": "bounded"})
+ assert (await client.put(base + "/scopes", json={"scope": child_scope})).status_code == 409
+ await upload(base, root_scope, parent)
+ old_child = await upload(base, child_scope, child)
+ # Simulate external cancellation; dedicated cancellation HTTP is not implemented yet.
+ async with database() as db:
+ job = await db.get(SyncJob, started.json()["import_id"])
+ job.status = "cancelled"
+ await db.commit()
+ later = base.replace("/first", "/second")
+ assert (
+ await client.put(later, json={"snapshot_id": "two", "coverage": "bounded"})
+ ).status_code == 200
+ await upload(later, root_scope, parent)
+ new_child = await upload(later, child_scope, child)
+ assert new_child["version"]["sweep_id"] != old_child["version"]["sweep_id"]
+ assert new_child["changed"] == 0 and new_child["unchanged"] == 1
diff --git a/backend/airweave/domains/native_ingestion/tests/test_search.py b/backend/airweave/domains/native_ingestion/tests/test_search.py
new file mode 100644
index 0000000000..9e4ad00b8e
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_search.py
@@ -0,0 +1,96 @@
+"""Native publications use real SQL authority without a provider registry entry."""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+from uuid import uuid4
+
+import pytest
+from fastapi import HTTPException
+from sqlalchemy import select, update
+
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.tests.helpers import publish_prepared
+from airweave.domains.entities.canonical.tests.test_search_visibility import hit
+from airweave.domains.native_ingestion.tests.test_ingestion import bind, ingest
+from airweave.domains.native_ingestion.tests.test_projection import knowledge
+from airweave.domains.search.canonical_visibility import visible_results
+from airweave.domains.search.owned import OwnedSearchService
+from airweave.domains.search.owned_models import OwnedSearchRequest
+from airweave.domains.search.types import SearchResults
+from airweave.models.collection import Collection
+from airweave.models.source_connection import SourceConnection
+
+
+async def test_native_search_requires_current_publication_and_source_access(database, source):
+ """Index transport is synthetic; capture, scope checks and publication are real SQL."""
+ _, fence = source
+ await bind(database, fence)
+ await ingest(database, fence, knowledge())
+ async with database() as db:
+ collection = await db.scalar(
+ select(Collection)
+ .join(
+ SourceConnection, SourceConnection.readable_collection_id == Collection.readable_id
+ )
+ .where(
+ SourceConnection.sync_id == fence.sync_id,
+ Collection.organization_id == fence.organization_id,
+ )
+ )
+ assert collection is not None
+ store = CanonicalProjectionStore()
+ async with database() as db:
+ work = (await store.pending(db, fence.organization_id, fence.sync_id))[0]
+ generation = uuid4()
+ async with database() as db:
+ assert await publish_prepared(store, db, work, generation, 1, collection.id)
+ locator = ProjectionLocator(
+ record_id=work.record.id,
+ revision=work.record.revision,
+ pipeline_version=work.pipeline_version,
+ generation=generation,
+ part_index=0,
+ )
+ candidate = hit(fence, locator.encode())
+ candidate.airweave_system_metadata.source_name = "almanac"
+ registry = Mock()
+ registry.get.side_effect = AssertionError("Native content has no provider lifecycle")
+ executor = Mock()
+ executor.prepare_query = AsyncMock(return_value=None)
+ executor.execute = AsyncMock(return_value=SearchResults(results=[candidate]))
+ service = OwnedSearchService(executor, registry)
+ ctx = SimpleNamespace(organization=SimpleNamespace(id=fence.organization_id))
+ request = OwnedSearchRequest(query="original", sync_ids=[fence.sync_id], mode="keyword")
+ result = await service.search(database, ctx, request)
+ assert len(result.items) == 1
+ assert result.items[0].provider == "almanac"
+ assert result.items[0].identity.record_type == "knowledge"
+ async with database() as db:
+ assert await visible_results(
+ db, fence.organization_id, collection.readable_id, [candidate], registry
+ ) == [candidate]
+ legacy = candidate.model_copy(deep=True)
+ legacy.airweave_system_metadata.original_entity_id = "unversioned-native-id"
+ assert (
+ await visible_results(
+ db, fence.organization_id, collection.readable_id, [legacy], registry
+ )
+ == []
+ )
+
+ async def revoke_during_retrieval(**kwargs):
+ async with database() as db:
+ await db.execute(
+ update(SourceConnection)
+ .where(SourceConnection.sync_id == fence.sync_id)
+ .values(is_authenticated=False)
+ )
+ await db.commit()
+ return SearchResults(results=[candidate])
+
+ executor.execute.side_effect = revoke_during_retrieval
+ with pytest.raises(HTTPException) as error:
+ await service.search(database, ctx, request)
+ assert error.value.status_code == 404
+ registry.get.assert_not_called()
diff --git a/backend/airweave/domains/native_ingestion/tests/test_sources.py b/backend/airweave/domains/native_ingestion/tests/test_sources.py
new file mode 100644
index 0000000000..c1f73d8301
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_sources.py
@@ -0,0 +1,141 @@
+"""Actual SQL races, rollback and HTTP backend-authority gate for native sources."""
+
+import asyncio
+from types import SimpleNamespace
+from uuid import uuid4
+
+import httpx
+import pytest
+from fastapi import FastAPI, HTTPException
+from sqlalchemy import func, select
+
+from airweave.api import deps
+from airweave.api.v1.endpoints.native_sources import router
+from airweave.db.session import get_db
+from airweave.domains.native_ingestion.source_models import (
+ EnsureNativeSource,
+ native_source_id,
+ native_sync_id,
+)
+from airweave.domains.native_ingestion.source_service import NativeSources
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.models import Collection, Organization, SourceConnection, Sync
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+@pytest.fixture
+async def native_setup(database):
+ org, foreign = uuid4(), uuid4()
+ async with database() as db:
+ db.add_all([Organization(id=org, name="Native"), Organization(id=foreign, name="Other")])
+ metadata = VectorDbDeploymentMetadata(
+ dense_embedder="test", sparse_embedder="test", embedding_dimensions=3
+ )
+ db.add(metadata)
+ await db.flush()
+ for name, owner in (("one", org), ("two", org), ("foreign", foreign)):
+ db.add(
+ Collection(
+ name=name,
+ readable_id=name,
+ organization_id=owner,
+ vector_db_deployment_metadata_id=metadata.id,
+ )
+ )
+ await db.commit()
+ return org, foreign, NativeSources(NativeSourceStore())
+
+
+async def ensure(database, service, org, request):
+ async with database() as db:
+ return await service.ensure(db, org, request)
+
+
+async def test_concurrent_same_identity_and_immutable_collection(database, native_setup):
+ org, _, service = native_setup
+ request = EnsureNativeSource(owner_id="owner", dataset="knowledge", collection="one")
+ results = await asyncio.gather(*(ensure(database, service, org, request) for _ in range(4)))
+ assert all(result == results[0] for result in results)
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Sync)) == 1
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ with pytest.raises(HTTPException) as conflict:
+ await ensure(database, service, org, request.model_copy(update={"collection": "two"}))
+ assert conflict.value.status_code == 409
+ async with database() as db:
+ assert await service.get(db, org, results[0].source_connection_id) == results[0]
+
+
+async def test_creation_rollback_and_partial_state_fail_closed(database, native_setup, monkeypatch):
+ org, _, service = native_setup
+ request = EnsureNativeSource(owner_id="owner", dataset="sessions", collection="one")
+ original = service.store.require
+
+ async def fail(*args, **kwargs):
+ raise RuntimeError("synthetic failure after source insert")
+
+ monkeypatch.setattr(service.store, "require", fail)
+ with pytest.raises(RuntimeError, match="after source insert"):
+ await ensure(database, service, org, request)
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(Sync)) == 0
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 0
+ monkeypatch.setattr(service.store, "require", original)
+ async with database() as db:
+ db.add(
+ Sync(
+ id=native_sync_id(native_source_id(org, request.binding())),
+ organization_id=org,
+ name="Unexplained partial",
+ )
+ )
+ await db.commit()
+ with pytest.raises(HTTPException) as conflict:
+ await ensure(database, service, org, request)
+ assert conflict.value.status_code == 409
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 0
+
+
+async def test_owner_namespace_and_collection_organization_isolation(database, native_setup):
+ org, foreign, service = native_setup
+ request = EnsureNativeSource(owner_id="same-owner", dataset="knowledge", collection="one")
+ ours = await ensure(database, service, org, request)
+ theirs = await ensure(
+ database, service, foreign, request.model_copy(update={"collection": "foreign"})
+ )
+ assert ours.source_connection_id != theirs.source_connection_id
+ async with database() as db:
+ with pytest.raises(HTTPException) as missing:
+ await service.get(db, foreign, ours.source_connection_id)
+ assert missing.value.status_code == 404
+ with pytest.raises(HTTPException) as missing:
+ await ensure(database, service, foreign, request)
+ assert missing.value.status_code == 404
+
+
+async def test_http_session_cannot_ensure_or_read_but_backend_key_can(database, native_setup):
+ org, _, service = native_setup
+ ctx = SimpleNamespace(organization=SimpleNamespace(id=org), is_api_key_auth=False)
+ app = FastAPI()
+ app.include_router(router, prefix="/native/sources")
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: ctx
+ app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(native_sources=service)
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app), base_url="http://test"
+ ) as client:
+ body = {"owner_id": "owner", "dataset": "knowledge", "collection": "one"}
+ assert (await client.put("/native/sources", json=body)).status_code == 403
+ assert (await client.get(f"/native/sources/{uuid4()}")).status_code == 403
+ ctx.is_api_key_auth = True
+ created = await client.put("/native/sources", json=body)
+ assert created.status_code == 200, created.text
+ loaded = await client.get(f"/native/sources/{created.json()['source_connection_id']}")
+ assert loaded.status_code == 200 and loaded.json() == created.json()
+ assert (await client.put("/native/sources", json=body)).json() == created.json()
diff --git a/backend/airweave/domains/native_ingestion/tests/test_terminal_api.py b/backend/airweave/domains/native_ingestion/tests/test_terminal_api.py
new file mode 100644
index 0000000000..545c88ed7a
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_terminal_api.py
@@ -0,0 +1,124 @@
+"""Durable native terminal outcomes and fencing through actual HTTP/SQL."""
+
+# ruff: noqa: F811
+from uuid import uuid4
+
+from sqlalchemy import select
+
+from airweave.domains.native_ingestion.tests.test_import_api import native_api # noqa: F401
+from airweave.models.sync import Sync
+from airweave.models.sync_job import SyncJob
+
+
+async def test_completion_retry_terminal_scope_read_and_new_import(database, native_api):
+ client, ctx, source, _ = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/first"
+ body = {"snapshot_id": "snapshot", "coverage": "bounded"}
+ await client.put(base, json=body)
+ assert (await client.post(base + "/complete")).status_code == 409
+ scope = {"record_type": "knowledge"}
+ opened = (await client.put(base + "/scopes", json={"scope": scope})).json()
+ page = {
+ "scope": scope,
+ "page_id": str(uuid4()),
+ "expected": opened["version"],
+ "snapshots": [],
+ "final": True,
+ }
+ ack = (await client.put(base + "/pages", json=page)).json()
+ await client.post(base + "/scopes/reconcile", json={"scope": scope, "expected": ack["version"]})
+ done = await client.post(base + "/complete")
+ assert done.status_code == 200, done.text
+ summary = done.json()["summary"]
+ assert summary["capture_complete"] and summary["indexing"] == "not_verified"
+ assert summary["coverage"] == "bounded" and summary["completed_scopes"] == 1
+ assert (await client.post(base + "/complete")).json() == done.json()
+ assert (await client.post(base + "/scopes/read", json={"scope": scope})).status_code == 200
+ assert (await client.put(base + "/pages", json=page)).status_code == 409
+ newer = base.replace("/first", "/next")
+ await client.put(newer, json=body)
+ await client.put(newer + "/scopes", json={"scope": scope})
+ assert (await client.post(base + "/scopes/read", json={"scope": scope})).status_code == 409
+ assert (await client.post(base + "/complete")).json() == done.json()
+ assert (await client.get(base)).json() == done.json()
+ assert (await client.post(base + "/cancel")).json() == done.json()
+ async with database() as db:
+ sync = await db.get(Sync, source.sync_id)
+ assert (await db.get(SyncJob, sync.writer_job_id)).status == "running"
+ ctx.organization.id = uuid4()
+ assert (await client.post(base + "/complete")).status_code == 404
+ assert (await client.post(newer + "/cancel")).status_code == 404
+
+
+async def test_cancel_retries_do_not_cancel_new_writer(native_api, database):
+ client, _, source, _ = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/cancelled"
+ body = {"snapshot_id": "snapshot", "coverage": "complete"}
+ await client.put(base, json=body)
+ scope = {"record_type": "knowledge"}
+ opened = (await client.put(base + "/scopes", json={"scope": scope})).json()
+ cancelled = await client.post(base + "/cancel")
+ assert cancelled.status_code == 200
+ assert not cancelled.json()["summary"]["capture_complete"]
+ assert (await client.post(base + "/complete")).status_code == 409
+ assert (
+ await client.put(
+ base + "/pages",
+ json={
+ "scope": scope,
+ "page_id": str(uuid4()),
+ "expected": opened["version"],
+ "snapshots": [],
+ "final": True,
+ },
+ )
+ ).status_code == 409
+ newer = base.replace("/cancelled", "/next")
+ started = await client.put(newer, json=body)
+ assert started.status_code == 200
+ assert (await client.post(base + "/cancel")).json() == cancelled.json()
+ async with database() as db:
+ job = await db.scalar(select(SyncJob).where(SyncJob.id == started.json()["import_id"]))
+ assert job.status == "running"
+ assert (await db.get(Sync, source.sync_id)).writer_job_id == job.id
+
+
+async def test_terminal_summary_failure_rolls_back_cycle_completion(
+ native_api, database, monkeypatch
+):
+ import airweave.domains.native_ingestion.import_store as module
+
+ client, _, source, _ = native_api
+ base = f"/native/sources/{source.source_connection_id}/imports/rollback"
+ await client.put(base, json={"snapshot_id": "snapshot", "coverage": "bounded"})
+ scope = {"record_type": "knowledge"}
+ opened = (await client.put(base + "/scopes", json={"scope": scope})).json()
+ page = (
+ await client.put(
+ base + "/pages",
+ json={
+ "scope": scope,
+ "page_id": str(uuid4()),
+ "expected": opened["version"],
+ "snapshots": [],
+ "final": True,
+ },
+ )
+ ).json()
+ await client.post(
+ base + "/scopes/reconcile", json={"scope": scope, "expected": page["version"]}
+ )
+ original = module.NativeImportSummary
+
+ def fail(**kwargs):
+ raise RuntimeError("synthetic summary persistence failure")
+
+ monkeypatch.setattr(module, "NativeImportSummary", fail)
+ import pytest
+
+ with pytest.raises(RuntimeError, match="summary persistence"):
+ await client.post(base + "/complete")
+ assert (await client.get(base)).json()["status"] == "running"
+ monkeypatch.setattr(module, "NativeImportSummary", original)
+ # Successful retry also proves the cursor was not left complete by the failed transaction.
+ assert (await client.post(base + "/complete")).status_code == 200
diff --git a/backend/airweave/domains/native_ingestion/tests/test_transcript_reconciliation.py b/backend/airweave/domains/native_ingestion/tests/test_transcript_reconciliation.py
new file mode 100644
index 0000000000..4ab1498434
--- /dev/null
+++ b/backend/airweave/domains/native_ingestion/tests/test_transcript_reconciliation.py
@@ -0,0 +1,203 @@
+"""Explicit complete transcripts reconcile children within bounded observed roots."""
+
+# Imported pytest fixture names are resolved through function parameters.
+# ruff: noqa: F811
+
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import select
+
+from airweave.domains.entities.canonical.cycle_store import CycleConflict
+from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity
+from airweave.domains.entities.canonical.scan_store import ScanConflict
+from airweave.domains.entities.canonical.store import CanonicalRecordStore
+from airweave.domains.native_ingestion.errors import NativeAdmissionError
+from airweave.domains.native_ingestion.import_models import StartNativeImport
+from airweave.domains.native_ingestion.import_service import NativeImports
+from airweave.domains.native_ingestion.import_store import NativeImportStore
+from airweave.domains.native_ingestion.models import SessionVersion
+from airweave.domains.native_ingestion.page_models import CommitNativePage
+from airweave.domains.native_ingestion.scope_models import BeginNativeScope, ReconcileNativeScope
+from airweave.domains.native_ingestion.source_store import NativeSourceStore
+from airweave.domains.native_ingestion.tests.test_imports import native # noqa: F401
+from airweave.domains.native_ingestion.tests.test_ingestion import ingest, snapshot
+from airweave.models.entity import Entity
+
+pytestmark = pytest.mark.parametrize("native", ["sessions"], indirect=True)
+
+
+def imports():
+ return NativeImports(NativeImportStore(NativeSourceStore(), CanonicalRecordStore()))
+
+
+def session(native_id, revision=1):
+ return snapshot(
+ native_id,
+ owner_id="owner",
+ identity=RecordIdentity(record_type="session", native_id=native_id),
+ version=SessionVersion(revision=revision, content_revision=revision),
+ )
+
+
+def message(parent, native_id):
+ return snapshot(
+ native_id,
+ owner_id="owner",
+ identity=RecordIdentity(
+ record_type="message", native_id=native_id, container_id=parent.identity.native_id
+ ),
+ parent=parent.identity,
+ version=parent.version,
+ )
+
+
+def child_scope(parent):
+ return CompletedScope(
+ record_type="message", container_id=parent.identity.native_id, parent=parent.identity
+ )
+
+
+async def begin(database, native, key, scope):
+ async with database() as db:
+ return await imports().begin_scope(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ key,
+ BeginNativeScope(scope=scope),
+ )
+
+
+async def page(database, native, key, state, items, *, final=True):
+ async with database() as db:
+ return await imports().page(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ key,
+ CommitNativePage(
+ page_id=uuid4(),
+ scope=state.scope,
+ expected=state.version,
+ snapshots=items,
+ final=final,
+ ),
+ )
+
+
+async def reconcile(database, native, key, scope, version):
+ async with database() as db:
+ return await imports().reconcile_scope(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ key,
+ ReconcileNativeScope(scope=scope, expected=version),
+ )
+
+
+async def capture_scope(database, native, key, scope, items):
+ state = await begin(database, native, key, scope)
+ ack = await page(database, native, key, state, items)
+ return await reconcile(database, native, key, scope, ack.version)
+
+
+async def start(database, native, key, **intent):
+ async with database() as db:
+ return await imports().start(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ key,
+ StartNativeImport(snapshot_id=key, coverage="bounded", **intent),
+ )
+
+
+async def seed(database, native):
+ first, second = session("first"), session("second")
+ await start(database, native, "baseline")
+ await capture_scope(
+ database, native, "baseline", CompletedScope(record_type="session"), (first, second)
+ )
+ await capture_scope(
+ database,
+ native,
+ "baseline",
+ child_scope(first),
+ (message(first, "kept"), message(first, "omitted")),
+ )
+ await capture_scope(
+ database, native, "baseline", child_scope(second), (message(second, "unselected"),)
+ )
+ async with database() as db:
+ await imports().finish(db, native.organization_id, native.source_connection_id, "baseline")
+ return first, second
+
+
+async def rows(database, native):
+ async with database() as db:
+ values = await db.scalars(select(Entity).where(Entity.sync_id == native.sync_id))
+ return {item.native_id: item for item in values}
+
+
+@pytest.mark.parametrize("complete", [False, True])
+async def test_selected_transcript_policy_preserves_unselected_roots(database, native, complete):
+ await seed(database, native)
+ intent = {"transcript_coverage": "complete"} if complete else {}
+ started = await start(database, native, "selected", **intent)
+ assert await start(database, native, "selected", **intent) == started
+ parent = session("first", revision=2)
+ root = await capture_scope(
+ database, native, "selected", CompletedScope(record_type="session"), (parent,)
+ )
+ assert root.coverage == "bounded"
+ scope = await capture_scope(
+ database, native, "selected", child_scope(parent), (message(parent, "kept"),)
+ )
+ assert scope.coverage == ("complete" if complete else "bounded")
+ async with database() as db:
+ read = await imports().read_scope(
+ db,
+ native.organization_id,
+ native.source_connection_id,
+ "selected",
+ BeginNativeScope(scope=scope.scope),
+ )
+ assert read.coverage == scope.coverage
+ result = await imports().finish(
+ db, native.organization_id, native.source_connection_id, "selected"
+ )
+ assert result.summary.coverage == "bounded" and result.summary.capture_complete
+ records = await rows(database, native)
+ assert (records["omitted"].deleted_at is not None) == complete
+ assert records["omitted"].removal_reason == ("absent" if complete else None)
+ assert records["kept"].deleted_at is None
+ for name in ("second", "unselected"):
+ assert records[name].deleted_at is None and records[name].record_revision == 1
+
+
+async def test_incomplete_failed_and_parent_changed_scan_cannot_reconcile_absence(database, native):
+ old_parent, _ = await seed(database, native)
+ await start(database, native, "selected", transcript_coverage="complete")
+ parent = session("first", revision=2)
+ await capture_scope(
+ database, native, "selected", CompletedScope(record_type="session"), (parent,)
+ )
+ state = await begin(database, native, "selected", child_scope(parent))
+ ack = await page(database, native, "selected", state, (message(parent, "kept"),), final=False)
+ state = state.model_copy(update={"version": ack.version})
+ with pytest.raises(NativeAdmissionError, match="version must match"):
+ await page(database, native, "selected", state, (message(old_parent, "stale"),))
+ with pytest.raises(ScanConflict, match="fully collected"):
+ await reconcile(database, native, "selected", state.scope, ack.version)
+ async with database() as db:
+ saved = await imports().store.active(
+ db, native.organization_id, native.source_connection_id, "selected"
+ )
+ await ingest(database, saved.fence, session("first", revision=3))
+ with pytest.raises((CycleConflict, ScanConflict)):
+ await page(database, native, "selected", state, ())
+ records = await rows(database, native)
+ assert records["omitted"].deleted_at is None and records["omitted"].record_revision == 1
+ assert "stale" not in records
diff --git a/backend/airweave/domains/oauth/callback_service.py b/backend/airweave/domains/oauth/callback_service.py
index a682eda78b..7c889a1c86 100644
--- a/backend/airweave/domains/oauth/callback_service.py
+++ b/backend/airweave/domains/oauth/callback_service.py
@@ -531,6 +531,7 @@ async def _complete_connection_common( # noqa: C901
collection_id=collection.id,
collection_readable_id=collection.readable_id,
source_entry=source_entry,
+ source_config=validated_config,
schedule_config=schedule_config,
run_immediately=True,
ctx=ctx,
diff --git a/backend/airweave/domains/oauth/tests/test_callback_service.py b/backend/airweave/domains/oauth/tests/test_callback_service.py
index a2b1a762c1..91c0581ed6 100644
--- a/backend/airweave/domains/oauth/tests/test_callback_service.py
+++ b/backend/airweave/domains/oauth/tests/test_callback_service.py
@@ -25,8 +25,8 @@
from airweave.domains.organizations.fakes.repository import FakeOrganizationRepository
from airweave.domains.source_connections.fakes.repository import FakeSourceConnectionRepository
from airweave.domains.sources.exceptions import SourceNotFoundError, SourceValidationError
-from airweave.domains.syncs.jobs.fakes.repository import FakeSyncJobRepository
from airweave.domains.syncs.fakes.repository import FakeSyncRepository
+from airweave.domains.syncs.jobs.fakes.repository import FakeSyncJobRepository
from airweave.models.connection_init_session import ConnectionInitSession, ConnectionInitStatus
from airweave.models.organization import Organization
from airweave.models.source_connection import SourceConnection
@@ -1122,6 +1122,7 @@ async def test_triggers_workflow_when_pending_job_exists(self):
job_id = uuid4()
sync_job = SyncJob(
+ provisioning_generation=0,
id=job_id,
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1255,6 +1256,7 @@ async def test_running_job_skips_workflow(self):
from airweave.models.sync_job import SyncJob
sync_job = SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.RUNNING,
@@ -1316,6 +1318,7 @@ async def test_missing_connection_id_raises_value_error(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1407,6 +1410,7 @@ async def test_pending_job_with_connection_executes_event_payload_path(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1575,6 +1579,7 @@ async def test_no_collection_skips_workflow(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1743,6 +1748,7 @@ async def test_sync_deferred_when_claim_token_set(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1823,6 +1829,7 @@ async def test_sync_triggered_when_no_claim_token(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -1953,6 +1960,7 @@ async def test_two_step_contract_callback_defers_then_verify_triggers(self):
sync_id,
[
SyncJob(
+ provisioning_generation=0,
id=uuid4(),
sync_id=sync_id,
status=SyncJobStatus.PENDING,
@@ -2241,9 +2249,7 @@ async def test_complete_connection_common_activates_sync(self):
db = AsyncMock()
ctx = _ctx()
- with patch(
- "airweave.domains.oauth.callback_service.UnitOfWork"
- ) as mock_uow_cls:
+ with patch("airweave.domains.oauth.callback_service.UnitOfWork") as mock_uow_cls:
mock_uow = AsyncMock()
mock_uow.session = AsyncMock()
mock_uow.__aenter__ = AsyncMock(return_value=mock_uow)
diff --git a/backend/airweave/domains/ocr/fallback.py b/backend/airweave/domains/ocr/fallback.py
index 1afa9f407e..7392b00a4d 100644
--- a/backend/airweave/domains/ocr/fallback.py
+++ b/backend/airweave/domains/ocr/fallback.py
@@ -54,9 +54,9 @@ def __init__(
async def convert_batch(self, file_paths: List[str]) -> Dict[str, Optional[str]]:
"""Convert files to markdown, trying providers in order.
- Skips circuit-broken providers. On failure, records the failure
- and tries the next provider. If all providers are exhausted,
- returns None for every file.
+ Retains successful results and sends only unresolved files to later providers.
+ Skips circuit-broken providers. Unresolved files remain None when the chain
+ is exhausted; successful files are never processed or charged twice.
Args:
file_paths: Local file paths to convert.
@@ -64,29 +64,36 @@ async def convert_batch(self, file_paths: List[str]) -> Dict[str, Optional[str]]
Returns:
Mapping of ``file_path -> markdown`` (``None`` on failure).
"""
+ combined = {path: None for path in file_paths}
for provider_key, provider in self._providers:
+ pending = [path for path, result in combined.items() if result is None]
+ if not pending:
+ break
if not await self._circuit_breaker.is_available(provider_key):
logger.info(f"[FallbackOCR] Skipping '{provider_key}' (circuit-broken)")
continue
try:
- results = await provider.convert_batch(file_paths)
-
- if results and all(v is None for v in results.values()):
+ results = await provider.convert_batch(pending)
+ succeeded = {
+ path: results[path] for path in pending if results.get(path) is not None
+ }
+ if succeeded:
+ combined.update(succeeded)
+ await self._circuit_breaker.record_success(provider_key)
+ else:
await self._circuit_breaker.record_failure(provider_key)
logger.warning(
- f"[FallbackOCR] '{provider_key}' returned all-None results "
- f"({len(results)} files), trying next provider"
+ f"[FallbackOCR] '{provider_key}' produced no usable results "
+ f"for {len(pending)} files, trying next provider"
)
- continue
-
- await self._circuit_breaker.record_success(provider_key)
- return results
except Exception as exc:
await self._circuit_breaker.record_failure(provider_key)
logger.warning(
- f"[FallbackOCR] '{provider_key}' failed ({exc}), trying next provider"
+ f"[FallbackOCR] '{provider_key}' failed ({type(exc).__name__}), "
+ "trying next provider"
)
- logger.error("[FallbackOCR] All OCR providers unavailable or failed")
- return {path: None for path in file_paths}
+ if any(result is None for result in combined.values()):
+ logger.warning("[FallbackOCR] Some files have no successful OCR result")
+ return combined
diff --git a/backend/airweave/domains/ocr/local.py b/backend/airweave/domains/ocr/local.py
new file mode 100644
index 0000000000..db3ca08808
--- /dev/null
+++ b/backend/airweave/domains/ocr/local.py
@@ -0,0 +1,109 @@
+"""Optional local OCR provider with one bounded, disposable process per input."""
+
+import asyncio
+import re
+import sys
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from airweave.domains.ocr.local_worker import (
+ MAX_INPUT_BYTES,
+ MAX_OUTPUT_BYTES,
+ SUPPORTED_EXTENSIONS,
+)
+
+OCR_TIMEOUT_SECONDS = 120
+WORKER_PATH = Path(__file__).with_name("local_worker.py")
+
+
+async def _stop(process: asyncio.subprocess.Process) -> None:
+ """Reap the worker after killing it, including a process that exited concurrently."""
+ if process.returncode is None:
+ try:
+ process.kill()
+ except ProcessLookupError:
+ pass
+ await process.wait()
+
+
+class LocalOcrProvider:
+ """Use explicitly provisioned Tesseract language files without runtime downloads."""
+
+ def __init__(self, tessdata_path: str | Path, languages: tuple[str, ...] = ("eng",)) -> None:
+ """Fail configuration early for absent models or unsafe language selectors."""
+ directory = Path(tessdata_path).resolve()
+ if not directory.is_dir() or not languages:
+ raise ValueError("Local OCR requires a tessdata directory and languages")
+ for language in languages:
+ if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,63}", language):
+ raise ValueError("Invalid local OCR language selector")
+ model = directory / f"{language}.traineddata"
+ if not model.is_file() or model.stat().st_size == 0:
+ raise ValueError("Local OCR language model is unavailable")
+ try:
+ with model.open("rb") as handle:
+ handle.read(1)
+ except OSError:
+ raise ValueError("Local OCR language model is unreadable") from None
+ self._tessdata = directory
+ self._languages = "+".join(languages)
+
+ async def convert_batch(self, file_paths: list[str]) -> dict[str, str | None]:
+ """Process sequentially; each failed file remains available to the fallback chain."""
+ results = {}
+ for path in file_paths:
+ results[path] = await self._convert(path)
+ return results
+
+ async def _convert(self, file_path: str) -> str | None:
+ try:
+ path = Path(file_path).resolve()
+ if (
+ path.suffix.lower() not in SUPPORTED_EXTENSIONS
+ or not path.is_file()
+ or path.stat().st_size > MAX_INPUT_BYTES
+ ):
+ return None
+ except (OSError, RuntimeError):
+ return None
+ with TemporaryDirectory(prefix="airweave-local-ocr-") as directory:
+ output = Path(directory) / "text.txt"
+ spawn = asyncio.create_task(
+ asyncio.create_subprocess_exec(
+ sys.executable,
+ "-I",
+ str(WORKER_PATH),
+ "--input",
+ str(path),
+ "--output",
+ str(output),
+ "--tessdata",
+ str(self._tessdata),
+ "--languages",
+ self._languages,
+ stdin=asyncio.subprocess.DEVNULL,
+ stdout=asyncio.subprocess.DEVNULL,
+ stderr=asyncio.subprocess.DEVNULL,
+ )
+ )
+ try:
+ process = await asyncio.shield(spawn)
+ except asyncio.CancelledError:
+ process = await spawn
+ await _stop(process)
+ raise
+ except OSError:
+ return None
+ try:
+ await asyncio.wait_for(process.wait(), timeout=OCR_TIMEOUT_SECONDS)
+ if process.returncode != 0:
+ return None
+ with output.open("rb") as handle:
+ data = handle.read(MAX_OUTPUT_BYTES + 1)
+ if len(data) > MAX_OUTPUT_BYTES:
+ return None
+ return data.decode("utf-8").strip() or None
+ except (OSError, UnicodeDecodeError, TimeoutError):
+ return None
+ finally:
+ await _stop(process)
diff --git a/backend/airweave/domains/ocr/local_worker.py b/backend/airweave/domains/ocr/local_worker.py
new file mode 100644
index 0000000000..46c0c894fd
--- /dev/null
+++ b/backend/airweave/domains/ocr/local_worker.py
@@ -0,0 +1,89 @@
+"""Isolated bounded PyMuPDF/Tesseract process; no application imports or downloads."""
+
+import argparse
+import math
+from pathlib import Path
+
+DPI = 150
+# Match storage/limits.py without importing the application in this isolated worker.
+MAX_INPUT_BYTES = 200 * 1024 * 1024
+MAX_PAGES = 200
+MAX_PAGE_PIXELS = 20_000_000
+MAX_OUTPUT_BYTES = 8 * 1024 * 1024
+SUPPORTED_EXTENSIONS = frozenset(
+ {".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tif", ".tiff"}
+)
+
+
+def _check_page(width: float, height: float) -> None:
+ """Reject excessive render surfaces before constructing any page pixmap."""
+ if not all(math.isfinite(value) and value > 0 for value in (width, height)):
+ raise ValueError("Invalid page dimensions")
+ pixels = math.ceil(width * DPI / 72) * math.ceil(height * DPI / 72)
+ if pixels > MAX_PAGE_PIXELS:
+ raise ValueError("Page render limit exceeded")
+
+
+def _require_single_frame(path: Path) -> None:
+ """PyMuPDF image opening must not silently omit later container frames."""
+ from PIL import Image
+
+ with Image.open(path) as image:
+ try:
+ image.seek(1)
+ except EOFError:
+ return
+ raise ValueError("Multi-frame images are not supported")
+
+
+def extract(path: Path, tessdata: Path, languages: str) -> bytes:
+ """Return complete UTF-8 text or fail; partial results are never published."""
+ import pymupdf
+
+ if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
+ raise ValueError("Unsupported input format")
+ if path.stat().st_size > MAX_INPUT_BYTES:
+ raise ValueError("Input size limit exceeded")
+ with pymupdf.open(path) as document:
+ if document.needs_pass or not 1 <= document.page_count <= MAX_PAGES:
+ raise ValueError("Unsupported document bounds or encryption")
+ if not document.is_pdf:
+ _require_single_frame(path)
+ for page in document:
+ _check_page(page.rect.width, page.rect.height)
+ output = bytearray()
+ for page in document:
+ textpage = page.get_textpage_ocr(
+ language=languages, dpi=DPI, full=True, tessdata=str(tessdata)
+ )
+ text = page.get_text("text", textpage=textpage).strip().encode("utf-8")
+ separator = b"\n\n" if output and text else b""
+ if len(output) + len(separator) + len(text) > MAX_OUTPUT_BYTES:
+ raise ValueError("OCR output limit exceeded")
+ output.extend(separator)
+ output.extend(text)
+ if not output:
+ raise ValueError("OCR produced no text")
+ return bytes(output)
+
+
+def main() -> int:
+ """Only a successfully completed extraction may create the private result file."""
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--input", required=True, type=Path)
+ parser.add_argument("--output", required=True, type=Path)
+ parser.add_argument("--tessdata", required=True, type=Path)
+ parser.add_argument("--languages", required=True)
+ args = parser.parse_args()
+ try:
+ content = extract(args.input, args.tessdata, args.languages)
+ with args.output.open("xb") as destination:
+ destination.write(content)
+ return 0
+ except Exception:
+ # Provider/native exception messages may contain document text or private paths.
+ return 1
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/backend/airweave/domains/ocr/tests/test_fallback.py b/backend/airweave/domains/ocr/tests/test_fallback.py
index dcafe63ebb..ab76d22b99 100644
--- a/backend/airweave/domains/ocr/tests/test_fallback.py
+++ b/backend/airweave/domains/ocr/tests/test_fallback.py
@@ -137,8 +137,8 @@ async def test_fallback_ocr(case: Case):
@pytest.mark.asyncio
-async def test_partial_none_does_not_trigger_fallback():
- """When a provider returns some None and some real results, don't fall back."""
+async def test_partial_failure_retries_only_unresolved_files():
+ """Keep successful local work while later providers handle unsupported files."""
primary = FakeOcrProvider(
overrides={"/tmp/a.pdf": "# A", "/tmp/b.pdf": None},
default_markdown=None,
@@ -153,10 +153,10 @@ async def test_partial_none_does_not_trigger_fallback():
result = await fallback.convert_batch(["/tmp/a.pdf", "/tmp/b.pdf"])
assert result["/tmp/a.pdf"] == "# A"
- assert result["/tmp/b.pdf"] is None
- assert cb.successes == ["primary"]
+ assert result["/tmp/b.pdf"] == "# Secondary"
+ assert cb.successes == ["primary", "secondary"]
assert cb.failures == []
- assert secondary.call_count == 0
+ assert secondary.calls == [["/tmp/b.pdf"]]
def test_rejects_empty_providers():
diff --git a/backend/airweave/domains/ocr/tests/test_local.py b/backend/airweave/domains/ocr/tests/test_local.py
new file mode 100644
index 0000000000..fb79815556
--- /dev/null
+++ b/backend/airweave/domains/ocr/tests/test_local.py
@@ -0,0 +1,182 @@
+"""Local OCR isolation, bounded input/output, and optional real model qualification."""
+
+import asyncio
+import os
+from pathlib import Path
+
+import pymupdf
+import pytest
+from PIL import Image
+
+from airweave.domains.ocr import local, local_worker
+from airweave.domains.ocr.local import LocalOcrProvider
+from airweave.domains.ocr.protocols import OcrProvider
+
+
+@pytest.fixture
+def models(tmp_path):
+ directory = tmp_path / "tessdata"
+ directory.mkdir()
+ (directory / "eng.traineddata").write_bytes(b"synthetic model for process tests")
+ return directory
+
+
+def pdf(path, *, pages=1, size=300, text="Local OCR test"):
+ with pymupdf.open() as document:
+ for _ in range(pages):
+ page = document.new_page(width=size, height=size)
+ page.insert_text((30, 50), text)
+ document.save(path)
+ return path
+
+
+def test_configuration_requires_explicit_safe_models(models):
+ assert isinstance(LocalOcrProvider(models), OcrProvider)
+ for languages in ((), ("../eng",), ("eng+hin",), ("hin",)):
+ with pytest.raises(ValueError):
+ LocalOcrProvider(models, languages=languages)
+ with pytest.raises(ValueError):
+ LocalOcrProvider(models / "missing")
+
+
+async def test_unsupported_input_does_not_start_process(models, tmp_path, monkeypatch):
+ async def forbidden(*args, **kwargs):
+ raise AssertionError("Unsupported input started OCR")
+
+ monkeypatch.setattr(asyncio, "create_subprocess_exec", forbidden)
+ path = tmp_path / "private.txt"
+ path.write_text("not an OCR image")
+ large = tmp_path / "large.pdf"
+ with large.open("wb") as handle:
+ handle.truncate(local_worker.MAX_INPUT_BYTES + 1)
+ assert await LocalOcrProvider(models).convert_batch([str(path), str(large)]) == {
+ str(path): None,
+ str(large): None,
+ }
+ with pytest.raises(ValueError, match="Input size"):
+ local_worker.extract(large, models, "eng")
+
+
+@pytest.mark.parametrize("cancel", [False, True], ids=["timeout", "cancellation"])
+async def test_worker_is_killed_reaped_and_scratch_removed(models, tmp_path, monkeypatch, cancel):
+ worker = tmp_path / "sleeper.py"
+ worker.write_text("import time\ntime.sleep(60)\n")
+ monkeypatch.setattr(local, "WORKER_PATH", worker)
+ monkeypatch.setattr(local, "OCR_TIMEOUT_SECONDS", 0.1 if not cancel else 120)
+ actual_spawn = asyncio.create_subprocess_exec
+ started = asyncio.Event()
+ processes, outputs = [], []
+
+ async def spawn(*args, **kwargs):
+ assert kwargs["stderr"] == asyncio.subprocess.DEVNULL
+ assert kwargs["stdout"] == asyncio.subprocess.DEVNULL
+ assert args[1] == "-I"
+ process = await actual_spawn(*args, **kwargs)
+ processes.append(process)
+ outputs.append(Path(args[args.index("--output") + 1]))
+ started.set()
+ return process
+
+ monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
+ path = pdf(tmp_path / "private.pdf")
+ task = asyncio.create_task(LocalOcrProvider(models).convert_batch([str(path)]))
+ await started.wait()
+ if cancel:
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ else:
+ assert await task == {str(path): None}
+ assert processes[0].returncode is not None
+ assert not outputs[0].parent.exists()
+
+
+def test_render_limits_and_encryption_fail_before_ocr(tmp_path, models, monkeypatch):
+ def forbidden(*args, **kwargs):
+ raise AssertionError("Unsafe page reached rendering")
+
+ monkeypatch.setattr(pymupdf.Page, "get_textpage_ocr", forbidden)
+ paths = [pdf(tmp_path / "too-many.pdf", pages=201), pdf(tmp_path / "too-big.pdf", size=10000)]
+ encrypted = tmp_path / "encrypted.pdf"
+ with pymupdf.open() as document:
+ document.new_page()
+ document.save(
+ encrypted, encryption=pymupdf.PDF_ENCRYPT_AES_256, owner_pw="owner", user_pw="secret"
+ )
+ paths.append(encrypted)
+ for path in paths:
+ with pytest.raises(ValueError):
+ local_worker.extract(path, models, "eng")
+
+
+def test_multiframe_images_are_not_silently_truncated(tmp_path, models):
+ path = tmp_path / "frames.tiff"
+ Image.new("RGB", (10, 10)).save(path, save_all=True, append_images=[Image.new("RGB", (10, 10))])
+ with pytest.raises(ValueError, match="Multi-frame"):
+ local_worker.extract(path, models, "eng")
+
+
+def test_output_limit_rejects_whole_result(tmp_path, models, monkeypatch):
+ def synthetic_ocr(page, *, full, dpi, **kwargs):
+ assert full is True and dpi == 150
+ return page.get_textpage()
+
+ monkeypatch.setattr(pymupdf.Page, "get_textpage_ocr", synthetic_ocr)
+ monkeypatch.setattr(local_worker, "MAX_OUTPUT_BYTES", 3)
+ with pytest.raises(ValueError, match="output limit"):
+ local_worker.extract(pdf(tmp_path / "text.pdf"), models, "eng")
+
+
+async def test_real_local_ocr_with_explicit_models(tmp_path):
+ directory = os.environ.get("LOCAL_OCR_TEST_TESSDATA")
+ if not directory:
+ pytest.skip("Set LOCAL_OCR_TEST_TESSDATA to explicitly qualify real local OCR")
+ path = tmp_path / "scan.pdf"
+ with pymupdf.open() as born_digital:
+ page = born_digital.new_page(width=400, height=150)
+ page.insert_text((30, 65), "LOCAL OCR ORIGINAL 12345", fontsize=18)
+ raster = page.get_pixmap(matrix=pymupdf.Matrix(2, 2)).tobytes("png")
+ with pymupdf.open() as scanned:
+ page = scanned.new_page(width=400, height=150)
+ page.insert_image(page.rect, stream=raster)
+ scanned.save(path)
+ image_path = tmp_path / "scan.png"
+ image_path.write_bytes(raster)
+ tiff_path = tmp_path / "scan.tiff"
+ with Image.open(image_path) as image:
+ image.save(tiff_path)
+ result = await LocalOcrProvider(directory).convert_batch(
+ [str(path), str(image_path), str(tiff_path)]
+ )
+ assert "LOCAL OCR ORIGINAL 12345" in result[str(path)]
+ assert "LOCAL OCR ORIGINAL 12345" in result[str(image_path)]
+ assert "LOCAL OCR ORIGINAL 12345" in result[str(tiff_path)]
+
+
+async def test_failed_file_does_not_discard_other_results(models, tmp_path, monkeypatch):
+ worker = tmp_path / "controlled.py"
+ worker.write_text(
+ "import pathlib, sys\n"
+ "args = dict(zip(sys.argv[1::2], sys.argv[2::2]))\n"
+ "if pathlib.Path(args['--input']).name == 'broken.pdf':\n"
+ " print('synthetic private diagnostic', file=sys.stderr)\n"
+ " raise SystemExit(1)\n"
+ "pathlib.Path(args['--output']).write_text('Recognized text')\n"
+ )
+ monkeypatch.setattr(local, "WORKER_PATH", worker)
+ paths = [pdf(tmp_path / name) for name in ("first.pdf", "broken.pdf", "last.pdf")]
+ actual_spawn = asyncio.create_subprocess_exec
+ processes = []
+
+ async def spawn(*args, **kwargs):
+ assert all(process.returncode is not None for process in processes)
+ process = await actual_spawn(*args, **kwargs)
+ processes.append(process)
+ return process
+
+ monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
+ assert await LocalOcrProvider(models).convert_batch(list(map(str, paths))) == {
+ str(paths[0]): "Recognized text",
+ str(paths[1]): None,
+ str(paths[2]): "Recognized text",
+ }
diff --git a/backend/airweave/domains/owned_provisioning/__init__.py b/backend/airweave/domains/owned_provisioning/__init__.py
new file mode 100644
index 0000000000..4f048b3711
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/__init__.py
@@ -0,0 +1 @@
+"""Durable, idempotent desired state for Almanac-owned source connections."""
diff --git a/backend/airweave/domains/owned_provisioning/guard.py b/backend/airweave/domains/owned_provisioning/guard.py
new file mode 100644
index 0000000000..4193a9961c
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/guard.py
@@ -0,0 +1,48 @@
+"""Legacy editing must not bypass the generation-controlled account lifecycle."""
+
+from uuid import UUID
+
+from fastapi import HTTPException
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.models.owned_provisioning import OwnedProvisioning
+from airweave.models.source_connection import SourceConnection
+
+
+def require_provider_source(short_name: str) -> None:
+ """Native imports cannot use provider configuration, OAuth or pull-sync routes."""
+ if short_name == "almanac":
+ raise HTTPException(status_code=409, detail="Manage this source through its native import")
+
+
+async def require_unmanaged_source(
+ db: AsyncSession, source_id: UUID, organization: UUID, *, short_name: str
+) -> None:
+ """Owned source identity/configuration is writable only through ensure."""
+ require_provider_source(short_name)
+ owned = await db.scalar(
+ select(OwnedProvisioning.id).where(
+ OwnedProvisioning.source_connection_id == source_id,
+ OwnedProvisioning.organization_id == organization,
+ )
+ )
+ if owned is not None:
+ raise HTTPException(
+ status_code=409, detail="Manage this source through its Almanac account"
+ )
+
+
+async def require_provider_sync(db: AsyncSession, sync_id: UUID, organization: UUID) -> None:
+ """Check the job's actual sync owner, not a caller-supplied source URL."""
+ native = await db.scalar(
+ select(SourceConnection.id)
+ .where(
+ SourceConnection.sync_id == sync_id,
+ SourceConnection.organization_id == organization,
+ SourceConnection.short_name == "almanac",
+ )
+ .limit(1)
+ )
+ if native is not None:
+ raise HTTPException(status_code=409, detail="Cancel this job through its native import")
diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py
new file mode 100644
index 0000000000..05afea45a3
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/models.py
@@ -0,0 +1,246 @@
+"""Backend-only desired account capture contract; credentials remain in Composio."""
+
+from typing import Literal
+from uuid import UUID
+
+from croniter import croniter
+from pydantic import (
+ BaseModel,
+ ConfigDict,
+ Field,
+ JsonValue,
+ ValidationInfo,
+ field_validator,
+ model_validator,
+)
+
+
+class ManagedSource(BaseModel):
+ """Native identity is supplied only by Almanac's verified account authority."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ provider: Literal[
+ "gmail",
+ "google_calendar",
+ "google_drive",
+ "slack",
+ "outlook_mail",
+ "outlook_calendar",
+ "stripe",
+ "linear",
+ "attio",
+ "github",
+ "notion",
+ ]
+ expected_identity: str = Field(min_length=1, max_length=512)
+ expected_user_identity: str | None = Field(default=None, min_length=1, max_length=512)
+ collection: str = Field(min_length=1, max_length=255)
+ auth_provider: str = Field(min_length=1, max_length=255)
+ connected_account_id: str = Field(min_length=1, max_length=255)
+ auth_config_id: str = Field(min_length=1, max_length=255)
+ user_id: str = Field(min_length=1, max_length=255)
+ config: dict[str, JsonValue] = Field(default_factory=dict)
+ cron: str = Field(min_length=1, max_length=100)
+
+ @model_validator(mode="after")
+ def native_principal(self):
+ """Workspace-scoped grants retain the native member or bot, not broker user_id."""
+ if (self.provider in {"slack", "notion"}) != (self.expected_user_identity is not None):
+ raise ValueError("Slack and Notion require an expected native user or bot identity")
+ if self.provider in {"stripe", "linear", "attio", "github", "notion"}:
+ self.source_config() # Validate native scope before admission, not during delivery.
+ return self
+
+ @field_validator("expected_identity")
+ @classmethod
+ def canonical_identity(cls, value: str, info: ValidationInfo) -> str:
+ """Native IDs have one canonical spelling across admission and reconnect."""
+ if info.data.get("provider") == "github":
+ if not value.isascii() or not value.isdecimal() or int(value) <= 0:
+ raise ValueError("GitHub identity must be a positive numeric user ID")
+ return str(int(value))
+ if info.data.get("provider") in {"linear", "attio", "notion"}:
+ return str(UUID(value))
+ return value
+
+ @field_validator("expected_user_identity")
+ @classmethod
+ def canonical_member(cls, value: str | None, info: ValidationInfo) -> str | None:
+ """Notion bot UUID is part of the verified native account pair."""
+ if value is not None and info.data.get("provider") == "notion":
+ return str(UUID(value))
+ return value
+
+ @field_validator("cron")
+ @classmethod
+ def valid_schedule(cls, value: str) -> str:
+ """Use the existing scheduler dependency rather than inventing cron parsing."""
+ if len(value.split()) != 5 or not croniter.is_valid(value):
+ raise ValueError("Use a valid five-field cron schedule")
+ return value
+
+ def source_config(self) -> dict[str, JsonValue]:
+ """Expected identity cannot be overridden inside unstructured provider config."""
+ if self.provider == "notion":
+ from airweave.platform.configs.config import NotionConfig
+
+ return NotionConfig.model_validate(
+ {
+ **self.config,
+ "expected_workspace_id": self.expected_identity,
+ "expected_bot_id": self.expected_user_identity,
+ }
+ ).model_dump(mode="json")
+ if self.provider == "github":
+ from airweave.platform.configs.config import GitHubConfig
+
+ parsed = GitHubConfig.model_validate(
+ {**self.config, "expected_user_id": int(self.expected_identity)}
+ )
+ value = parsed.model_dump(mode="json")
+ value["repositories"] = sorted(
+ value["repositories"], key=lambda item: item["repository_id"]
+ )
+ return value
+ if self.provider in {"linear", "attio"}:
+ from airweave.platform.configs.config import AttioConfig, LinearConfig
+
+ config_type = LinearConfig if self.provider == "linear" else AttioConfig
+ parsed = config_type.model_validate(
+ {**self.config, "workspace_id": self.expected_identity}
+ )
+ value = parsed.model_dump(mode="json")
+ if self.provider == "linear":
+ value["team_ids"] = sorted(value["team_ids"])
+ return value
+ if self.provider == "stripe":
+ from airweave.platform.configs.config import StripeCaptureConfig, StripeConfig
+
+ binding = StripeCaptureConfig.model_validate(
+ {**self.config, "expected_account_id": self.expected_identity}
+ )
+ return StripeConfig(original_capture=binding).model_dump(mode="json")
+ if self.provider in {"outlook_mail", "outlook_calendar"}:
+ return {
+ **self.config,
+ "expected_principal_id": self.expected_identity,
+ "capture_originals": True,
+ }
+ if self.provider == "slack":
+ return {
+ **self.config,
+ "expected_team_id": self.expected_identity,
+ "expected_user_id": self.expected_user_identity,
+ }
+ field = {
+ "gmail": "expected_mailbox",
+ "google_calendar": "expected_primary_calendar_id",
+ "google_drive": "expected_permission_id",
+ }[self.provider]
+ return {**self.config, field: self.expected_identity}
+
+ def auth_config(self) -> dict[str, str]:
+ """Selectors contain no provider tokens or API keys."""
+ return {
+ "account_id": self.connected_account_id,
+ "auth_config_id": self.auth_config_id,
+ "user_id": self.user_id,
+ }
+
+
+class EnsureSource(BaseModel):
+ """Monotonic account authority request, safe for exact retries."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ generation: int = Field(ge=1, strict=True)
+ state: Literal["active", "paused", "disconnected"]
+ source: ManagedSource | None = None
+
+ @model_validator(mode="after")
+ def active_source(self) -> "EnsureSource":
+ """Inactive requests never replace identity or credentials."""
+ if self.state == "active" and self.source is None:
+ raise ValueError("An active account requires its verified source specification")
+ if self.state != "active" and self.source is not None:
+ raise ValueError("Inactive state does not accept new credentials or identity")
+ return self
+
+
+class ProvisionedSource(BaseModel):
+ """Committed identifiers and delivery state, not proof of capture completion."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ account_id: UUID
+ organization_id: UUID
+ generation: int
+ observed_generation: int
+ state: Literal["pending", "ready", "paused", "disconnected"]
+ source_connection_id: UUID | None
+ sync_id: UUID | None
+ expected_identity: str | None
+ expected_user_identity: str | None = None
+
+
+def _workspace_principal(provider: str, config: dict) -> str:
+ """Both workspace sources retain canonical UUIDs in their validated native config."""
+ from airweave.platform.configs.config import AttioConfig, LinearConfig
+
+ schema = LinearConfig if provider == "linear" else AttioConfig
+ return str(schema.model_validate(config).workspace_id)
+
+
+def native_principal(provider: str, config: dict) -> tuple[str, str | None]:
+ """Read the original attested principal from the protected source config."""
+ if provider == "notion":
+ from airweave.platform.configs.config import NotionConfig
+
+ notion = NotionConfig.model_validate(config)
+ return str(notion.expected_workspace_id), str(notion.expected_bot_id)
+ if provider == "github":
+ from airweave.platform.configs.config import GitHubConfig
+
+ return str(GitHubConfig.model_validate(config).expected_user_id), None
+ if provider in {"linear", "attio"}:
+ return _workspace_principal(provider, config), None
+ return _account_principal(provider, config)
+
+
+def _account_principal(provider: str, config: dict) -> tuple[str, str | None]:
+ """Validate the established provider-specific account identity contracts."""
+ from airweave.platform.configs.config import (
+ GmailConfig,
+ GoogleCalendarConfig,
+ GoogleDriveConfig,
+ OutlookCalendarConfig,
+ OutlookMailConfig,
+ SlackConfig,
+ StripeConfig,
+ )
+
+ user = None
+ match provider:
+ case "gmail":
+ identity = GmailConfig.model_validate(config).expected_mailbox
+ case "google_calendar":
+ identity = GoogleCalendarConfig.model_validate(config).expected_primary_calendar_id
+ case "google_drive":
+ identity = GoogleDriveConfig.model_validate(config).expected_permission_id
+ case "outlook_mail" | "outlook_calendar":
+ config_type = OutlookMailConfig if provider == "outlook_mail" else OutlookCalendarConfig
+ outlook = config_type.model_validate(config)
+ if not outlook.capture_originals:
+ raise ValueError("Owned Outlook source must capture originals")
+ identity = outlook.expected_principal_id
+ case "stripe":
+ stripe = StripeConfig.model_validate(config)
+ if stripe.original_capture is None:
+ raise ValueError("Owned Stripe source must capture originals")
+ identity = stripe.original_capture.expected_account_id
+ case "slack":
+ slack = SlackConfig.model_validate(config)
+ identity, user = slack.expected_team_id, slack.expected_user_id
+ case _:
+ raise ValueError("Unsupported owned source provider")
+ if identity is None:
+ raise ValueError("Owned source has no trusted native identity")
+ return identity, user
diff --git a/backend/airweave/domains/owned_provisioning/service.py b/backend/airweave/domains/owned_provisioning/service.py
new file mode 100644
index 0000000000..c2a7496cdc
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/service.py
@@ -0,0 +1,223 @@
+"""Reconcile committed account intent using existing native validation and Temporal."""
+
+import asyncio
+from uuid import UUID
+
+from fastapi import HTTPException
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from temporalio.exceptions import WorkflowAlreadyStartedError
+
+from airweave import schemas
+from airweave.api.context import ApiContext
+from airweave.core.datetime_utils import utc_now_naive
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.owned_provisioning.models import (
+ EnsureSource,
+ ProvisionedSource,
+ native_principal,
+)
+from airweave.domains.owned_provisioning.store import ProvisioningStore, locked_intent
+from airweave.domains.sources.protocols import SourceLifecycleServiceProtocol
+from airweave.domains.syncs.jobs.protocols import SyncJobRepositoryProtocol
+from airweave.domains.syncs.protocols import SyncRepositoryProtocol
+from airweave.domains.temporal.protocols import (
+ TemporalScheduleServiceProtocol,
+ TemporalWorkflowServiceProtocol,
+)
+from airweave.models.collection import Collection
+from airweave.models.connection import Connection
+from airweave.models.owned_provisioning import OwnedProvisioning
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+from airweave.models.sync_job import SyncJob
+from airweave.schemas.sync_job import SyncJobCreate
+
+EXECUTION_TIMEOUT_SECONDS = 20
+
+
+class OwnedProvisioningService:
+ """Retries repair one desired relationship; there is no independent job queue."""
+
+ def __init__(
+ self,
+ store: ProvisioningStore,
+ lifecycle: SourceLifecycleServiceProtocol,
+ jobs: SyncJobRepositoryProtocol,
+ syncs: SyncRepositoryProtocol,
+ schedules: TemporalScheduleServiceProtocol,
+ workflows: TemporalWorkflowServiceProtocol,
+ ):
+ """Use existing source verification, job admission and Temporal services."""
+ self.store, self.lifecycle = store, lifecycle
+ self.jobs, self.syncs = jobs, syncs
+ self.schedules, self.workflows = schedules, workflows
+
+ async def ensure(
+ self, db: AsyncSession, ctx: ApiContext, account: UUID, request: EnsureSource
+ ) -> ProvisionedSource:
+ """Persist before verification/execution; failures are retried with the same request."""
+ await self.store.ensure(db, ctx, account, request)
+ return await self.reconcile(db, ctx, account, request.generation)
+
+ async def get(self, db: AsyncSession, ctx: ApiContext, account: UUID) -> ProvisionedSource:
+ """Read only current committed relationship; never infer completion from IDs."""
+ row = await db.scalar(
+ select(OwnedProvisioning)
+ .where(
+ OwnedProvisioning.organization_id == ctx.organization.id,
+ OwnedProvisioning.client_namespace == "almanac",
+ OwnedProvisioning.account_id == account,
+ )
+ .execution_options(populate_existing=True)
+ )
+ if row is None:
+ raise HTTPException(status_code=404, detail="Owned account source not found")
+ source = (
+ await db.get(SourceConnection, row.source_connection_id)
+ if row.source_connection_id
+ else None
+ )
+ identity, user = (
+ native_principal(source.short_name, source.config_fields) if source else (None, None)
+ )
+ return ProvisionedSource(
+ account_id=row.account_id,
+ organization_id=row.organization_id,
+ generation=row.generation,
+ observed_generation=row.observed_generation,
+ state=(
+ "pending"
+ if row.observed_generation != row.generation
+ else "ready"
+ if row.desired_state == "active"
+ else row.desired_state
+ ),
+ source_connection_id=row.source_connection_id,
+ sync_id=row.sync_id,
+ expected_identity=identity,
+ expected_user_identity=user,
+ )
+
+ async def reconcile(
+ self, db: AsyncSession, ctx: ApiContext, account: UUID, generation: int
+ ) -> ProvisionedSource:
+ """Provider requests outside locks; late success must pass generation CAS."""
+ async with UnitOfWork(db):
+ row = await locked_intent(db, ctx.organization.id, account)
+ self._current(row, generation)
+ done = row.observed_generation == generation
+ source_id, sync_id = row.source_connection_id, row.sync_id
+ active = row.desired_state == "active"
+ needs_verification = active and row.verified_at is None
+ if done:
+ return await self.get(db, ctx, account)
+ if needs_verification:
+ # create performs the connector's native principal validation, not
+ # merely connected-account metadata validation in Composio.
+ verified_source = await self.lifecycle.create(db, source_id, ctx)
+ await verified_source.http_client.aclose()
+ await db.commit()
+ async with UnitOfWork(db) as uow:
+ row = await locked_intent(db, ctx.organization.id, account)
+ self._current(row, generation)
+ sync = await db.scalar(
+ select(Sync)
+ .where(Sync.id == sync_id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ source = await db.get(SourceConnection, source_id)
+ sync.provisioning_ready_generation = generation
+ sync.status = "active"
+ source.is_authenticated = True
+ row.verified_at = utc_now_naive()
+ # Admission reloads Sync with populate_existing; production disables
+ # autoflush. Keep readiness and admission in this same transaction.
+ await db.flush()
+ if row.initial_job_id is None:
+ job = await self.jobs.create(db, SyncJobCreate(sync_id=sync.id), ctx, uow=uow)
+ await db.flush()
+ row.initial_job_id = job.id
+ try:
+ async with asyncio.timeout(EXECUTION_TIMEOUT_SECONDS):
+ await self._execution(db, ctx, account, generation)
+ except TimeoutError:
+ raise HTTPException(
+ status_code=503, detail="Owned capture setup is still pending"
+ ) from None
+ return await self.get(db, ctx, account)
+
+ @staticmethod
+ def _current(row: OwnedProvisioning, generation: int) -> None:
+ if row.generation != generation:
+ raise HTTPException(status_code=409, detail="Account generation was superseded")
+
+ async def _execution(
+ self, db: AsyncSession, ctx: ApiContext, account: UUID, generation: int
+ ) -> None:
+ # Only bounded Temporal RPCs occur under this relationship lock. This
+ # serializes older schedule creation against newer disconnect cleanup.
+ async with UnitOfWork(db) as uow:
+ row = await locked_intent(db, ctx.organization.id, account)
+ self._current(row, generation)
+ if row.observed_generation == generation:
+ return
+ for job_id in row.cancellation_job_ids:
+ result = await self.workflows.cancel_sync_job_workflow(job_id, ctx)
+ if not result["success"]:
+ raise HTTPException(
+ status_code=503, detail="Owned capture cancellation is pending"
+ )
+ row.cancellation_job_ids = []
+ if row.desired_state != "active":
+ if row.sync_id is not None:
+ await self.schedules.delete_all_schedules_for_sync(
+ row.sync_id, db, ctx, uow=uow, strict=True
+ )
+ else:
+ await self._start(db, ctx, row, uow)
+ row.observed_generation = generation
+
+ async def _start(
+ self, db: AsyncSession, ctx: ApiContext, row: OwnedProvisioning, uow: UnitOfWork
+ ) -> None:
+ spec = EnsureSource.model_validate(row.request_payload).source
+ source = await db.get(SourceConnection, row.source_connection_id)
+ # A prior RPC may have succeeded before its DB link committed. Remove
+ # deterministic handles even when temporal_schedule_id is absent.
+ await self.schedules.delete_all_schedules_for_sync(
+ row.sync_id, db, ctx, uow=uow, strict=True
+ )
+ await self.schedules.create_or_update_schedule(
+ row.sync_id,
+ spec.cron,
+ db,
+ ctx,
+ uow,
+ collection_readable_id=source.readable_collection_id,
+ connection_id=source.connection_id,
+ )
+ job = await db.get(SyncJob, row.initial_job_id)
+ if job.status not in ("pending", "running"):
+ return # A lost reply may be retried after the initial run already finished.
+ sync = await self.syncs.get(db, row.sync_id, ctx)
+ collection = await db.scalar(
+ select(Collection).where(
+ Collection.organization_id == ctx.organization.id,
+ Collection.readable_id == source.readable_collection_id,
+ )
+ )
+ connection = await db.get(Connection, source.connection_id)
+ try:
+ await self.workflows.run_source_connection_workflow(
+ sync=sync,
+ sync_job=schemas.SyncJob.model_validate(job, from_attributes=True),
+ collection=schemas.CollectionRecord.model_validate(
+ collection, from_attributes=True
+ ),
+ connection=schemas.Connection.model_validate(connection, from_attributes=True),
+ ctx=ctx,
+ )
+ except WorkflowAlreadyStartedError:
+ pass # Same durable job ID is the exact same execution, not a new attempt.
diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py
new file mode 100644
index 0000000000..6becbd27a5
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/store.py
@@ -0,0 +1,248 @@
+"""Commit desired generation and source configuration before remote execution."""
+
+import hashlib
+import json
+from uuid import UUID, uuid4
+
+from fastapi import HTTPException
+from sqlalchemy import select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.api.context import ApiContext
+from airweave.core.datetime_utils import utc_now_naive
+from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.owned_provisioning.models import EnsureSource, native_principal
+from airweave.domains.source_connections.protocols import SourceConnectionCreateServiceProtocol
+from airweave.domains.sources.protocols import SourceValidationServiceProtocol
+from airweave.models.connection import Connection
+from airweave.models.owned_provisioning import OwnedProvisioning
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+from airweave.models.sync_job import SyncJob
+from airweave.platform.configs.config import StripeConfig
+from airweave.schemas.source_connection import (
+ AuthProviderAuthentication,
+ ScheduleConfig,
+ SourceConnectionCreate,
+)
+
+
+async def locked_intent(db: AsyncSession, organization: UUID, account: UUID) -> OwnedProvisioning:
+ """Every generation/execution mutation serializes on this one relationship."""
+ row = await db.scalar(
+ select(OwnedProvisioning)
+ .where(
+ OwnedProvisioning.organization_id == organization,
+ OwnedProvisioning.client_namespace == "almanac",
+ OwnedProvisioning.account_id == account,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if row is None:
+ raise HTTPException(status_code=404, detail="Owned account source not found")
+ return row
+
+
+class ProvisioningStore:
+ """Source domain creation participates in our transaction; it cannot commit early."""
+
+ def __init__(
+ self,
+ create: SourceConnectionCreateServiceProtocol,
+ validation: SourceValidationServiceProtocol,
+ ):
+ """Reuse transactional source creation and source-specific validation."""
+ self.create = create
+ self.validation = validation
+
+ async def ensure(
+ self, db: AsyncSession, ctx: ApiContext, account: UUID, request: EnsureSource
+ ) -> None:
+ """Lost replies replay one source; lower generations cannot overwrite it."""
+ payload = request.model_dump(mode="json")
+ fingerprint = hashlib.sha256(
+ json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+ ).hexdigest()
+ async with UnitOfWork(db) as uow:
+ now = utc_now_naive()
+ await db.execute(
+ insert(OwnedProvisioning)
+ .values(
+ id=uuid4(),
+ organization_id=ctx.organization.id,
+ client_namespace="almanac",
+ account_id=account,
+ generation=request.generation,
+ request_hash=fingerprint,
+ request_payload=payload,
+ desired_state=request.state,
+ observed_generation=0,
+ cancellation_job_ids=[],
+ created_at=now,
+ modified_at=now,
+ )
+ .on_conflict_do_nothing(constraint="uq_owned_provisioning_account")
+ )
+ row = await locked_intent(db, ctx.organization.id, account)
+ if request.generation < row.generation:
+ raise HTTPException(status_code=409, detail="Account generation was superseded")
+ if request.generation == row.generation and row.request_hash != fingerprint:
+ raise HTTPException(
+ status_code=409, detail="Generation already names another request"
+ )
+ if request.generation == row.generation and (
+ row.sync_id is not None or row.desired_state != "active"
+ ):
+ return
+ if row.desired_state == "disconnected" and request.state != "disconnected":
+ raise HTTPException(
+ status_code=409, detail="Connect a new account after disconnect"
+ )
+ if request.state == "active" and row.source_connection_id is not None:
+ source = await db.scalar(
+ select(SourceConnection)
+ .where(
+ SourceConnection.id == row.source_connection_id,
+ SourceConnection.organization_id == ctx.organization.id,
+ )
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ if (
+ source.short_name,
+ native_principal(source.short_name, source.config_fields),
+ source.readable_collection_id,
+ ) != (
+ request.source.provider,
+ (request.source.expected_identity, request.source.expected_user_identity),
+ request.source.collection,
+ ):
+ raise HTTPException(
+ status_code=409, detail="Reconnect changes original account identity"
+ )
+ if request.source is not None:
+ await self._configure(db, uow, ctx, row, request)
+ if row.sync_id is not None:
+ sync = await db.scalar(
+ select(Sync)
+ .where(Sync.id == row.sync_id, Sync.organization_id == ctx.organization.id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ )
+ sync.provisioning_generation = request.generation
+ sync.status = "paused"
+ sync.writer_epoch += 1
+ jobs = (
+ await db.scalars(
+ select(SyncJob)
+ .where(
+ SyncJob.sync_id == sync.id,
+ SyncJob.status.in_(("pending", "running", "cancelling")),
+ )
+ .with_for_update()
+ )
+ ).all()
+ row.cancellation_job_ids = list(
+ dict.fromkeys([*row.cancellation_job_ids, *(str(job.id) for job in jobs)])
+ )
+ for job in jobs:
+ job.status = "cancelled"
+ source = await db.get(SourceConnection, row.source_connection_id)
+ source.is_authenticated = False
+ row.generation = request.generation
+ row.request_payload = payload
+ row.request_hash = fingerprint
+ row.desired_state = request.state
+ row.initial_job_id = None
+ row.verified_at = None
+ await db.flush()
+
+ async def _configure(
+ self,
+ db: AsyncSession,
+ uow: UnitOfWork,
+ ctx: ApiContext,
+ row: OwnedProvisioning,
+ request: EnsureSource,
+ ) -> None:
+ spec = request.source
+ auth = await db.scalar(
+ select(Connection).where(
+ Connection.organization_id == ctx.organization.id,
+ Connection.readable_id == spec.auth_provider,
+ Connection.short_name == "composio",
+ )
+ )
+ if auth is None:
+ raise HTTPException(
+ status_code=400,
+ detail="Owned capture requires this organization's Composio connection",
+ )
+ config = self.validation.validate_config(spec.provider, spec.source_config(), ctx)
+ if row.source_connection_id is None:
+ created = await self.create.create_deferred(
+ db,
+ ctx=ctx,
+ uow=uow,
+ obj_in=SourceConnectionCreate(
+ name="Almanac " + spec.provider,
+ short_name=spec.provider,
+ readable_collection_id=spec.collection,
+ config=config,
+ schedule=ScheduleConfig(cron=spec.cron),
+ sync_immediately=False,
+ authentication=AuthProviderAuthentication(
+ provider_readable_id=spec.auth_provider, provider_config=spec.auth_config()
+ ),
+ ),
+ )
+ if created.sync_id is None:
+ raise ValueError("Owned source creation must create its durable sync")
+ row.source_connection_id = created.id
+ row.sync_id = created.sync_id
+ else:
+ source = await db.scalar(
+ select(SourceConnection)
+ .where(
+ SourceConnection.id == row.source_connection_id,
+ SourceConnection.organization_id == ctx.organization.id,
+ )
+ .with_for_update()
+ )
+ if spec.provider == "stripe" and (
+ StripeConfig.model_validate(source.config_fields or {}).original_capture
+ != StripeConfig.model_validate(config).original_capture
+ ):
+ raise HTTPException(
+ status_code=409,
+ detail="Reconnect changes Stripe account mode, context or API version; "
+ "connect a new account",
+ )
+ if spec.provider == "github":
+ from airweave.platform.configs.config import GitHubConfig
+
+ if GitHubConfig.model_validate(source.config_fields) != GitHubConfig.model_validate(
+ config
+ ):
+ raise HTTPException(
+ status_code=409,
+ detail="Reconnect changes GitHub repository selection or capture options; "
+ "connect a new account",
+ )
+ if spec.provider == "linear":
+ from airweave.platform.configs.config import LinearConfig
+
+ previous_teams = set(LinearConfig.model_validate(source.config_fields).team_ids)
+ current_teams = set(LinearConfig.model_validate(config).team_ids)
+ if previous_teams != current_teams:
+ raise HTTPException(
+ status_code=409,
+ detail="Reconnect changes selected Linear teams; connect a new account",
+ )
+ source.config_fields = config
+ source.readable_auth_provider_id = spec.auth_provider
+ source.auth_provider_config = spec.auth_config()
+ sync = await db.get(Sync, row.sync_id)
+ sync.cron_schedule = spec.cron
diff --git a/backend/airweave/domains/owned_provisioning/tests/conftest.py b/backend/airweave/domains/owned_provisioning/tests/conftest.py
new file mode 100644
index 0000000000..2916906a10
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/tests/conftest.py
@@ -0,0 +1,14 @@
+"""Use isolated PostgreSQL with the production session autoflush policy."""
+
+import pytest
+
+from airweave.domains.entities.canonical.tests.conftest import (
+ database as canonical_database, # noqa: F401
+)
+
+
+@pytest.fixture
+async def database(canonical_database): # noqa: F811 - pytest injects the imported fixture
+ """Keep implicit flushes from masking production transaction ordering."""
+ canonical_database.configure(autoflush=False)
+ return canonical_database
diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py
new file mode 100644
index 0000000000..160e2fa6e2
--- /dev/null
+++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py
@@ -0,0 +1,999 @@
+"""Real SQL lifecycle and lost-reply recovery; no private provider or Temporal calls."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+from uuid import uuid4
+
+import pytest
+from fastapi import HTTPException
+from sqlalchemy import func, select
+
+from airweave import schemas
+from airweave.core.shared_models import AuthMethod, IntegrationType
+from airweave.domains.collections.repository import CollectionRepository
+from airweave.domains.connections.repository import ConnectionRepository
+from airweave.domains.owned_provisioning.models import EnsureSource, ManagedSource, native_principal
+from airweave.domains.owned_provisioning.service import OwnedProvisioningService
+from airweave.domains.owned_provisioning.store import ProvisioningStore
+from airweave.domains.source_connections.repository import SourceConnectionRepository
+from airweave.domains.source_connections.tests.test_create import _ctx, _entry, _service
+from airweave.domains.syncs.jobs.repository import SyncJobRepository
+from airweave.domains.syncs.repository import SyncRepository
+from airweave.domains.syncs.service import SyncService
+from airweave.models import Collection, Connection, Organization, SourceConnection, Sync, SyncJob
+from airweave.models.owned_provisioning import OwnedProvisioning
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+pytestmark = pytest.mark.integration
+
+
+@pytest.fixture
+async def setup(database):
+ ctx = _ctx()
+ ctx.auth_method = AuthMethod.API_KEY
+ metadata_id = uuid4()
+ async with database() as db:
+ db.add(Organization(id=ctx.organization.id, name="Synthetic provisioning"))
+ db.add(
+ VectorDbDeploymentMetadata(
+ id=metadata_id,
+ dense_embedder="test",
+ sparse_embedder="test",
+ embedding_dimensions=3,
+ )
+ )
+ await db.flush()
+ db.add(
+ Collection(
+ name="Owned",
+ readable_id="owned",
+ organization_id=ctx.organization.id,
+ vector_db_deployment_metadata_id=metadata_id,
+ )
+ )
+ db.add(
+ Connection(
+ name="Composio",
+ readable_id="composio",
+ short_name="composio",
+ organization_id=ctx.organization.id,
+ integration_type=IntegrationType.AUTH_PROVIDER,
+ )
+ )
+ await db.commit()
+
+ lifecycle = SimpleNamespace(
+ create=AsyncMock(
+ return_value=SimpleNamespace(http_client=SimpleNamespace(aclose=AsyncMock()))
+ )
+ )
+ schedules = SimpleNamespace(
+ create_or_update_schedule=AsyncMock(), delete_all_schedules_for_sync=AsyncMock()
+ )
+ workflows = SimpleNamespace(
+ run_source_connection_workflow=AsyncMock(),
+ cancel_sync_job_workflow=AsyncMock(return_value={"success": True, "workflow_found": True}),
+ )
+ entry = _entry()
+ entry.short_name = "gmail"
+ creator = _service(entry)
+ creator._sc_repo = SourceConnectionRepository(creator._source_registry)
+ creator._collection_repo = CollectionRepository(creator._source_registry, creator._sc_repo)
+ creator._connection_repo = ConnectionRepository()
+ creator._source_validation.seed_config_result(
+ "gmail", {"expected_mailbox": "owner@example.test"}
+ )
+ syncs = SyncRepository()
+ jobs = SyncJobRepository()
+ creator._sync_service = SyncService(
+ syncs, jobs, Mock(), Mock(), Mock(), workflows, schedules, Mock()
+ )
+ creator._sync_service.resolve_destination_ids = AsyncMock(return_value=[])
+ creator._response_builder.build_response = AsyncMock(
+ side_effect=lambda db, source, ctx: SimpleNamespace(id=source.id, sync_id=source.sync_id)
+ )
+ service = OwnedProvisioningService(
+ ProvisioningStore(
+ creator, SimpleNamespace(validate_config=Mock(side_effect=lambda p, c, ctx: c))
+ ),
+ lifecycle,
+ jobs,
+ syncs,
+ schedules,
+ workflows,
+ )
+ request = EnsureSource(
+ generation=1,
+ state="active",
+ source=ManagedSource(
+ provider="gmail",
+ expected_identity="owner@example.test",
+ collection="owned",
+ auth_provider="composio",
+ connected_account_id="ca_first",
+ auth_config_id="ac_test",
+ user_id="owner",
+ cron="0 * * * *",
+ ),
+ )
+ return ctx, service, request, uuid4(), lifecycle, schedules, workflows
+
+
+async def test_lost_reply_and_execution_failure_reuse_source_and_job(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ schedules.create_or_update_schedule.side_effect = RuntimeError("Temporal unavailable")
+ async with database() as db:
+ with pytest.raises(RuntimeError, match="Temporal unavailable"):
+ await service.ensure(db, ctx, account, request)
+ async with database() as db:
+ pending = await service.get(db, ctx, account)
+ assert pending.state == "pending"
+ row = await db.scalar(
+ select(OwnedProvisioning).where(OwnedProvisioning.account_id == account)
+ )
+ first_job = row.initial_job_id
+ schedules.create_or_update_schedule.side_effect = None
+ async with database() as db:
+ ready = await service.ensure(db, ctx, account, request)
+ async with database() as db:
+ replay = await service.ensure(db, ctx, account, request)
+ assert ready == replay and ready.state == "ready"
+ assert ready.sync_id == pending.sync_id
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(SourceConnection)
+ .where(SourceConnection.organization_id == ctx.organization.id)
+ )
+ == 1
+ )
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == ready.sync_id)
+ )
+ == 1
+ )
+ assert (await db.get(SyncJob, first_job)).provisioning_generation == 1
+ with pytest.raises(HTTPException) as conflict:
+ await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(
+ update={
+ "source": request.source.model_copy(
+ update={"connected_account_id": "ca_other"}
+ )
+ }
+ ),
+ )
+ assert conflict.value.status_code == 409
+ lifecycle.create.assert_awaited_once()
+ workflows.run_source_connection_workflow.assert_awaited_once()
+ # Reconciliation always repairs orphan deterministic schedules, even with no DB link.
+ assert schedules.delete_all_schedules_for_sync.await_count == 2
+
+
+async def test_disconnect_wins_over_inflight_verification_and_blocks_admission(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ started, release = asyncio.Event(), asyncio.Event()
+ resource = SimpleNamespace(http_client=SimpleNamespace(aclose=AsyncMock()))
+
+ async def verify(*args):
+ started.set()
+ await release.wait()
+ return resource
+
+ lifecycle.create.side_effect = verify
+
+ async def connect():
+ async with database() as db:
+ return await service.ensure(db, ctx, account, request)
+
+ task = asyncio.create_task(connect())
+ await started.wait()
+ async with database() as db:
+ pending = await service.get(db, ctx, account)
+ with pytest.raises(HTTPException) as blocked:
+ await service.jobs.create(db, schemas.SyncJobCreate(sync_id=pending.sync_id), ctx)
+ assert blocked.value.status_code == 409
+ async with database() as db:
+ disconnected = await service.ensure(
+ db, ctx, account, EnsureSource(generation=2, state="disconnected")
+ )
+ assert disconnected.state == "disconnected"
+ release.set()
+ with pytest.raises(HTTPException) as stale:
+ await task
+ assert stale.value.status_code == 409
+ resource.http_client.aclose.assert_awaited_once()
+ workflows.run_source_connection_workflow.assert_not_awaited()
+ async with database() as db:
+ with pytest.raises(HTTPException):
+ await service.ensure(db, ctx, account, request)
+ assert (await db.get(Sync, pending.sync_id)).provisioning_ready_generation == 0
+
+
+async def test_rotation_timeout_retains_cleanup_and_original_source(database, setup, monkeypatch):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ original_job = await db.scalar(select(SyncJob.id).where(SyncJob.sync_id == first.sync_id))
+ rotated = request.model_copy(
+ update={
+ "generation": 2,
+ "source": request.source.model_copy(update={"connected_account_id": "ca_new"}),
+ }
+ )
+
+ async def stalled(*args):
+ await asyncio.Event().wait()
+
+ workflows.cancel_sync_job_workflow.side_effect = stalled
+ monkeypatch.setattr(
+ "airweave.domains.owned_provisioning.service.EXECUTION_TIMEOUT_SECONDS", 0.02
+ )
+ async with database() as db:
+ with pytest.raises(HTTPException) as timeout:
+ await service.ensure(db, ctx, account, rotated)
+ assert timeout.value.status_code == 503
+ async with database() as db:
+ row = await db.scalar(
+ select(OwnedProvisioning).where(OwnedProvisioning.account_id == account)
+ )
+ assert row.cancellation_job_ids == [str(original_job)]
+ assert row.generation == 2 and row.observed_generation == 1
+ assert (
+ row.sync_id == first.sync_id and row.source_connection_id == first.source_connection_id
+ )
+ workflows.cancel_sync_job_workflow.side_effect = None
+ monkeypatch.setattr("airweave.domains.owned_provisioning.service.EXECUTION_TIMEOUT_SECONDS", 20)
+ async with database() as db:
+ second = await service.ensure(db, ctx, account, rotated)
+ assert second.state == "ready" and second.sync_id == first.sync_id
+ assert (
+ await db.scalar(
+ select(func.count()).select_from(SyncJob).where(SyncJob.sync_id == first.sync_id)
+ )
+ == 2
+ )
+ other = _ctx()
+ with pytest.raises(HTTPException) as denied:
+ await service.get(db, other, account)
+ assert denied.value.status_code == 404
+
+
+async def test_http_boundary_and_legacy_edit_cannot_bypass_generation(database, setup):
+ import httpx
+ from fastapi import FastAPI
+
+ from airweave.api import deps
+ from airweave.api.v1.endpoints.owned_provisioning import router
+ from airweave.db.session import get_db
+ from airweave.domains.source_connections.tests.test_delete import (
+ _build_service as delete_service,
+ )
+ from airweave.domains.source_connections.tests.test_update import (
+ _build_service as update_service,
+ )
+ from airweave.schemas.source_connection import SourceConnectionUpdate
+
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ app = FastAPI()
+ app.include_router(router, prefix="/owned-sources")
+
+ async def session():
+ async with database() as db:
+ yield db
+
+ app.dependency_overrides[get_db] = session
+ app.dependency_overrides[deps.get_context] = lambda: ctx
+ app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(
+ owned_provisioning=service
+ )
+ async with httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=app), base_url="http://test"
+ ) as http:
+ ctx.auth_method = AuthMethod.SYSTEM
+ denied = await http.put(f"/owned-sources/{account}", json=request.model_dump(mode="json"))
+ assert denied.status_code == 403
+ lifecycle.create.assert_not_awaited()
+ ctx.auth_method = AuthMethod.API_KEY
+ connected = await http.put(
+ f"/owned-sources/{account}", json=request.model_dump(mode="json")
+ )
+ assert connected.status_code == 200, connected.text
+ source_id = connected.json()["source_connection_id"]
+ from uuid import UUID
+
+ source_id = UUID(source_id)
+ repo = SourceConnectionRepository(service.store.create._source_registry)
+ async with database() as db:
+ with pytest.raises(HTTPException) as guarded:
+ await update_service(sc_repo=repo).update(
+ db,
+ id=source_id,
+ obj_in=SourceConnectionUpdate(
+ config={"expected_mailbox": "someone-else@example.test"}
+ ),
+ ctx=ctx,
+ )
+ assert guarded.value.status_code == 409
+ async with database() as db:
+ with pytest.raises(HTTPException) as guarded:
+ await delete_service(sc_repo=repo).delete(db, id=source_id, ctx=ctx)
+ assert guarded.value.status_code == 409
+
+
+async def test_creation_failure_rolls_back_relationship_source_and_sync(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ response = service.store.create._response_builder.build_response
+ original = response.side_effect
+ response.side_effect = RuntimeError("creation interrupted before commit")
+ async with database() as db:
+ with pytest.raises(RuntimeError, match="creation interrupted"):
+ await service.ensure(db, ctx, account, request)
+ async with database() as db:
+ for model in (OwnedProvisioning, SourceConnection, Sync, SyncJob):
+ assert (
+ await db.scalar(
+ select(func.count())
+ .select_from(model)
+ .where(model.organization_id == ctx.organization.id)
+ )
+ == 0
+ )
+ lifecycle.create.assert_not_awaited()
+ schedules.create_or_update_schedule.assert_not_awaited()
+ response.side_effect = original
+ async with database() as db:
+ ready = await service.ensure(db, ctx, account, request)
+ assert ready.state == "ready"
+
+
+async def test_pause_resume_preserves_identity_source_and_cleanup(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ async with database() as db:
+ paused = await service.ensure(db, ctx, account, EnsureSource(generation=2, state="paused"))
+ assert paused.state == "paused"
+ assert paused.expected_identity == request.source.expected_identity
+ assert (paused.source_connection_id, paused.sync_id) == (
+ first.source_connection_id,
+ first.sync_id,
+ )
+ sync = await db.get(Sync, first.sync_id)
+ assert sync.status == "paused" and sync.provisioning_generation == 2
+ with pytest.raises(HTTPException):
+ await service.jobs.create(db, schemas.SyncJobCreate(sync_id=first.sync_id), ctx)
+ workflows.cancel_sync_job_workflow.assert_awaited_once()
+ async with database() as db:
+ wrong_source = request.source.model_copy(update={"expected_identity": "other@example.test"})
+ wrong = request.model_copy(update={"generation": 3, "source": wrong_source})
+ with pytest.raises(HTTPException, match="Reconnect changes original account identity"):
+ await service.ensure(db, ctx, account, wrong)
+ async with database() as db:
+ resumed = await service.ensure(
+ db, ctx, account, request.model_copy(update={"generation": 3})
+ )
+ assert resumed.state == "ready"
+ assert (resumed.source_connection_id, resumed.sync_id) == (
+ first.source_connection_id,
+ first.sync_id,
+ )
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert lifecycle.create.await_count == 2
+ async with database() as db:
+ stopped = await service.ensure(
+ db, ctx, account, EnsureSource(generation=4, state="disconnected")
+ )
+ assert stopped.expected_identity == request.source.expected_identity
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Connect a new account after disconnect"):
+ await service.ensure(db, ctx, account, request.model_copy(update={"generation": 5}))
+
+
+async def test_initial_pause_can_activate_but_initial_disconnect_is_terminal(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ async with database() as db:
+ paused = await service.ensure(db, ctx, account, EnsureSource(generation=1, state="paused"))
+ assert paused.state == "paused" and paused.expected_identity is None
+ assert paused.source_connection_id is None and paused.sync_id is None
+ async with database() as db:
+ active = await service.ensure(
+ db, ctx, account, request.model_copy(update={"generation": 2})
+ )
+ assert active.state == "ready"
+ other = uuid4()
+ async with database() as db:
+ await service.ensure(db, ctx, other, EnsureSource(generation=1, state="disconnected"))
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Connect a new account after disconnect"):
+ await service.ensure(db, ctx, other, request.model_copy(update={"generation": 2}))
+
+
+async def test_slack_workspace_user_survives_pause_and_rejects_different_member(database, setup):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = "slack"
+ creator._source_validation.seed_config_result(
+ "slack", {"expected_team_id": "T1", "expected_user_id": "U1"}
+ )
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": "slack",
+ "expected_identity": "T1",
+ "expected_user_identity": "U1",
+ "config": {"expected_team_id": "wrong", "expected_user_id": "wrong"},
+ }
+ )
+ request = request.model_copy(update={"source": spec})
+ assert spec.source_config() == {"expected_team_id": "T1", "expected_user_id": "U1"}
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ assert (first.expected_identity, first.expected_user_identity) == ("T1", "U1")
+ async with database() as db:
+ paused = await service.ensure(db, ctx, account, EnsureSource(generation=2, state="paused"))
+ assert (paused.expected_identity, paused.expected_user_identity) == ("T1", "U1")
+ async with database() as db:
+ wrong = request.model_copy(
+ update={
+ "generation": 3,
+ "source": spec.model_copy(update={"expected_user_identity": "U2"}),
+ }
+ )
+ with pytest.raises(HTTPException, match="Reconnect changes original account identity"):
+ await service.ensure(db, ctx, account, wrong)
+ async with database() as db:
+ resumed = await service.ensure(
+ db, ctx, account, request.model_copy(update={"generation": 3})
+ )
+ assert resumed.source_connection_id == first.source_connection_id
+ assert (resumed.expected_identity, resumed.expected_user_identity) == ("T1", "U1")
+
+
+@pytest.mark.parametrize("provider", ["outlook_mail", "outlook_calendar"])
+async def test_outlook_fresh_source_retry_and_principal_cannot_downgrade_owned_capture(
+ database, setup, provider
+):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = provider
+ expected_config = {"expected_principal_id": "native-owner", "capture_originals": True}
+ creator._source_validation.seed_config_result(provider, expected_config)
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": provider,
+ "expected_identity": "native-owner",
+ "config": {"expected_principal_id": "wrong-owner", "capture_originals": False},
+ }
+ )
+ request = request.model_copy(update={"source": spec})
+ assert spec.source_config() == expected_config
+ assert native_principal(provider, spec.source_config()) == ("native-owner", None)
+ with pytest.raises(ValueError, match="must capture originals"):
+ native_principal(provider, {"expected_principal_id": "native-owner"})
+ with pytest.raises(ValueError):
+ native_principal(provider, {"capture_originals": True})
+
+ async with database() as db:
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 0
+ first = await service.ensure(db, ctx, account, request)
+ assert first.state == "ready" and first.expected_identity == "native-owner"
+ source = await db.get(SourceConnection, first.source_connection_id)
+ assert source.short_name == provider and source.config_fields == expected_config
+ sync = await db.get(Sync, first.sync_id)
+ original_epoch = sync.writer_epoch
+ assert sync.provisioning_generation == sync.provisioning_ready_generation == 1
+ assert sync.index_pipeline_version == 2
+ async with database() as db:
+ replay = await service.ensure(db, ctx, account, request)
+ assert replay == first
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert await db.scalar(select(func.count()).select_from(Sync)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == original_epoch
+ async with database() as db:
+ wrong = request.model_copy(
+ update={
+ "generation": 2,
+ "source": spec.model_copy(update={"expected_identity": "other-native-owner"}),
+ }
+ )
+ with pytest.raises(HTTPException, match="Reconnect changes original account identity"):
+ await service.ensure(db, ctx, account, wrong)
+ async with database() as db:
+ current = await service.get(db, ctx, account)
+ assert current == first
+ source = await db.get(SourceConnection, first.source_connection_id)
+ assert source.config_fields == expected_config
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == original_epoch
+ lifecycle.create.assert_awaited_once()
+ workflows.run_source_connection_workflow.assert_awaited_once()
+
+
+async def test_verified_source_job_admission_without_autoflush(database, setup):
+ """Production sessions must persist readiness before admission reloads the Sync."""
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ async with database(autoflush=False) as db:
+ result = await service.ensure(db, ctx, account, request)
+ assert result.state == "ready"
+ async with database(autoflush=False) as db:
+ sync = await db.get(Sync, result.sync_id)
+ source = await db.get(SourceConnection, result.source_connection_id)
+ row = await db.scalar(
+ select(OwnedProvisioning).where(OwnedProvisioning.account_id == account)
+ )
+ job = await db.get(SyncJob, row.initial_job_id)
+ assert sync.provisioning_ready_generation == request.generation
+ assert sync.status == "active" and source.is_authenticated
+ assert row.verified_at is not None
+ assert job.provisioning_generation == request.generation
+ assert (await service.ensure(db, ctx, account, request)) == result
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ lifecycle.create.assert_awaited_once()
+ workflows.run_source_connection_workflow.assert_awaited_once()
+
+
+async def test_job_admission_failure_rolls_back_verified_readiness(database, setup, monkeypatch):
+ """Flushed verification and a real admitted job remain in the same transaction."""
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ create_job = service.jobs.create
+
+ async def fail_after_admission(db, obj_in, ctx, uow=None):
+ job = await create_job(db, obj_in, ctx, uow=uow)
+ await db.flush()
+ assert job.provisioning_generation == request.generation
+ raise RuntimeError("Synthetic failure after real job admission")
+
+ monkeypatch.setattr(service.jobs, "create", fail_after_admission)
+ async with database(autoflush=False) as db:
+ with pytest.raises(RuntimeError, match="Synthetic failure after real job admission"):
+ await service.ensure(db, ctx, account, request)
+ async with database(autoflush=False) as db:
+ row = await db.scalar(
+ select(OwnedProvisioning).where(OwnedProvisioning.account_id == account)
+ )
+ sync = await db.get(Sync, row.sync_id)
+ source = await db.get(SourceConnection, row.source_connection_id)
+ assert row.verified_at is None and row.initial_job_id is None
+ assert sync.provisioning_ready_generation != request.generation
+ assert sync.status != "active" and not source.is_authenticated
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 0
+ assert (await service.get(db, ctx, account)).state == "pending"
+ schedules.create_or_update_schedule.assert_not_awaited()
+ workflows.run_source_connection_workflow.assert_not_awaited()
+ monkeypatch.setattr(service.jobs, "create", create_job)
+ async with database(autoflush=False) as db:
+ assert (await service.ensure(db, ctx, account, request)).state == "ready"
+
+
+async def test_stripe_provisioning_retry_and_reconnect_preserve_account_mode_and_version(
+ database, setup
+):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = "stripe"
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": "stripe",
+ "expected_identity": "acct_selected",
+ "config": {
+ "expected_account_id": "acct_untrusted",
+ "livemode": False,
+ "api_version": "2025-06-30.basil",
+ },
+ }
+ )
+ expected_config = spec.source_config()
+ creator._source_validation.seed_config_result("stripe", expected_config)
+ assert native_principal("stripe", expected_config) == ("acct_selected", None)
+ request = request.model_copy(update={"source": spec})
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ sync = await db.get(Sync, first.sync_id)
+ assert sync.index_pipeline_version == 2
+ original_epoch = sync.writer_epoch
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, request) == first
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == original_epoch
+ for change in (
+ {"livemode": True},
+ {"api_version": "2026-01-28.clover"},
+ {"connected_account_id": "acct_selected"},
+ ):
+ changed = ManagedSource.model_validate(
+ {
+ **spec.model_dump(),
+ "config": {**spec.config, **change},
+ }
+ )
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Reconnect changes Stripe"):
+ await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(
+ update={
+ "generation": 2,
+ "source": changed,
+ }
+ ),
+ )
+ async with database() as db:
+ assert await service.get(db, ctx, account) == first
+ assert (
+ await db.get(SourceConnection, first.source_connection_id)
+ ).config_fields == expected_config
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == original_epoch
+ # New broker credentials may reconnect the same native identity/configuration.
+ async with database() as db:
+ reconnected = await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(
+ update={
+ "generation": 2,
+ "source": spec.model_copy(update={"connected_account_id": "ca_reconnected"}),
+ }
+ ),
+ )
+ assert reconnected.source_connection_id == first.source_connection_id
+ assert reconnected.sync_id == first.sync_id and reconnected.generation == 2
+ source = await db.get(SourceConnection, first.source_connection_id)
+ assert source.auth_provider_config["account_id"] == "ca_reconnected"
+ assert source.config_fields == expected_config
+
+
+@pytest.mark.parametrize(
+ "config",
+ [
+ {},
+ {"livemode": "false", "api_version": "2025-06-30.basil"},
+ {
+ "livemode": False,
+ "api_version": "2025-06-30.basil",
+ "connected_account_id": "acct_other",
+ },
+ ],
+)
+def test_stripe_provisioning_requires_explicit_typed_mode_and_version(config):
+ with pytest.raises(ValueError):
+ ManagedSource(
+ provider="stripe",
+ expected_identity="acct_selected",
+ collection="owned",
+ auth_provider="composio",
+ connected_account_id="ca_selected",
+ auth_config_id="ac_selected",
+ user_id="owner",
+ cron="0 * * * *",
+ config=config,
+ )
+
+
+@pytest.mark.parametrize("provider", ["linear", "attio"])
+async def test_workspace_provisioning_retry_reconnect_and_immutable_scope(
+ database, setup, provider
+):
+ ctx, service, request, account, lifecycle, schedules, workflows = setup
+ workspace, team1, team2 = uuid4(), uuid4(), uuid4()
+ config = {"workspace_id": str(uuid4())}
+ if provider == "linear":
+ config["team_ids"] = [str(team2), str(team1)]
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": provider,
+ "expected_identity": workspace.hex.upper(),
+ "config": config,
+ }
+ )
+ assert spec.expected_identity == str(workspace)
+ expected = spec.source_config()
+ assert expected["workspace_id"] == str(workspace)
+ assert native_principal(provider, expected) == (str(workspace), None)
+ if provider == "linear":
+ assert expected["team_ids"] == sorted([str(team1), str(team2)])
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = provider
+ creator._source_validation.seed_config_result(provider, expected)
+ request = request.model_copy(update={"source": spec})
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ epoch = (await db.get(Sync, first.sync_id)).writer_epoch
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, request) == first
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ assert (
+ await db.get(SourceConnection, first.source_connection_id)
+ ).config_fields == expected
+ changes = [{"expected_identity": str(uuid4())}]
+ if provider == "linear":
+ changes.append({"config": {"team_ids": [str(uuid4())]}})
+ for change in changes:
+ invalid = ManagedSource.model_validate({**spec.model_dump(), **change})
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Reconnect changes"):
+ await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(
+ update={
+ "generation": 2,
+ "source": invalid,
+ }
+ ),
+ )
+ async with database() as db:
+ unchanged = await service.get(db, ctx, account)
+ assert unchanged == first
+ assert unchanged.observed_generation == first.observed_generation == 1
+ assert (
+ await db.get(SourceConnection, first.source_connection_id)
+ ).config_fields == expected
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == epoch
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ reconnect_config = dict(config)
+ if provider == "linear":
+ reconnect_config["team_ids"] = list(reversed(config["team_ids"]))
+ reconnect = request.model_copy(
+ update={
+ "generation": 2,
+ "source": ManagedSource.model_validate(
+ {
+ **spec.model_dump(),
+ "connected_account_id": "ca_reconnected",
+ "config": reconnect_config,
+ }
+ ),
+ }
+ )
+ async with database() as db:
+ second = await service.ensure(db, ctx, account, reconnect)
+ assert second.state == "ready" and second.generation == 2
+ assert second.source_connection_id == first.source_connection_id
+ assert second.sync_id == first.sync_id
+ assert (await db.get(Sync, first.sync_id)).writer_epoch > epoch
+ original = await db.get(SourceConnection, first.source_connection_id)
+ assert original.config_fields == expected
+ assert original.auth_provider_config["account_id"] == "ca_reconnected"
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, reconnect) == second
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 2
+ assert lifecycle.create.await_count == 2
+
+
+@pytest.mark.parametrize(
+ "provider,identity,config",
+ [
+ ("linear", str(uuid4()), {}),
+ ("linear", str(uuid4()), {"team_ids": []}),
+ ("linear", str(uuid4()), {"team_ids": ["not-a-uuid"]}),
+ ("linear", "not-a-uuid", {"team_ids": [str(uuid4())]}),
+ ("attio", "not-a-uuid", {}),
+ ],
+)
+def test_workspace_provisioning_rejects_invalid_intent(provider, identity, config):
+ with pytest.raises(ValueError):
+ ManagedSource(
+ provider=provider,
+ expected_identity=identity,
+ config=config,
+ collection="owned",
+ auth_provider="composio",
+ connected_account_id="ca_test",
+ auth_config_id="ac_test",
+ user_id="owner",
+ cron="0 * * * *",
+ )
+
+
+GITHUB_REPOSITORIES = [
+ {"repository_id": 20, "owner_id": 30, "full_name": "team/two", "ref": "main"},
+ {"repository_id": 10, "owner_id": 30, "full_name": "team/one"},
+]
+
+
+async def test_github_provisioning_reconnect_preserves_principal_and_exact_scope(database, setup):
+ ctx, service, request, account, lifecycle, _, _ = setup
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": "github",
+ "expected_identity": "00042",
+ "config": {"expected_user_id": 999, "repositories": GITHUB_REPOSITORIES},
+ }
+ )
+ assert spec.expected_identity == "42"
+ config = spec.source_config()
+ assert config["expected_user_id"] == 42
+ assert [row["repository_id"] for row in config["repositories"]] == [10, 20]
+ assert native_principal("github", config) == ("42", None)
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = "github"
+ creator._source_validation.seed_config_result("github", config)
+ request = request.model_copy(update={"source": spec})
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ epoch = (await db.get(Sync, first.sync_id)).writer_epoch
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, request) == first
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ variants = [
+ {"expected_identity": "43"},
+ *(
+ {
+ "config": {
+ **spec.config,
+ "repositories": [{**GITHUB_REPOSITORIES[0], **change}, GITHUB_REPOSITORIES[1]],
+ }
+ }
+ for change in (
+ {"repository_id": 21},
+ {"owner_id": 31},
+ {"ref": "other"},
+ {"full_name": "team/renamed"},
+ )
+ ),
+ {"config": {**spec.config, "include_code": False}},
+ {"config": {**spec.config, "include_conversations": False}},
+ ]
+ for change in variants:
+ candidate = ManagedSource.model_validate({**spec.model_dump(), **change})
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Reconnect changes"):
+ await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(update={"generation": 2, "source": candidate}),
+ )
+ async with database() as db:
+ unchanged = await service.get(db, ctx, account)
+ assert unchanged == first and unchanged.observed_generation == 1
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == epoch
+ assert (
+ await db.get(SourceConnection, first.source_connection_id)
+ ).config_fields == config
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ reconnected = ManagedSource.model_validate(
+ {
+ **spec.model_dump(),
+ "connected_account_id": "ca_reconnected",
+ "config": {**spec.config, "repositories": list(reversed(GITHUB_REPOSITORIES))},
+ }
+ )
+ retry = request.model_copy(update={"generation": 2, "source": reconnected})
+ async with database() as db:
+ second = await service.ensure(db, ctx, account, retry)
+ assert second.source_connection_id == first.source_connection_id
+ assert second.sync_id == first.sync_id and second.observed_generation == 2
+ assert (await db.get(Sync, first.sync_id)).writer_epoch > epoch
+ connection = await db.get(SourceConnection, first.source_connection_id)
+ assert connection.auth_provider_config["account_id"] == "ca_reconnected"
+ assert connection.config_fields == config
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, retry) == second
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 2
+ assert lifecycle.create.await_count == 2
+
+
+@pytest.mark.parametrize(
+ "identity,config",
+ [
+ ("octocat", {"repositories": GITHUB_REPOSITORIES}),
+ ("0", {"repositories": GITHUB_REPOSITORIES}),
+ ("-42", {"repositories": GITHUB_REPOSITORIES}),
+ ("４２", {"repositories": GITHUB_REPOSITORIES}),
+ ("42", {}),
+ ("42", {"repositories": []}),
+ ("42", {"repo_name": "team/one"}),
+ ],
+)
+def test_github_provisioning_rejects_invalid_principal_or_missing_selection(identity, config):
+ with pytest.raises(ValueError):
+ ManagedSource(
+ provider="github",
+ expected_identity=identity,
+ config=config,
+ collection="owned",
+ auth_provider="composio",
+ connected_account_id="ca_test",
+ auth_config_id="ac_test",
+ user_id="owner",
+ cron="0 * * * *",
+ )
+
+
+async def test_notion_reconnect_retains_workspace_and_bot_with_exact_retry(database, setup):
+ ctx, service, request, account, lifecycle, _, _ = setup
+ workspace, bot = uuid4(), uuid4()
+ spec = ManagedSource.model_validate(
+ {
+ **request.source.model_dump(),
+ "provider": "notion",
+ "expected_identity": workspace.hex.upper(),
+ "expected_user_identity": bot.hex.upper(),
+ "config": {"expected_workspace_id": str(uuid4()), "expected_bot_id": str(uuid4())},
+ }
+ )
+ expected = spec.source_config()
+ assert expected == {"expected_workspace_id": str(workspace), "expected_bot_id": str(bot)}
+ assert native_principal("notion", expected) == (str(workspace), str(bot))
+ creator = service.store.create
+ creator._source_registry.get.return_value.short_name = "notion"
+ creator._source_validation.seed_config_result("notion", expected)
+ request = request.model_copy(update={"source": spec})
+ async with database() as db:
+ first = await service.ensure(db, ctx, account, request)
+ epoch = (await db.get(Sync, first.sync_id)).writer_epoch
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, request) == first
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ for identity_field in ("expected_identity", "expected_user_identity"):
+ changed = ManagedSource.model_validate({**spec.model_dump(), identity_field: str(uuid4())})
+ async with database() as db:
+ with pytest.raises(HTTPException, match="Reconnect changes original account identity"):
+ await service.ensure(
+ db,
+ ctx,
+ account,
+ request.model_copy(update={"generation": 2, "source": changed}),
+ )
+ async with database() as db:
+ assert await service.get(db, ctx, account) == first
+ assert (await db.get(Sync, first.sync_id)).writer_epoch == epoch
+ assert (
+ await db.get(SourceConnection, first.source_connection_id)
+ ).config_fields == expected
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 1
+ reconnect = request.model_copy(
+ update={
+ "generation": 2,
+ "source": spec.model_copy(update={"connected_account_id": "ca_reconnected"}),
+ }
+ )
+ async with database() as db:
+ second = await service.ensure(db, ctx, account, reconnect)
+ assert second.source_connection_id == first.source_connection_id
+ assert second.sync_id == first.sync_id and second.observed_generation == 2
+ original = await db.get(SourceConnection, first.source_connection_id)
+ assert original.config_fields == expected
+ assert original.auth_provider_config["account_id"] == "ca_reconnected"
+ async with database() as db:
+ assert await service.ensure(db, ctx, account, reconnect) == second
+ assert await db.scalar(select(func.count()).select_from(SourceConnection)) == 1
+ assert await db.scalar(select(func.count()).select_from(SyncJob)) == 2
+ assert lifecycle.create.await_count == 2
+
+
+@pytest.mark.parametrize(
+ "workspace,bot",
+ [("workspace-name", str(uuid4())), (str(uuid4()), None), (str(uuid4()), "bot-name")],
+)
+def test_notion_provisioning_requires_native_uuid_pair(workspace, bot):
+ with pytest.raises(ValueError):
+ ManagedSource(
+ provider="notion",
+ expected_identity=workspace,
+ expected_user_identity=bot,
+ collection="owned",
+ auth_provider="composio",
+ connected_account_id="ca_test",
+ auth_config_id="ac_test",
+ user_id="owner",
+ cron="0 * * * *",
+ )
diff --git a/backend/airweave/domains/search/adapters/vector_db/filter_translator.py b/backend/airweave/domains/search/adapters/vector_db/filter_translator.py
index 806ca89d5e..0c10e3601c 100644
--- a/backend/airweave/domains/search/adapters/vector_db/filter_translator.py
+++ b/backend/airweave/domains/search/adapters/vector_db/filter_translator.py
@@ -6,7 +6,7 @@
from __future__ import annotations
-from datetime import datetime
+from datetime import datetime, timezone
from typing import Any, List, Optional, Union
from airweave.core.logging import ContextualLogger
@@ -170,7 +170,7 @@ def _escape_string(self, value: str) -> str:
"""Escape special characters for YQL string literals."""
return value.replace("\\", "\\\\").replace("'", "\\'")
- def _parse_datetime_to_epoch(self, value: str, field: str) -> int:
+ def _parse_datetime_to_epoch(self, value: str, field: str) -> int | float:
"""Parse ISO datetime string to epoch seconds."""
try:
# Strip surrounding quotes — LLMs sometimes wrap datetime strings
@@ -178,7 +178,12 @@ def _parse_datetime_to_epoch(self, value: str, field: str) -> int:
if value.endswith("Z"):
value = value[:-1] + "+00:00"
dt = datetime.fromisoformat(value)
- return int(dt.timestamp())
+ # Existing filters accept zone-less ISO values; interpret them as UTC,
+ # never as the deployment host timezone. New product inputs are aware.
+ if dt.tzinfo is None:
+ dt = dt.replace(tzinfo=timezone.utc)
+ seconds = dt.timestamp()
+ return int(seconds) if seconds.is_integer() else seconds
except (ValueError, AttributeError) as e:
raise FilterTranslationError(
f"Failed to parse datetime '{value}' for field '{field}': {e}"
diff --git a/backend/airweave/domains/search/adapters/vector_db/tests/test_partial_coverage.py b/backend/airweave/domains/search/adapters/vector_db/tests/test_partial_coverage.py
new file mode 100644
index 0000000000..db0a78cadb
--- /dev/null
+++ b/backend/airweave/domains/search/adapters/vector_db/tests/test_partial_coverage.py
@@ -0,0 +1,38 @@
+"""Engine coverage survives the real adapter; only remote HTTP is mocked."""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB
+from airweave.domains.search.types import CompiledQuery
+
+
+@pytest.mark.parametrize(
+ "coverage,errors,expected",
+ [
+ ({"coverage": 100, "full": True}, [], False),
+ ({"coverage": 75}, [], True),
+ ({"coverage": 100, "degraded": {"timeout": True}}, [], True),
+ ({"coverage": 100, "full": False}, [], True),
+ ({"coverage": 100}, [{"code": 12}], True),
+ ],
+)
+async def test_partial_engine_response_preserved(coverage, errors, expected):
+ app = MagicMock()
+ response = MagicMock()
+ response.is_successful.return_value = True
+ response.json = {"root": {"coverage": coverage, "errors": errors}}
+ response.hits = []
+ app.query.return_value = response
+ client = VespaVectorDB(app, MagicMock(), MagicMock())
+ result = await client.execute_query(
+ CompiledQuery(
+ vector_db="vespa",
+ display="test",
+ raw={"yql": "select * from sources * where true", "params": {}},
+ )
+ )
+ assert result.engine_partial is expected
+ assert result.retrieval_incomplete is expected
+ assert result.engine_coverage_percent == coverage["coverage"]
diff --git a/backend/airweave/domains/search/adapters/vector_db/tests/test_vespa_client_name_substring.py b/backend/airweave/domains/search/adapters/vector_db/tests/test_vespa_client_name_substring.py
index cab65bb5ba..1e697854c2 100644
--- a/backend/airweave/domains/search/adapters/vector_db/tests/test_vespa_client_name_substring.py
+++ b/backend/airweave/domains/search/adapters/vector_db/tests/test_vespa_client_name_substring.py
@@ -128,3 +128,23 @@ async def test_clause_absent_when_name_substring_none(self) -> None:
await client.filter_search(filter_groups=[], collection_id="col-123")
yql = app.query.call_args.kwargs["body"]["yql"]
assert "name matches" not in yql
+
+
+@pytest.mark.parametrize("fragment", ["… 原文", None, {"invalid": "shape"}])
+def test_optional_query_snippet_never_replaces_full_chunk(fragment):
+ client, _ = _make_client()
+ result = client._convert_hits_to_results(
+ [
+ {
+ "fields": {
+ "entity_id": "chunk",
+ "name": "Original",
+ "textual_representation": "Full retained chunk",
+ "query_snippet": fragment,
+ }
+ }
+ ]
+ ).results[0]
+ assert result.textual_representation == "Full retained chunk"
+ assert result.query_snippet == (fragment if isinstance(fragment, str) else None)
+ assert "query_snippet" not in result.model_dump()
diff --git a/backend/airweave/domains/search/adapters/vector_db/vespa_client.py b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py
index b8cef0ba0b..f9b5eefc3c 100644
--- a/backend/airweave/domains/search/adapters/vector_db/vespa_client.py
+++ b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py
@@ -155,7 +155,14 @@ async def execute_query(
f"coverage={coverage_pct:.1f}%"
)
- return self._convert_hits_to_results(hits)
+ results = self._convert_hits_to_results(hits)
+ results.engine_coverage_percent = coverage.get("coverage")
+ results.engine_partial = bool(
+ coverage.get("degraded") or coverage.get("full") is False
+ or coverage.get("coverage", 100) < 100 or root.get("errors")
+ )
+ results.retrieval_incomplete = results.engine_partial
+ return results
async def count(
self,
@@ -377,6 +384,10 @@ def _build_params(
RetrievalStrategy.SEMANTIC,
RetrievalStrategy.HYBRID,
):
+ # Filter-first graph traversal can miss all eligible neighbors in
+ # selectively scoped collections. Keep pre-filtering, but compute
+ # distances before checking the filter to preserve candidate recall.
+ params["ranking.matching.filterFirstThreshold"] = 0.0
for i, dense_emb in enumerate(embeddings.dense_embeddings):
params[f"input.query(q{i})"] = {"values": dense_emb.vector}
@@ -452,9 +463,16 @@ def _convert_hits_to_results(self, hits: List[Dict[str, Any]]) -> SearchResults:
textual_representation=self._get_required_field(
fields, "textual_representation", entity_id
),
+ query_snippet=(
+ fields.get("query_snippet")
+ if isinstance(fields.get("query_snippet"), str)
+ else None
+ ),
airweave_system_metadata=self._extract_system_metadata(fields, entity_id),
access=self._extract_access_control(fields),
- web_url=self._get_required_field(raw_source_fields, "web_url", entity_id),
+ # Native snapshots can be opened by canonical record ID without
+ # a provider web URL. Absence is valid, not a malformed index hit.
+ web_url=str(raw_source_fields.get("web_url") or ""),
url=fields.get("url"),
raw_source_fields=raw_source_fields,
)
diff --git a/backend/airweave/domains/search/agentic/agent.py b/backend/airweave/domains/search/agentic/agent.py
index 75b0da96a2..e7bf6d43a1 100644
--- a/backend/airweave/domains/search/agentic/agent.py
+++ b/backend/airweave/domains/search/agentic/agent.py
@@ -47,6 +47,7 @@
from airweave.core.protocols.reranker import RerankerProtocol
from airweave.core.protocols.tokenizer import TokenizerProtocol
from airweave.domains.collections.protocols import CollectionRepositoryProtocol
+from airweave.domains.search.adapters.vector_db.exceptions import VectorDBError
from airweave.domains.search.adapters.vector_db.protocol import VectorDBProtocol
from airweave.domains.search.agentic.context_manager import ContextManager
from airweave.domains.search.agentic.exceptions import (
@@ -85,12 +86,15 @@
ToolErrorResult,
ToolName,
)
+from airweave.domains.search.canonical_visibility import visible_results
from airweave.domains.search.config import SearchConfig
from airweave.domains.search.protocols import (
CollectionMetadataBuilderProtocol,
SearchPlanExecutorProtocol,
)
from airweave.domains.search.types import SearchResults
+from airweave.domains.search.visible_vector_db import VisibleVectorDB
+from airweave.domains.sources.protocols import SourceRegistryProtocol
if TYPE_CHECKING:
from airweave.schemas.search_v2 import AgenticSearchRequest
@@ -110,6 +114,7 @@ def __init__(
collection_repo: CollectionRepositoryProtocol,
event_bus: EventBus,
config: SearchConfig,
+ source_registry: SourceRegistryProtocol,
) -> None:
"""Initialize with all dependencies (injected by service)."""
self._llm = llm
@@ -117,6 +122,7 @@ def __init__(
self._reranker = reranker
self._executor = executor
self._vector_db = vector_db
+ self._source_registry = source_registry
self._metadata_builder = metadata_builder
self._collection_repo = collection_repo
self._event_bus = event_bus
@@ -238,6 +244,14 @@ async def _run( # noqa: C901 — agent loop orchestration is inherently complex
for iteration in range(max_iter):
diag.iteration = iteration
+ # Never resend cached text after its authoritative publication changes.
+ current = await visible_results(
+ db, ctx.organization.id, readable_id, list(state.results.values()),
+ self._source_registry,
+ )
+ if len(current) != len(state.results):
+ raise VectorDBError("Source content changed during search; retry the query")
+
# 1. Call LLM
llm_start = time.monotonic()
response = await self._llm.chat(
@@ -395,6 +409,14 @@ async def _run( # noqa: C901 — agent loop orchestration is inherently complex
state.results[eid] for eid in state.collected_ids if eid in state.results
]
+ collected_before_validation = len(collected_results)
+ collected_results = await visible_results(
+ db, ctx.organization.id, readable_id, collected_results, self._source_registry
+ )
+ excluded_at_finish = collected_before_validation - len(collected_results)
+ state.retrieval_incomplete |= excluded_at_finish > 0
+ state.excluded_candidates += excluded_at_finish
+
# Optional reranking
if self._reranker and collected_results:
rerank_start = time.monotonic()
@@ -472,7 +494,11 @@ async def _run( # noqa: C901 — agent loop orchestration is inherently complex
)
)
- return SearchResults(results=collected_results)
+ return SearchResults(
+ results=collected_results,
+ retrieval_incomplete=state.retrieval_incomplete,
+ excluded_candidates=state.excluded_candidates,
+ )
@staticmethod
def _append_iteration_messages(
@@ -599,6 +625,9 @@ def _build_dispatcher(
user_principal: str | None = None,
) -> ToolDispatcher:
"""Construct tools and dispatcher for this request."""
+ visible_db = VisibleVectorDB(
+ self._vector_db, db, ctx, collection_readable_id, collection_id, self._source_registry
+ )
return ToolDispatcher(
{
ToolName.SEARCH: SearchTool(
@@ -611,27 +640,27 @@ def _build_dispatcher(
user_principal=user_principal,
),
ToolName.READ: ReadTool(
- vector_db=self._vector_db,
+ vector_db=visible_db,
collection_id=collection_id,
surrounding_chunks=self._config.READ_SURROUNDING_CHUNKS,
),
ToolName.ADD_TO_RESULTS: AddToResultsTool(),
ToolName.REMOVE_FROM_RESULTS: RemoveFromResultsTool(),
ToolName.COUNT: CountTool(
- vector_db=self._vector_db,
+ vector_db=visible_db,
collection_id=collection_id,
user_filter=user_filter,
),
ToolName.GET_CHILDREN: GetChildrenTool(
- vector_db=self._vector_db,
+ vector_db=visible_db,
collection_id=collection_id,
),
ToolName.GET_SIBLINGS: GetSiblingsTool(
- vector_db=self._vector_db,
+ vector_db=visible_db,
collection_id=collection_id,
),
ToolName.GET_PARENT: GetParentTool(
- vector_db=self._vector_db,
+ vector_db=visible_db,
collection_id=collection_id,
),
ToolName.REVIEW_RESULTS: ReviewResultsTool(),
diff --git a/backend/airweave/domains/search/agentic/context_manager.py b/backend/airweave/domains/search/agentic/context_manager.py
index b10cb59710..643e2c4687 100644
--- a/backend/airweave/domains/search/agentic/context_manager.py
+++ b/backend/airweave/domains/search/agentic/context_manager.py
@@ -288,10 +288,17 @@ def fit_tool_result(
def _fit_search_result(self, result: SearchToolResult, available_tokens: int) -> str:
"""Fit search summaries into budget, dropping from bottom."""
+ notice = ""
+ if result.retrieval_incomplete or result.excluded_candidates:
+ notice = (
+ f"Retrieval incomplete: {result.excluded_candidates} stale or inaccessible "
+ "candidates were excluded. An empty or short page does not prove "
+ "that no other matches exist.\n\n"
+ )
if not result.summaries:
- return "No results found."
+ return notice + ("No currently visible results." if notice else "No results found.")
- header = f"**{len(result.summaries)} results** ({result.new_count} new):\n\n"
+ header = notice + f"**{len(result.summaries)} results** ({result.new_count} new):\n\n"
header_tokens = self._tokenizer.count_tokens(header)
remaining = available_tokens - header_tokens
diff --git a/backend/airweave/domains/search/agentic/service.py b/backend/airweave/domains/search/agentic/service.py
index 7f15fd193e..4db02e07d8 100644
--- a/backend/airweave/domains/search/agentic/service.py
+++ b/backend/airweave/domains/search/agentic/service.py
@@ -24,6 +24,7 @@
SearchPlanExecutorProtocol,
)
from airweave.domains.search.types import SearchResults
+from airweave.domains.sources.protocols import SourceRegistryProtocol
if TYPE_CHECKING:
from airweave.schemas.search_v2 import AgenticSearchRequest
@@ -42,6 +43,7 @@ def __init__(
metadata_builder: CollectionMetadataBuilderProtocol,
collection_repo: CollectionRepositoryProtocol,
event_bus: EventBus,
+ source_registry: SourceRegistryProtocol,
) -> None:
"""Initialize with all dependencies needed to construct agents."""
self._llm = llm
@@ -49,6 +51,7 @@ def __init__(
self._reranker = reranker
self._executor = executor
self._vector_db = vector_db
+ self._source_registry = source_registry
self._metadata_builder = metadata_builder
self._collection_repo = collection_repo
self._event_bus = event_bus
@@ -61,6 +64,7 @@ def with_llm(self, llm: LLMProtocol) -> "AgenticSearchService":
reranker=self._reranker,
executor=self._executor,
vector_db=self._vector_db,
+ source_registry=self._source_registry,
metadata_builder=self._metadata_builder,
collection_repo=self._collection_repo,
event_bus=self._event_bus,
@@ -81,6 +85,7 @@ async def search(
reranker=self._reranker,
executor=self._executor,
vector_db=self._vector_db,
+ source_registry=self._source_registry,
metadata_builder=self._metadata_builder,
collection_repo=self._collection_repo,
event_bus=self._event_bus,
diff --git a/backend/airweave/domains/search/agentic/state.py b/backend/airweave/domains/search/agentic/state.py
index cf0b877ece..6249449b47 100644
--- a/backend/airweave/domains/search/agentic/state.py
+++ b/backend/airweave/domains/search/agentic/state.py
@@ -29,6 +29,9 @@ def __init__(self) -> None:
self.collected_ids: set[str] = set()
self.should_finish: bool = False
self.return_warned: bool = False
+ self.retrieval_incomplete: bool = False
+ # Candidate occurrences across calls, not a unique-record total.
+ self.excluded_candidates: int = 0
def add_to_collected(self, entity_ids: list[str]) -> tuple[list[str], list[str], list[str]]:
"""Add entity IDs to the collected result set.
diff --git a/backend/airweave/domains/search/agentic/tests/test_agent.py b/backend/airweave/domains/search/agentic/tests/test_agent.py
index bf13b4b20e..56a9dbacde 100644
--- a/backend/airweave/domains/search/agentic/tests/test_agent.py
+++ b/backend/airweave/domains/search/agentic/tests/test_agent.py
@@ -33,11 +33,24 @@
from airweave.domains.search.fakes.executor import FakeSearchPlanExecutor
from airweave.domains.search.fakes.metadata_builder import FakeCollectionMetadataBuilder
from airweave.domains.search.types import SearchResults
+from airweave.domains.sources.fakes.registry import FakeSourceRegistry
from airweave.models.collection import Collection
from airweave.schemas.search_v2 import AgenticSearchRequest
# ── Test constants ────────────────────────────────────────────────────
+
+@pytest.fixture(autouse=True)
+def isolated_publication_gate(monkeypatch):
+ """Synthetic orchestration fixtures do not query the real publication database."""
+
+ async def passthrough(db, organization_id, collection, results, registry):
+ return results
+
+ monkeypatch.setattr("airweave.domains.search.visible_vector_db.visible_results", passthrough)
+ monkeypatch.setattr("airweave.domains.search.agentic.agent.visible_results", passthrough)
+
+
DEFAULT_ORG_ID = uuid4()
DEFAULT_COLLECTION_ID = uuid4()
DEFAULT_READABLE_ID = "test-col"
@@ -167,6 +180,7 @@ def _build_agent(
reranker=None,
executor=executor or FakeSearchPlanExecutor(),
vector_db=FakeVectorDB(),
+ source_registry=FakeSourceRegistry(),
metadata_builder=metadata_builder,
collection_repo=collection_repo,
event_bus=event_bus or FakeEventBus(),
@@ -448,9 +462,7 @@ async def test_llm_fatal_error_emits_failed_event(self) -> None:
async def test_llm_all_providers_failed_emits_failed_event(self) -> None:
"""LLMAllProvidersFailedError → SearchFailedEvent published."""
llm = FakeLLM(make_model_spec())
- llm.seed_error(
- LLMAllProvidersFailedError([("p1", LLMFatalError("down", provider="p1"))])
- )
+ llm.seed_error(LLMAllProvidersFailedError([("p1", LLMFatalError("down", provider="p1"))]))
event_bus = FakeEventBus()
agent = _build_agent(llm=llm, event_bus=event_bus)
@@ -622,3 +634,96 @@ async def test_no_reranker_returns_unranked(self) -> None:
# Completed event should still be emitted
completed = event_bus.get_events("search.completed")
assert len(completed) == 1
+
+
+@pytest.mark.asyncio
+async def test_revoked_cached_text_stops_before_next_llm_call(monkeypatch):
+ """A result found on turn one must be revalidated before turn two sees its text."""
+ from airweave.domains.search.adapters.vector_db.exceptions import VectorDBError
+
+ executor = FakeSearchPlanExecutor()
+ executor.seed_result(SearchResults(results=[make_result("revoked", content="private text")]))
+ llm = FakeLLM(make_model_spec())
+ llm.seed_tool_response(_make_search_response(tool_calls=[_search_tool_call()]))
+ llm.seed_tool_response(_make_search_response(tool_calls=[_return_results_tool_call()]))
+ gate = AsyncMock(return_value=[])
+ monkeypatch.setattr("airweave.domains.search.agentic.agent.visible_results", gate)
+ agent = _build_agent(llm=llm, executor=executor)
+ with pytest.raises(VectorDBError, match="changed during search"):
+ await agent.run(AsyncMock(), _make_ctx(), DEFAULT_READABLE_ID, _make_request())
+ assert len([call for call in llm._calls if call[0] == "chat"]) == 1
+ assert gate.await_count == 2
+ assert gate.await_args.args[3][0].entity_id == "revoked"
+
+
+def test_all_original_read_and_navigation_tools_receive_visibility_wrapper():
+ """Composition must not accidentally hand the raw adapter to one navigation tool."""
+ from airweave.domains.search.visible_vector_db import VisibleVectorDB
+
+ agent = _build_agent()
+ dispatcher = agent._build_dispatcher(
+ str(DEFAULT_COLLECTION_ID),
+ [],
+ AsyncMock(),
+ _make_ctx(),
+ DEFAULT_READABLE_ID,
+ )
+ for name in ("read", "count", "get_children", "get_siblings", "get_parent"):
+ assert isinstance(dispatcher._tools[name]._vector_db, VisibleVectorDB)
+
+
+@pytest.mark.asyncio
+async def test_final_response_preserves_search_exclusions_and_finish_time_drop(monkeypatch):
+ result = make_result("fresh-until-finish")
+ executor = FakeSearchPlanExecutor()
+ executor.seed_result(
+ SearchResults(results=[result], retrieval_incomplete=True, excluded_candidates=2)
+ )
+ llm = FakeLLM(make_model_spec())
+ llm.seed_tool_response(_make_search_response(tool_calls=[_search_tool_call()]))
+ llm.seed_tool_response(
+ _make_search_response(
+ tool_calls=[
+ _add_results_tool_call(entity_ids=[result.entity_id]),
+ _return_results_tool_call(),
+ ]
+ )
+ )
+ gate = AsyncMock(side_effect=[[], [result], []])
+ monkeypatch.setattr("airweave.domains.search.agentic.agent.visible_results", gate)
+ response = await _build_agent(llm=llm, executor=executor).run(
+ AsyncMock(),
+ _make_ctx(),
+ DEFAULT_READABLE_ID,
+ _make_request(),
+ )
+ assert response.results == []
+ assert response.retrieval_incomplete is True
+ assert response.excluded_candidates == 3
+
+
+@pytest.mark.asyncio
+async def test_unavailable_count_allows_agent_to_continue_with_other_tools(monkeypatch):
+ from airweave.domains.search.visible_vector_db import UnavailableExactCount
+
+ monkeypatch.setattr(
+ "airweave.domains.search.visible_vector_db.VisibleVectorDB.count",
+ AsyncMock(side_effect=UnavailableExactCount("exact count unavailable")),
+ )
+ llm = FakeLLM(make_model_spec())
+ llm.seed_tool_response(
+ _make_search_response(
+ tool_calls=[
+ LLMToolCall(id="count-1", name="count", arguments={"filter_groups": []}),
+ ]
+ )
+ )
+ llm.seed_tool_response(_make_search_response(tool_calls=[_return_results_tool_call()]))
+ response = await _build_agent(llm=llm).run(
+ AsyncMock(),
+ _make_ctx(),
+ DEFAULT_READABLE_ID,
+ _make_request(),
+ )
+ assert response.results == []
+ assert len([call for call in llm._calls if call[0] == "chat"]) == 2
diff --git a/backend/airweave/domains/search/agentic/tests/test_context_manager.py b/backend/airweave/domains/search/agentic/tests/test_context_manager.py
index 8f562b93b6..c0d3dec71c 100644
--- a/backend/airweave/domains/search/agentic/tests/test_context_manager.py
+++ b/backend/airweave/domains/search/agentic/tests/test_context_manager.py
@@ -311,8 +311,9 @@ def test_triggers_emergency_compress_when_nearly_full(self) -> None:
messages = [{"role": "user", "content": "query"}]
for i in range(80):
messages.append({"role": "assistant", "content": f"thinking {i}" * 30})
- messages.append({"role": "tool", "tool_call_id": f"tc-{i}",
- "content": f"result {i}" * 30})
+ messages.append(
+ {"role": "tool", "tool_call_id": f"tc-{i}", "content": f"result {i}" * 30}
+ )
# Messages should be large enough to trigger emergency compress
assert cm.input_tokens(messages) > cm._context_window - cm._MIN_OUTPUT_TOKENS
@@ -350,11 +351,20 @@ def test_counts_tool_calls_in_input(self) -> None:
messages_without = [{"role": "user", "content": "query"}]
messages_with = [
{"role": "user", "content": "query"},
- {"role": "assistant", "content": "", "tool_calls": [
- {"id": "tc-1", "type": "function", "function": {
- "name": "search", "arguments": "x" * 4_000,
- }}
- ]},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "tc-1",
+ "type": "function",
+ "function": {
+ "name": "search",
+ "arguments": "x" * 4_000,
+ },
+ }
+ ],
+ },
]
result_without = cm.max_output_tokens(messages_without)
result_with = cm.max_output_tokens(messages_with)
@@ -429,8 +439,9 @@ def test_actually_reduces_token_count(self) -> None:
messages = [{"role": "user", "content": "query"}]
for i in range(20):
messages.append({"role": "assistant", "content": f"thinking {'x' * 500} {i}"})
- messages.append({"role": "tool", "tool_call_id": f"tc-{i}",
- "content": f"result {'y' * 500} {i}"})
+ messages.append(
+ {"role": "tool", "tool_call_id": f"tc-{i}", "content": f"result {'y' * 500} {i}"}
+ )
tokens_before = cm.input_tokens(messages)
result = cm.emergency_compress(messages)
@@ -451,7 +462,6 @@ def test_reserves_output_space(self) -> None:
messages = [{"role": "user", "content": "query"}]
budget = cm.available_budget(messages)
- max_out = cm.max_output_tokens(messages)
# budget + output + input should not exceed context_window
input_tokens = cm.input_tokens(messages)
@@ -498,16 +508,54 @@ def test_counts_thinking(self) -> None:
"""_thinking field on assistant messages is counted."""
cm = make_context_mgr()
without = cm.input_tokens([{"role": "assistant", "content": "hi"}])
- with_thinking = cm.input_tokens([
- {"role": "assistant", "content": "hi", "_thinking": "x" * 1000}
- ])
+ with_thinking = cm.input_tokens(
+ [{"role": "assistant", "content": "hi", "_thinking": "x" * 1000}]
+ )
assert with_thinking > without
def test_counts_tool_calls(self) -> None:
"""tool_calls on assistant messages are counted."""
cm = make_context_mgr()
without = cm.input_tokens([{"role": "assistant", "content": "hi"}])
- with_tc = cm.input_tokens([{"role": "assistant", "content": "hi", "tool_calls": [
- {"id": "tc-1", "type": "function", "function": {"name": "s", "arguments": "x" * 500}}
- ]}])
+ with_tc = cm.input_tokens(
+ [
+ {
+ "role": "assistant",
+ "content": "hi",
+ "tool_calls": [
+ {
+ "id": "tc-1",
+ "type": "function",
+ "function": {"name": "s", "arguments": "x" * 500},
+ }
+ ],
+ }
+ ]
+ )
assert with_tc > without
+
+
+def test_empty_filtered_search_keeps_incomplete_notice():
+ cm = make_context_mgr()
+ result = SearchToolResult(
+ summaries=[], new_count=0, retrieval_incomplete=True, excluded_candidates=4
+ )
+ content = cm.fit_tool_result(result, 1000)
+ assert "Retrieval incomplete" in content
+ assert "4 stale or inaccessible" in content
+ assert "No results found" not in content
+
+
+def test_incomplete_notice_survives_context_budget_truncation():
+ cm = make_context_mgr()
+ result = SearchToolResult(
+ summaries=[
+ make_rendered_result(entity_id=str(i), text="long text " * 100) for i in range(3)
+ ],
+ new_count=3,
+ retrieval_incomplete=True,
+ excluded_candidates=2,
+ )
+ content = cm.fit_tool_result(result, 50)
+ assert "Retrieval incomplete" in content
+ assert "context budget" in content
diff --git a/backend/airweave/domains/search/agentic/tests/test_tools.py b/backend/airweave/domains/search/agentic/tests/test_tools.py
index e45a3347bb..10ed29b5ab 100644
--- a/backend/airweave/domains/search/agentic/tests/test_tools.py
+++ b/backend/airweave/domains/search/agentic/tests/test_tools.py
@@ -38,7 +38,14 @@ async def test_new_results_tracked_in_state(self) -> None:
executor = FakeSearchPlanExecutor()
executor.seed_result(SearchResults(results=[r1, r2]))
- tool = SearchTool(executor=executor, user_filter=[], collection_id="col-1", db=AsyncMock(), ctx=AsyncMock(), collection_readable_id="col-readable")
+ tool = SearchTool(
+ executor=executor,
+ user_filter=[],
+ collection_id="col-1",
+ db=AsyncMock(),
+ ctx=AsyncMock(),
+ collection_readable_id="col-readable",
+ )
state = make_state()
result = await tool.execute(
{
@@ -64,7 +71,14 @@ async def test_duplicate_results_not_counted_as_new(self) -> None:
executor.seed_result(SearchResults(results=[r1, r2]))
state = make_state(results={"ent-1": r1})
- tool = SearchTool(executor=executor, user_filter=[], collection_id="col-1", db=AsyncMock(), ctx=AsyncMock(), collection_readable_id="col-readable")
+ tool = SearchTool(
+ executor=executor,
+ user_filter=[],
+ collection_id="col-1",
+ db=AsyncMock(),
+ ctx=AsyncMock(),
+ collection_readable_id="col-readable",
+ )
result = await tool.execute(
{
"query": {"primary": "test"},
@@ -86,7 +100,14 @@ async def test_lineage_tracked(self) -> None:
executor.seed_result(SearchResults(results=[r1]))
state = make_state()
- tool = SearchTool(executor=executor, user_filter=[], collection_id="col-1", db=AsyncMock(), ctx=AsyncMock(), collection_readable_id="col-readable")
+ tool = SearchTool(
+ executor=executor,
+ user_filter=[],
+ collection_id="col-1",
+ db=AsyncMock(),
+ ctx=AsyncMock(),
+ collection_readable_id="col-readable",
+ )
await tool.execute(
{
"query": {"primary": "test"},
@@ -555,3 +576,43 @@ async def test_dispatcher_passes_through_tool_validation_error(self) -> None:
with pytest.raises(ToolValidationError, match="bad args"):
await dispatcher.dispatch(tc, make_state())
+
+
+@pytest.mark.asyncio
+async def test_search_accumulates_excluded_candidates_without_claiming_empty_complete():
+ executor = FakeSearchPlanExecutor()
+ executor.seed_result_sequence(
+ [
+ SearchResults(results=[], retrieval_incomplete=True, excluded_candidates=2),
+ SearchResults(results=[], retrieval_incomplete=True, excluded_candidates=3),
+ ]
+ )
+ tool = SearchTool(executor, [], "col", AsyncMock(), AsyncMock(), "readable")
+ state = make_state()
+ first = await tool.execute(
+ {"query": {"primary": "first"}, "limit": 10, "offset": 0, "retrieval_strategy": "hybrid"},
+ state,
+ )
+ second = await tool.execute(
+ {"query": {"primary": "second"}, "limit": 10, "offset": 0, "retrieval_strategy": "hybrid"},
+ state,
+ )
+ assert first.retrieval_incomplete is True
+ assert first.excluded_candidates == 2
+ assert second.excluded_candidates == 3
+ assert state.retrieval_incomplete is True
+ assert state.excluded_candidates == 5
+
+
+@pytest.mark.asyncio
+async def test_unavailable_exact_count_is_correctable_not_a_transient_outage():
+ from airweave.domains.search.agentic.exceptions import ToolExecutionError, ToolTransientError
+ from airweave.domains.search.visible_vector_db import UnavailableExactCount
+
+ vector = AsyncMock()
+ vector.count.side_effect = UnavailableExactCount("not supported")
+ tool = CountTool(vector, "collection", [])
+ with pytest.raises(ToolExecutionError, match="Search or browse") as failure:
+ await tool.execute({"filter_groups": []}, make_state())
+ assert not isinstance(failure.value, ToolTransientError)
+ assert "do not infer zero" in str(failure.value)
diff --git a/backend/airweave/domains/search/agentic/tools/count.py b/backend/airweave/domains/search/agentic/tools/count.py
index 56092edabd..070640b8df 100644
--- a/backend/airweave/domains/search/agentic/tools/count.py
+++ b/backend/airweave/domains/search/agentic/tools/count.py
@@ -5,10 +5,12 @@
from typing import Any
from airweave.domains.search.adapters.vector_db.protocol import VectorDBProtocol
+from airweave.domains.search.agentic.exceptions import ToolExecutionError
from airweave.domains.search.agentic.state import AgentState
from airweave.domains.search.agentic.tools.dispatcher import Tool
from airweave.domains.search.agentic.tools.types import CountToolResult
from airweave.domains.search.types.filters import FilterGroup
+from airweave.domains.search.visible_vector_db import UnavailableExactCount
COUNT_TOOL: dict[str, Any] = {
"type": "function",
@@ -76,5 +78,12 @@ async def execute(
else:
combined = validated
- count = await self._vector_db.count(combined, self._collection_id)
+ try:
+ count = await self._vector_db.count(combined, self._collection_id)
+ except UnavailableExactCount as exc:
+ raise ToolExecutionError(
+ "Exact current count is unavailable for this indexed collection. "
+ "Search or browse the available records instead; do not infer zero. "
+ "Retrying count does not resolve this capability limitation."
+ ) from exc
return CountToolResult(count=count)
diff --git a/backend/airweave/domains/search/agentic/tools/read.py b/backend/airweave/domains/search/agentic/tools/read.py
index 771f8b40fb..ea112b2b8d 100644
--- a/backend/airweave/domains/search/agentic/tools/read.py
+++ b/backend/airweave/domains/search/agentic/tools/read.py
@@ -85,6 +85,9 @@ async def execute(
for orig_id, group_results in groups.items():
chunks = await self._fetch_chunks(orig_id, group_results, state)
+ if not chunks:
+ not_found.extend(result.entity_id for result in group_results)
+ continue
all_read_results.extend(chunks)
# Side-load fetched chunks into state
@@ -156,14 +159,11 @@ async def _fetch_chunks(
)
]
- try:
- results = await self._vector_db.filter_search(
- filter_groups=filter_groups,
- collection_id=self._collection_id,
- limit=max(max_chunk - min_chunk + 1, 10),
- )
- except Exception:
- results = list(group_results)
+ results = await self._vector_db.filter_search(
+ filter_groups=filter_groups,
+ collection_id=self._collection_id,
+ limit=max(max_chunk - min_chunk + 1, 10),
+ )
results.sort(key=lambda r: r.airweave_system_metadata.chunk_index)
return results
diff --git a/backend/airweave/domains/search/agentic/tools/search.py b/backend/airweave/domains/search/agentic/tools/search.py
index 20a6e53c17..3c5710359e 100644
--- a/backend/airweave/domains/search/agentic/tools/search.py
+++ b/backend/airweave/domains/search/agentic/tools/search.py
@@ -72,6 +72,11 @@ async def execute(
user_principal=self._user_principal,
)
+ state.retrieval_incomplete |= (
+ results.retrieval_incomplete or results.excluded_candidates > 0
+ )
+ state.excluded_candidates += results.excluded_candidates
+
# Track new results in state
new_count = 0
for r in results.results:
@@ -93,4 +98,6 @@ async def execute(
new_count=new_count,
requested_limit=plan.limit,
requested_offset=plan.offset,
+ retrieval_incomplete=results.retrieval_incomplete,
+ excluded_candidates=results.excluded_candidates,
)
diff --git a/backend/airweave/domains/search/agentic/tools/types.py b/backend/airweave/domains/search/agentic/tools/types.py
index e18bc73e55..784749063b 100644
--- a/backend/airweave/domains/search/agentic/tools/types.py
+++ b/backend/airweave/domains/search/agentic/tools/types.py
@@ -57,6 +57,8 @@ class SearchToolResult:
new_count: int
requested_limit: int = 0 # for pagination warning
requested_offset: int = 0 # for pagination warning
+ retrieval_incomplete: bool = False
+ excluded_candidates: int = 0
@dataclass(frozen=True)
diff --git a/backend/airweave/domains/search/browse/service.py b/backend/airweave/domains/search/browse/service.py
index 8125d279ae..6639f5fac5 100644
--- a/backend/airweave/domains/search/browse/service.py
+++ b/backend/airweave/domains/search/browse/service.py
@@ -6,7 +6,6 @@
from __future__ import annotations
-import asyncio
import time
from typing import TYPE_CHECKING
@@ -23,6 +22,8 @@
FilterGroup,
FilterOperator,
)
+from airweave.domains.search.visible_vector_db import UnavailableExactCount, VisibleVectorDB
+from airweave.domains.sources.protocols import SourceRegistryProtocol
from airweave.schemas.search_v2 import BrowseResponse
if TYPE_CHECKING:
@@ -36,9 +37,11 @@ def __init__(
self,
vector_db: VectorDBProtocol,
collection_repo: CollectionRepositoryProtocol,
+ source_registry: SourceRegistryProtocol,
) -> None:
"""Initialize with vector DB and collection repository."""
self._vector_db = vector_db
+ self._source_registry = source_registry
self._collection_repo = collection_repo
async def browse(
@@ -62,20 +65,19 @@ async def browse(
collection_id = str(collection.id)
name_substring = (request.name_query or "").strip() or None
- results, total = await asyncio.gather(
- self._vector_db.filter_search(
- filter_groups=filter_groups,
- collection_id=collection_id,
- limit=request.limit,
- offset=request.offset,
- name_substring=name_substring,
- ),
- self._vector_db.count(
- filter_groups=filter_groups,
- collection_id=collection_id,
- name_substring=name_substring,
- ),
+ visible_db = VisibleVectorDB(
+ self._vector_db, db, ctx, readable_id, collection_id, self._source_registry
)
+ # One SQL session is never used concurrently by validation/count.
+ retrieved = await self._vector_db.filter_search(
+ filter_groups=filter_groups, collection_id=collection_id,
+ limit=request.limit, offset=request.offset, name_substring=name_substring,
+ )
+ results = await visible_db.validate(retrieved)
+ try:
+ total = await visible_db.count(filter_groups, collection_id, name_substring)
+ except UnavailableExactCount:
+ total = None
duration_ms = int((time.monotonic() - start) * 1000)
ctx.logger.info(
@@ -86,6 +88,7 @@ async def browse(
return BrowseResponse(
results=results,
total=total,
+ retrieval_incomplete=len(results) != len(retrieved),
limit=request.limit,
offset=request.offset,
)
diff --git a/backend/airweave/domains/search/browse/tests/test_service.py b/backend/airweave/domains/search/browse/tests/test_service.py
index 713da28713..4a5fce6743 100644
--- a/backend/airweave/domains/search/browse/tests/test_service.py
+++ b/backend/airweave/domains/search/browse/tests/test_service.py
@@ -3,6 +3,7 @@
from __future__ import annotations
from datetime import datetime, timezone
+from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4
@@ -16,15 +17,28 @@
from airweave.domains.search.adapters.vector_db.fakes.vector_db import FakeVectorDB
from airweave.domains.search.agentic.tests.conftest import make_result
from airweave.domains.search.browse.service import BrowseService
+from airweave.domains.search.tests.test_executor import _make_registry_entry
from airweave.domains.search.types.filters import (
+ FilterableField,
FilterCondition,
FilterGroup,
FilterOperator,
- FilterableField,
)
+from airweave.domains.sources.fakes.registry import FakeSourceRegistry
from airweave.models.collection import Collection
from airweave.schemas.search_v2 import BrowseRequest
+
+@pytest.fixture(autouse=True)
+def isolated_publication_gate(monkeypatch):
+ """Synthetic orchestration fixtures do not query the real publication database."""
+
+ async def passthrough(db, organization_id, collection, results, registry):
+ return results
+
+ monkeypatch.setattr("airweave.domains.search.visible_vector_db.visible_results", passthrough)
+
+
DEFAULT_ORG_ID = uuid4()
DEFAULT_COLLECTION_ID = uuid4()
DEFAULT_READABLE_ID = "test-col"
@@ -64,10 +78,20 @@ def _make_collection() -> Collection:
return col
+def _make_db():
+ db = AsyncMock()
+ db.scalars.return_value = [SimpleNamespace(short_name="legacy", sync_id=uuid4())]
+ return db
+
+
def _make_service() -> tuple[BrowseService, FakeVectorDB, FakeCollectionRepository]:
vector_db = FakeVectorDB()
collection_repo = FakeCollectionRepository()
- svc = BrowseService(vector_db=vector_db, collection_repo=collection_repo)
+ registry = FakeSourceRegistry()
+ registry.seed(_make_registry_entry("legacy"))
+ svc = BrowseService(
+ vector_db=vector_db, collection_repo=collection_repo, source_registry=registry
+ )
return svc, vector_db, collection_repo
@@ -84,7 +108,7 @@ async def test_happy_path_returns_response(self) -> None:
vector_db.seed_count(42)
response = await svc.browse(
- AsyncMock(), _make_ctx(), DEFAULT_READABLE_ID, BrowseRequest(limit=10, offset=0)
+ _make_db(), _make_ctx(), DEFAULT_READABLE_ID, BrowseRequest(limit=10, offset=0)
)
assert response.total == 42
@@ -98,7 +122,7 @@ async def test_collection_not_found_raises_404(self) -> None:
svc, _, _ = _make_service()
with pytest.raises(HTTPException) as exc_info:
- await svc.browse(AsyncMock(), _make_ctx(), "nonexistent", BrowseRequest())
+ await svc.browse(_make_db(), _make_ctx(), "nonexistent", BrowseRequest())
assert exc_info.value.status_code == 404
assert "nonexistent" in exc_info.value.detail
@@ -109,7 +133,7 @@ async def test_chunk_index_anchor_added_to_filter_groups(self) -> None:
svc, vector_db, repo = _make_service()
repo.seed_readable(DEFAULT_READABLE_ID, _make_collection())
- await svc.browse(AsyncMock(), _make_ctx(), DEFAULT_READABLE_ID, BrowseRequest())
+ await svc.browse(_make_db(), _make_ctx(), DEFAULT_READABLE_ID, BrowseRequest())
filter_call = next(c for c in vector_db._calls if c[0] == "filter_search")
filter_groups = filter_call[1]
@@ -124,7 +148,7 @@ async def test_sync_ids_translated_to_filter(self) -> None:
sync_id = str(uuid4())
await svc.browse(
- AsyncMock(),
+ _make_db(),
_make_ctx(),
DEFAULT_READABLE_ID,
BrowseRequest(sync_ids=[sync_id]),
@@ -145,7 +169,7 @@ async def test_entity_types_translated_to_filter(self) -> None:
repo.seed_readable(DEFAULT_READABLE_ID, _make_collection())
await svc.browse(
- AsyncMock(),
+ _make_db(),
_make_ctx(),
DEFAULT_READABLE_ID,
BrowseRequest(entity_types=["NotionPageEntity", "SlackMessageEntity"]),
@@ -186,7 +210,7 @@ async def test_user_filter_combined_with_anchor(self) -> None:
),
]
await svc.browse(
- AsyncMock(),
+ _make_db(),
_make_ctx(),
DEFAULT_READABLE_ID,
BrowseRequest(filter=user_filter),
@@ -210,7 +234,7 @@ async def test_name_query_passed_as_name_substring(self) -> None:
repo.seed_readable(DEFAULT_READABLE_ID, _make_collection())
await svc.browse(
- AsyncMock(),
+ _make_db(),
_make_ctx(),
DEFAULT_READABLE_ID,
BrowseRequest(name_query=" Quick "),
@@ -229,7 +253,7 @@ async def test_blank_name_query_treated_as_none(self) -> None:
repo.seed_readable(DEFAULT_READABLE_ID, _make_collection())
await svc.browse(
- AsyncMock(),
+ _make_db(),
_make_ctx(),
DEFAULT_READABLE_ID,
BrowseRequest(name_query=" "),
@@ -251,3 +275,24 @@ def test_sync_ids_rejects_oversized_list(self) -> None:
with pytest.raises(ValidationError):
BrowseRequest(sync_ids=[str(uuid4()) for _ in range(101)])
+
+
+@pytest.mark.asyncio
+async def test_browse_filters_before_output_and_reports_unknown_canonical_count(monkeypatch):
+ from airweave.domains.search.visible_vector_db import UnavailableExactCount
+
+ svc, vector_db, repo = _make_service()
+ repo.seed_readable(DEFAULT_READABLE_ID, _make_collection())
+ fresh, stale = make_result("fresh"), make_result("stale")
+ vector_db.seed_filter_results([fresh, stale])
+ gate = AsyncMock(return_value=[fresh])
+ monkeypatch.setattr("airweave.domains.search.visible_vector_db.visible_results", gate)
+ monkeypatch.setattr(
+ "airweave.domains.search.visible_vector_db.VisibleVectorDB.count",
+ AsyncMock(side_effect=UnavailableExactCount("unknown")),
+ )
+ response = await svc.browse(_make_db(), _make_ctx(), DEFAULT_READABLE_ID, BrowseRequest())
+ assert response.results == [fresh]
+ assert response.total is None
+ assert response.retrieval_incomplete is True
+ assert gate.await_args.args[3] == [fresh, stale]
diff --git a/backend/airweave/domains/search/builders/collection_metadata.py b/backend/airweave/domains/search/builders/collection_metadata.py
index fc8984bb95..509daf2363 100644
--- a/backend/airweave/domains/search/builders/collection_metadata.py
+++ b/backend/airweave/domains/search/builders/collection_metadata.py
@@ -8,6 +8,7 @@
from airweave.api.context import ApiContext
from airweave.domains.collections.protocols import CollectionRepositoryProtocol
+from airweave.domains.entities.canonical.query_store import CanonicalQueryStore
from airweave.domains.entities.protocols import (
EntityCountRepositoryProtocol,
EntityDefinitionRegistryProtocol,
@@ -181,6 +182,10 @@ class CollectionMetadataBuilder:
"A modern knowledge base and team wiki for documenting and sharing "
"information with powerful search and integrations."
),
+ "wispr": (
+ "Captured meeting notes and transcripts from Wispr. Search covers the "
+ "content exposed by the connected account; capture may be partial."
+ ),
"slack": (
"A team communication platform featuring channels, direct messages, and "
"integrations for real-time collaboration."
@@ -315,12 +320,25 @@ async def build(
# 3. Build metadata for each source connection
sources: list[SourceMetadata] = []
for sc in source_connections:
+ if not sc.is_authenticated:
+ continue
# Get entity definitions this source can produce (in-memory, sync)
entity_definitions = self._entity_definition_registry.list_for_source(sc.short_name)
- # Get entity counts for this source connection's sync
+ entry = self._source_registry.get(sc.short_name)
+ canonical_types = tuple(getattr(entry.source_class_ref, "canonical_record_types", ()))
+ captured_counts = None
+ if canonical_types:
+ counts = (
+ await CanonicalQueryStore().captured_counts(db, ctx.organization.id, sc.sync_id)
+ if sc.sync_id
+ else {}
+ )
+ captured_counts = {kind: counts.get(kind, 0) for kind in canonical_types}
+
+ # Legacy counts describe indexed entity classes, not captured records.
counts_by_short_name: dict[str, int] = {}
- if sc.sync_id:
+ if sc.sync_id and not canonical_types:
entity_counts = await self._entity_count_repo.get_counts_per_sync_and_type(
db, sc.sync_id
)
@@ -330,7 +348,9 @@ async def build(
# Build entity type metadata with fields and counts
entity_types: list[EntityTypeMetadata] = []
for entity_def in entity_definitions:
- count = counts_by_short_name.get(entity_def.short_name, 0)
+ count = (
+ None if canonical_types else counts_by_short_name.get(entity_def.short_name, 0)
+ )
fields = self._extract_fields(entity_def.entity_schema)
entity_types.append(
@@ -349,7 +369,8 @@ async def build(
short_name=sc.short_name,
description=self._get_source_description(sc.short_name),
entity_types=entity_types,
- federated=entry.federated_search,
+ federated=entry.federated_search and not canonical_types,
+ captured_record_counts=captured_counts,
)
)
diff --git a/backend/airweave/domains/search/canonical_visibility.py b/backend/airweave/domains/search/canonical_visibility.py
new file mode 100644
index 0000000000..2a6aebb03a
--- /dev/null
+++ b/backend/airweave/domains/search/canonical_visibility.py
@@ -0,0 +1,115 @@
+"""Reject stale index publications before any reranking or model consumption."""
+
+from uuid import UUID
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import publications_match
+from airweave.domains.entities.canonical.source import indexed_record_types
+from airweave.domains.search.types.results import SearchResult
+from airweave.domains.sources.protocols import SourceRegistryProtocol
+from airweave.models.entity import Entity
+from airweave.models.projection_generation import ProjectionGeneration
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+
+
+async def visible_results(
+ db: AsyncSession,
+ organization_id: UUID,
+ collection: str,
+ results: list[SearchResult],
+ registry: SourceRegistryProtocol,
+) -> list[SearchResult]:
+ """Validate request scope and exact publication; never trust index ownership fields."""
+ if not results:
+ return []
+ connections = await db.scalars(
+ select(SourceConnection).where(
+ SourceConnection.organization_id == organization_id,
+ SourceConnection.is_authenticated.is_(True),
+ SourceConnection.readable_collection_id == collection,
+ )
+ )
+ scopes = {
+ str(connection.sync_id): connection
+ for connection in connections
+ if connection.sync_id is not None
+ }
+ legacy: set[int] = set()
+ candidates: dict[int, tuple[UUID, ProjectionLocator]] = {}
+ for index, result in enumerate(results):
+ metadata = result.airweave_system_metadata
+ scope = scopes.get(metadata.sync_id)
+ if scope is None or metadata.source_name != scope.short_name:
+ continue
+ canonical = bool(indexed_record_types(scope.short_name, registry))
+ try:
+ locator = ProjectionLocator.parse(metadata.original_entity_id)
+ except ValueError:
+ continue
+ if locator is None:
+ if not canonical:
+ legacy.add(index)
+ else:
+ candidates[index] = (scope.sync_id, locator)
+ if not candidates:
+ return [result for index, result in enumerate(results) if index in legacy]
+ # Bound predicate size independently from a caller's requested search page.
+ eligible: dict[tuple[UUID, UUID, int, int, UUID], ExtractionCoverage | None] = {}
+ candidate_values = list(set(candidates.values()))
+ for start in range(0, len(candidate_values), 100):
+ batch = candidate_values[start : start + 100]
+ rows = await db.execute(
+ select(
+ Entity.sync_id,
+ Entity.id,
+ Entity.record_revision,
+ Entity.indexed_pipeline_version,
+ Entity.indexed_generation,
+ ProjectionGeneration.extraction_coverage,
+ )
+ .join(Sync, Sync.id == Entity.sync_id)
+ .join(ProjectionGeneration, ProjectionGeneration.id == Entity.indexed_generation)
+ .where(
+ Entity.organization_id == organization_id,
+ Sync.organization_id == organization_id,
+ Entity.sync_id.in_({sync_id for sync_id, _ in batch}),
+ publications_match(locator for _, locator in batch),
+ )
+ )
+ for row in rows:
+ eligible[tuple(row[:-1])] = (
+ ExtractionCoverage.model_validate(row[-1]) if row[-1] is not None else None
+ )
+ return [
+ result
+ for index, result in enumerate(results)
+ if index in legacy or _eligible_part(candidates.get(index), eligible)
+ ]
+
+
+def _eligible_part(
+ candidate: tuple[UUID, ProjectionLocator] | None,
+ eligible: dict[tuple[UUID, UUID, int, int, UUID], ExtractionCoverage | None],
+) -> bool:
+ if candidate is None:
+ return False
+ sync_id, locator = candidate
+ identity = (
+ sync_id,
+ locator.record_id,
+ locator.revision,
+ locator.pipeline_version,
+ locator.generation,
+ )
+ if identity not in eligible:
+ return False
+ coverage = eligible[identity]
+ return coverage is None or any(
+ part.part_index == locator.part_index and part.outcome == "indexed"
+ for part in coverage.parts
+ )
diff --git a/backend/airweave/domains/search/classic/service.py b/backend/airweave/domains/search/classic/service.py
index b2f98db5fc..b3402d0aa4 100644
--- a/backend/airweave/domains/search/classic/service.py
+++ b/backend/airweave/domains/search/classic/service.py
@@ -172,7 +172,9 @@ async def _execute(
documents=documents,
top_n=request.limit,
)
- results = SearchResults(results=[results.results[r.index] for r in reranked])
+ results = results.model_copy(
+ update={"results": [results.results[r.index] for r in reranked]}
+ )
# 7. Emit event + return
duration_ms = int((time.monotonic() - start_time) * 1000)
diff --git a/backend/airweave/domains/search/executor.py b/backend/airweave/domains/search/executor.py
index 8869e81f11..f94d05847a 100644
--- a/backend/airweave/domains/search/executor.py
+++ b/backend/airweave/domains/search/executor.py
@@ -19,6 +19,7 @@
from airweave.domains.embedders.protocols import DenseEmbedderProtocol, SparseEmbedderProtocol
from airweave.domains.search.adapters.vector_db.protocol import VectorDBProtocol
from airweave.domains.search.builders.search_plan import SearchPlanBuilder
+from airweave.domains.search.canonical_visibility import visible_results
from airweave.domains.search.exceptions import FederatedSearchError
from airweave.domains.search.protocols import SearchPlanExecutorProtocol
from airweave.domains.search.types import (
@@ -28,6 +29,7 @@
SearchPlan,
SearchResults,
)
+from airweave.domains.search.types.embeddings import PreparedQueryEmbeddings
from airweave.domains.search.types.filters import FilterableField, FilterCondition, FilterOperator
from airweave.domains.search.types.results import (
SearchAccessControl,
@@ -88,6 +90,8 @@ async def execute(
ctx: ApiContext,
collection_readable_id: str,
user_principal: Optional[str] = None,
+ indexed_only: bool = False,
+ prepared_query: PreparedQueryEmbeddings | None = None,
) -> SearchResults:
"""Execute the full search pipeline including federated sources."""
# 0. Resolve access control principals
@@ -99,7 +103,11 @@ async def execute(
complete_plan = SearchPlanBuilder.build(plan, user_filter)
# 2. Discover federated sources for this collection
- federated_sources = await self._discover_federated_sources(db, ctx, collection_readable_id)
+ federated_sources = (
+ []
+ if indexed_only
+ else await self._discover_federated_sources(db, ctx, collection_readable_id)
+ )
# 3. Adjust limit/offset for RRF pagination (if federated sources exist)
original_limit = complete_plan.limit
@@ -116,7 +124,9 @@ async def execute(
fetch_limit = original_offset + original_limit
vector_task = asyncio.create_task(
- self._execute_vector_search(complete_plan, collection_id, acl_principals)
+ self._execute_vector_search(
+ complete_plan, collection_id, acl_principals, prepared_query
+ )
)
fed_task = None
@@ -130,12 +140,30 @@ async def execute(
)
)
- vector_results = await vector_task
+ engine_results = await vector_task
+ retrieved = engine_results.results
+ coverage = {
+ "engine_partial": engine_results.engine_partial,
+ "engine_coverage_percent": engine_results.engine_coverage_percent,
+ }
+ vector_results = await visible_results(
+ db,
+ ctx.organization.id,
+ collection_readable_id,
+ retrieved,
+ self._source_registry,
+ )
+ excluded = len(retrieved) - len(vector_results)
fed_results = await fed_task if fed_task else []
# 5. If no federated sources, vector DB already has correct limit/offset
if not federated_sources:
- return SearchResults(results=vector_results)
+ return SearchResults(
+ results=vector_results,
+ retrieval_incomplete=excluded > 0 or engine_results.engine_partial,
+ **coverage,
+ excluded_candidates=excluded,
+ )
# 6. We over-fetched from vector DB (limit=offset+limit, offset=0) for RRF.
# Filter federated results in-memory and merge, or slice vector-only.
@@ -147,24 +175,23 @@ async def execute(
if fed_filtered:
merged = self._merge_with_rrf(vector_results, fed_filtered)
- return SearchResults(results=merged[original_offset : original_offset + original_limit])
+ return SearchResults(
+ results=merged[original_offset : original_offset + original_limit],
+ retrieval_incomplete=excluded > 0 or engine_results.engine_partial,
+ **coverage,
+ excluded_candidates=excluded,
+ )
# All federated results filtered out — slice vector results to original window
return SearchResults(
- results=vector_results[original_offset : original_offset + original_limit]
+ results=vector_results[original_offset : original_offset + original_limit],
+ retrieval_incomplete=excluded > 0 or engine_results.engine_partial,
+ **coverage,
+ excluded_candidates=excluded,
)
- async def _execute_vector_search(
- self,
- plan: SearchPlan,
- collection_id: str,
- acl_principals: Optional[list[str]] = None,
- ) -> list[SearchResult]:
- """Embed, compile, and execute vector DB search.
-
- Adapter exceptions (EmbedderError, VectorDBError) propagate directly
- to the caller — no wrapping needed since adapters own their error types.
- """
+ async def prepare_query(self, plan: SearchPlan) -> PreparedQueryEmbeddings:
+ """Prepare embeddings for this executor only; no source data or authorization cached."""
dense_embeddings = None
sparse_embedding = None
@@ -173,7 +200,7 @@ async def _execute_vector_search(
RetrievalStrategy.HYBRID,
):
texts = [plan.query.primary] + list(plan.query.variations)
- dense_embeddings = await self._dense_embedder.embed_many(texts)
+ dense_embeddings = await self._dense_embedder.embed_many(texts, purpose="query")
if plan.retrieval_strategy in (
RetrievalStrategy.KEYWORD,
@@ -186,13 +213,37 @@ async def _execute_vector_search(
sparse_embedding=sparse_embedding,
)
+ prepared = PreparedQueryEmbeddings(
+ primary=plan.query.primary,
+ variations=tuple(plan.query.variations),
+ strategy=plan.retrieval_strategy,
+ embeddings=embeddings,
+ )
+ prepared._owner = self
+ return prepared
+
+ async def _execute_vector_search(
+ self,
+ plan: SearchPlan,
+ collection_id: str,
+ acl_principals: Optional[list[str]] = None,
+ prepared_query: PreparedQueryEmbeddings | None = None,
+ ) -> SearchResults:
+ """Embed, compile, and execute vector DB search.
+
+ Adapter exceptions (EmbedderError, VectorDBError) propagate directly
+ to the caller — no wrapping needed since adapters own their error types.
+ """
+ prepared = prepared_query if prepared_query is not None else await self.prepare_query(plan)
+ embeddings = prepared.require_match(plan, self)
+
compiled_query = await self._vector_db.compile_query(
plan=plan,
embeddings=embeddings,
collection_id=collection_id,
acl_principals=acl_principals,
)
- return (await self._vector_db.execute_query(compiled_query)).results
+ return await self._vector_db.execute_query(compiled_query)
# ------------------------------------------------------------------
# Access control resolution
@@ -279,7 +330,9 @@ async def _discover_federated_sources(
federated_sources: list[BaseSource] = []
for sc in source_connections:
entry = self._source_registry.get(sc.short_name)
- if not entry.federated_search:
+ if not entry.federated_search or getattr(
+ entry.source_class_ref, "canonical_record_types", ()
+ ):
continue
try:
diff --git a/backend/airweave/domains/search/fakes/executor.py b/backend/airweave/domains/search/fakes/executor.py
index c49226d85e..a28f3d2daf 100644
--- a/backend/airweave/domains/search/fakes/executor.py
+++ b/backend/airweave/domains/search/fakes/executor.py
@@ -5,7 +5,8 @@
from typing import Any
from airweave.domains.search.protocols import SearchPlanExecutorProtocol
-from airweave.domains.search.types import FilterGroup, SearchPlan, SearchResults
+from airweave.domains.search.types import FilterGroup, QueryEmbeddings, SearchPlan, SearchResults
+from airweave.domains.search.types.embeddings import PreparedQueryEmbeddings
class FakeSearchPlanExecutor(SearchPlanExecutorProtocol):
@@ -32,6 +33,17 @@ def seed_result_sequence(self, sequence: list[SearchResults | Exception]) -> Non
"""Seed a sequence of results/errors for successive execute calls."""
self._sequence = list(sequence)
+ async def prepare_query(self, plan: SearchPlan) -> PreparedQueryEmbeddings:
+ """Mirror the internal preparation contract without inference."""
+ result = PreparedQueryEmbeddings(
+ primary=plan.query.primary,
+ variations=tuple(plan.query.variations),
+ strategy=plan.retrieval_strategy,
+ embeddings=QueryEmbeddings(),
+ )
+ result._owner = self
+ return result
+
async def execute(
self,
plan: SearchPlan,
@@ -41,6 +53,8 @@ async def execute(
ctx: Any = None,
collection_readable_id: str = "",
user_principal: str | None = None,
+ indexed_only: bool = False,
+ prepared_query: PreparedQueryEmbeddings | None = None,
) -> SearchResults:
"""Record the call and return seeded result, or raise seeded error."""
self._calls.append(("execute", plan, user_filter, collection_id))
diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py
new file mode 100644
index 0000000000..64944ae396
--- /dev/null
+++ b/backend/airweave/domains/search/owned.py
@@ -0,0 +1,740 @@
+"""Original-record indexed retrieval using the existing executor and SQL authority."""
+
+import asyncio
+import math
+import re
+from collections import defaultdict
+from uuid import UUID
+
+from fastapi import HTTPException
+from pydantic import AwareDatetime, BaseModel, ConfigDict
+from sqlalchemy import and_, case, func, or_, select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from airweave.api.context import ApiContext
+from airweave.core.protocols.reranker import RerankerProtocol, RerankerResult
+from airweave.core.protocols.tokenizer import TokenizerProtocol
+from airweave.domains.entities.canonical.coverage import capture_coverage
+from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage
+from airweave.domains.entities.canonical.projection_models import ProjectionLocator
+from airweave.domains.entities.canonical.projection_store import publications_match
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.search_metadata import (
+ SEARCH_METADATA_PIPELINE_VERSION,
+ epoch_microseconds,
+)
+from airweave.domains.entities.canonical.source import indexed_record_types
+from airweave.domains.entities.canonical.store import content_is_available
+from airweave.domains.search.owned_models import (
+ OwnedRanking,
+ OwnedSearchCoverage,
+ OwnedSearchGroup,
+ OwnedSearchHit,
+ OwnedSearchMatch,
+ OwnedSearchRequest,
+ OwnedSearchResponse,
+)
+from airweave.domains.search.protocols import SearchPlanExecutorProtocol
+from airweave.domains.search.types import FilterCondition, FilterGroup, SearchPlan, SearchQuery
+from airweave.domains.search.types.results import SearchResult, SearchResults
+from airweave.domains.sources.protocols import SourceRegistryProtocol
+from airweave.models.collection import Collection
+from airweave.models.entity import Entity
+from airweave.models.projection_generation import ProjectionGeneration
+from airweave.models.source_connection import SourceConnection
+from airweave.models.sync import Sync
+
+
+class _SourceScope(BaseModel):
+ """Detached authorization/index snapshot; revalidated before returning any hit."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ id: UUID
+ sync_id: UUID
+ short_name: str
+ readable_collection_id: str
+ collection_id: UUID
+ index_pipeline_version: int
+
+
+class _EnrichmentRecord(BaseModel):
+ """Detached search-card fields; original payloads and blob manifests stay in the store."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ id: UUID
+ sync_id: UUID
+ record_revision: int
+ indexed_generation: UUID
+ indexed_pipeline_version: int
+ entity_definition_short_name: str
+ native_id: str
+ container_id: str | None
+ parent_record_type: str | None
+ parent_native_id: str | None
+ parent_container_id: str | None
+ observed_at: AwareDatetime
+ source_created_at: AwareDatetime | None
+ source_updated_at: AwareDatetime | None
+ completeness: str
+ email_thread_id: str | None
+
+
+class OwnedSearchService:
+ """Canonical retrieval and optional shared ranking, without source provider calls."""
+
+ def __init__(
+ self,
+ executor: SearchPlanExecutorProtocol,
+ registry: SourceRegistryProtocol,
+ *,
+ reranker: RerankerProtocol | None = None,
+ tokenizer: TokenizerProtocol | None = None,
+ ):
+ """Reuse the executor and registry without another retrieval stack."""
+ self._executor = executor
+ self._registry = registry
+ self._reranker = reranker
+ self._tokenizer = tokenizer
+
+ async def search(
+ self,
+ sessions: async_sessionmaker[AsyncSession],
+ ctx: ApiContext,
+ request: OwnedSearchRequest,
+ ) -> OwnedSearchResponse:
+ """Resolve exact authorized scopes before any embedding/index request."""
+ async with sessions() as db:
+ scopes, groups = await self._resolve_scopes(db, ctx, request)
+ if self._filtered(request) and any(
+ scope.index_pipeline_version < SEARCH_METADATA_PIPELINE_VERSION
+ for scope in scopes.values()
+ ):
+ raise HTTPException(
+ 409,
+ {
+ "code": "reindex_required",
+ "message": "Selected sources need canonical search metadata re-projection",
+ },
+ )
+ scope_snapshot = self._scope_identity(scopes)
+ hits, scores, matched_text, exclusions, postfiltered = {}, {}, {}, 0, 0
+ engine_partial, full = False, False
+ plan = SearchPlan(
+ query=SearchQuery(primary=request.query),
+ limit=200,
+ offset=0,
+ retrieval_strategy=request.mode,
+ )
+ prepared_query = await self._executor.prepare_query(plan) if groups else None
+ for (collection_id, readable_id), sync_ids in groups.items():
+ # A new session has no checked-out connection before indexed retrieval:
+ # this path has neither principal ACL discovery nor provider federation.
+ async with sessions() as db:
+ results = await self._executor.execute(
+ plan=plan,
+ prepared_query=prepared_query,
+ user_filter=[FilterGroup(conditions=self._prefilters(request, sync_ids))],
+ collection_id=str(collection_id),
+ db=db,
+ ctx=ctx,
+ collection_readable_id=readable_id,
+ indexed_only=True,
+ )
+ group_hits, group_scores, group_text, rejected, filtered = await self._enrich(
+ db, ctx, request, sync_ids, scopes, results
+ )
+ engine_partial |= results.engine_partial
+ exclusions += results.excluded_candidates
+ full |= len(results.results) + results.excluded_candidates >= 200
+ hits.update(group_hits)
+ scores.update(group_scores)
+ matched_text.update(group_text)
+ exclusions += rejected
+ postfiltered += filtered
+ # Revalidate the whole union before any remote text disclosure: earlier
+ # collections may have changed while later collections were retrieved.
+ async with sessions() as db:
+ fresh_scopes, _ = await self._resolve_scopes(db, ctx, request)
+ if self._scope_identity(fresh_scopes) != scope_snapshot:
+ raise HTTPException(404, "Requested indexed sources changed during retrieval")
+ eligible = await self._final_publications(db, ctx, request, scores, scope_snapshot)
+ exclusions += len(hits.keys() - eligible)
+ candidates = sorted(hits.keys() & eligible, key=lambda key: (-scores[key][0], str(key)))
+ ranked, ranking = await self._rank(request.query, candidates, hits, matched_text)
+ # No SQL connection is held across the reranker await. No network I/O
+ # follows this final authorization/publication read boundary.
+ async with sessions() as db:
+ sources = await self._coverage(db, ctx, request)
+ fresh_scopes, _ = await self._resolve_scopes(db, ctx, request)
+ if self._scope_identity(fresh_scopes) != scope_snapshot:
+ raise HTTPException(404, "Requested indexed sources changed during retrieval")
+ final = await self._final_publications(db, ctx, request, scores, scope_snapshot)
+ exclusions += len(set(candidates) - final)
+ ranked = [key for key in ranked if key in final]
+ grouped = self._group_hits([hits[key] for key in ranked])
+ return OwnedSearchResponse(
+ items=tuple(grouped[: request.limit]),
+ ranking=ranking,
+ sources=sources,
+ candidate_window_full=full,
+ engine_partial=engine_partial,
+ excluded_candidates=exclusions,
+ postfilter_excluded=postfiltered,
+ retrieval_incomplete=ranking.shortlist_truncated
+ or engine_partial
+ or full
+ or exclusions > 0
+ or postfiltered > 0
+ or len(grouped) > request.limit
+ or any(
+ row.pending_records
+ or row.partially_indexed_records
+ or row.extraction_unavailable_records
+ or row.extraction_unknown_records
+ for row in sources
+ ),
+ )
+
+ @staticmethod
+ def _conversation(row: _EnrichmentRecord, hit: OwnedSearchHit) -> OwnedSearchGroup | None:
+ """Derive membership only from admitted canonical parent or validated Gmail identity."""
+ if hit.provider == "gmail" and hit.email_thread_id is not None:
+ return OwnedSearchGroup(
+ kind="email_thread", native_id=hit.email_thread_id, matched_records=1
+ )
+ if hit.provider != "almanac":
+ return None
+ if (
+ row.entity_definition_short_name == "session"
+ and row.container_id is None
+ and row.parent_record_type is None
+ and row.parent_native_id is None
+ and row.parent_container_id is None
+ ):
+ return OwnedSearchGroup(kind="session", native_id=row.native_id, matched_records=1)
+ if (
+ row.entity_definition_short_name == "message"
+ and row.parent_record_type == "session"
+ and row.parent_native_id
+ and row.parent_container_id is None
+ and row.container_id == row.parent_native_id
+ ):
+ return OwnedSearchGroup(
+ kind="session", native_id=row.parent_native_id, matched_records=1
+ )
+ return None
+
+ @staticmethod
+ def _group_hits(ranked: list[OwnedSearchHit]) -> list[OwnedSearchHit]:
+ """Group final eligible matches in rank order, preserving each original."""
+ output: list[OwnedSearchHit] = []
+ positions: dict[tuple[UUID, str, str, str], int] = {}
+ for hit in ranked:
+ group = hit.group
+ if group is None:
+ output.append(hit)
+ continue
+ key = (hit.sync_id, hit.provider, group.kind, group.native_id)
+ if key not in positions:
+ positions[key] = len(output)
+ output.append(
+ hit.model_copy(
+ update={
+ "group": group.model_copy(
+ update={"matched_records": 1, "additional_matches": ()}
+ )
+ }
+ )
+ )
+ continue
+ index = positions[key]
+ representative = output[index]
+ previous = representative.group
+ assert previous is not None
+ additional = previous.additional_matches
+ if len(additional) < 3:
+ exact = OwnedSearchMatch.model_validate(
+ hit.model_dump(exclude={"group"}, round_trip=True)
+ )
+ additional = (
+ *additional,
+ exact.model_copy(update={"excerpts": exact.excerpts[:1]}),
+ )
+ output[index] = representative.model_copy(
+ update={
+ "group": previous.model_copy(
+ update={
+ "matched_records": previous.matched_records + 1,
+ "additional_matches": additional,
+ }
+ )
+ }
+ )
+ return output
+
+ async def _rank(
+ self,
+ query: str,
+ candidates: list[UUID],
+ hits: dict[UUID, OwnedSearchHit],
+ matched_text: dict[UUID, str],
+ ) -> tuple[list[UUID], OwnedRanking]:
+ """One model call over a bounded mixed shortlist; failures keep retrieval order."""
+ shortlist = candidates[:200]
+ ranking = OwnedRanking(
+ candidates_considered=len(candidates),
+ shortlisted_candidates=len(shortlist),
+ shortlist_truncated=len(candidates) > len(shortlist),
+ )
+ if not shortlist:
+ return [], ranking.model_copy(update={"fallback_reason": None})
+ if self._reranker is None or self._tokenizer is None:
+ return shortlist, ranking
+ documents, truncated = [], 0
+ for key in shortlist:
+ title = hits[key].title[:256]
+ text = f"{title}\n\n{matched_text[key]}"
+ document = self._bounded_document(text)
+ documents.append(document)
+ truncated += int(document != text or title != hits[key].title)
+ ranking = ranking.model_copy(update={"input_truncated_documents": truncated})
+ try:
+ async with asyncio.timeout(10):
+ results = await self._reranker.rerank(query, documents, top_n=len(documents))
+ except TimeoutError:
+ return shortlist, ranking.model_copy(update={"fallback_reason": "timeout"})
+ except Exception:
+ # Never log provider errors: SDK messages can contain submitted text.
+ return shortlist, ranking.model_copy(update={"fallback_reason": "provider_error"})
+ try:
+ invalid = (
+ not isinstance(results, list)
+ or len(results) != len(shortlist)
+ or any(
+ not isinstance(result, RerankerResult)
+ or type(result.index) is not int
+ or not 0 <= result.index < len(shortlist)
+ or type(result.relevance_score) not in (int, float)
+ or not math.isfinite(result.relevance_score)
+ for result in results
+ )
+ or {result.index for result in results} != set(range(len(shortlist)))
+ )
+ except (TypeError, ValueError, OverflowError):
+ invalid = True
+ if invalid:
+ return shortlist, ranking.model_copy(update={"fallback_reason": "invalid_output"})
+ ordered = sorted(results, key=lambda result: (-result.relevance_score, result.index))
+ return [shortlist[result.index] for result in ordered], ranking.model_copy(
+ update={
+ "method": "shared_rerank",
+ "fallback_reason": None,
+ "candidates_reranked": len(shortlist),
+ }
+ )
+
+ def _bounded_document(self, text: str) -> str:
+ """Bound local tokenizer counts; provider tokenization and billing may differ."""
+ if self._tokenizer.count_tokens(text) <= 2048:
+ return text
+ accepted, rejected = 0, len(text)
+ while rejected - accepted > 1:
+ midpoint = (accepted + rejected) // 2
+ if self._tokenizer.count_tokens(text[:midpoint]) <= 2048:
+ accepted = midpoint
+ else:
+ rejected = midpoint
+ # Token count need not be monotonic in characters. The accepted prefix
+ # was measured explicitly, so this remains a safe bound, not a max-length claim.
+ return text[:accepted]
+
+ @staticmethod
+ def _filtered(request: OwnedSearchRequest) -> bool:
+ return bool(
+ request.record_types
+ or request.created_after
+ or request.created_before
+ or request.updated_after
+ or request.updated_before
+ )
+
+ @staticmethod
+ def _prefilters(request: OwnedSearchRequest, sync_ids: list[UUID]) -> list[FilterCondition]:
+ conditions = [
+ FilterCondition(
+ field="airweave_system_metadata.sync_id",
+ operator="in",
+ value=[str(item) for item in sync_ids],
+ )
+ ]
+ if request.record_types:
+ conditions.append(
+ FilterCondition(
+ field="airweave_system_metadata.canonical_record_type",
+ operator="in",
+ value=list(request.record_types),
+ )
+ )
+ for name, after, before in (
+ ("created", request.created_after, request.created_before),
+ ("updated", request.updated_after, request.updated_before),
+ ):
+ if after is not None or before is not None:
+ conditions.append(
+ FilterCondition(
+ field=f"airweave_system_metadata.source_{name}_known",
+ operator="equals",
+ value=1,
+ )
+ )
+ for boundary, operator in ((after, "greater_than_or_equal"), (before, "less_than")):
+ if boundary is not None:
+ conditions.append(
+ FilterCondition(
+ field=f"airweave_system_metadata.source_{name}_us",
+ operator=operator,
+ value=epoch_microseconds(boundary),
+ )
+ )
+ return conditions
+
+ async def _resolve_scopes(
+ self, db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest
+ ) -> tuple[dict[UUID, _SourceScope], dict[tuple[UUID, str], list[UUID]]]:
+ rows = (
+ (
+ await db.execute(
+ select(
+ SourceConnection.id,
+ SourceConnection.sync_id,
+ SourceConnection.short_name,
+ SourceConnection.readable_collection_id,
+ Collection.id.label("collection_id"),
+ Sync.index_pipeline_version,
+ )
+ .join(
+ Collection,
+ and_(
+ Collection.readable_id == SourceConnection.readable_collection_id,
+ Collection.organization_id == ctx.organization.id,
+ ),
+ )
+ .join(
+ Sync,
+ and_(
+ Sync.id == SourceConnection.sync_id,
+ Sync.organization_id == ctx.organization.id,
+ ),
+ )
+ .where(
+ SourceConnection.organization_id == ctx.organization.id,
+ SourceConnection.is_authenticated.is_(True),
+ SourceConnection.sync_id.in_(request.sync_ids),
+ )
+ )
+ )
+ .mappings()
+ .all()
+ )
+ snapshots = [_SourceScope.model_validate(row) for row in rows]
+ if len(snapshots) != len(request.sync_ids) or {row.sync_id for row in snapshots} != set(
+ request.sync_ids
+ ):
+ raise HTTPException(404, "Requested indexed sources are unavailable")
+ scopes = {row.sync_id: row for row in snapshots}
+ allowed_types = set()
+ groups = defaultdict(list)
+ for connection in snapshots:
+ types = indexed_record_types(connection.short_name, self._registry)
+ if not types:
+ raise HTTPException(422, "Source does not support owned indexed records")
+ allowed_types.update(types)
+ groups[(connection.collection_id, connection.readable_collection_id)].append(
+ connection.sync_id
+ )
+ if set(request.record_types) - allowed_types:
+ raise HTTPException(422, "Record type is unsupported by selected sources")
+ return scopes, groups
+
+ async def _enrich(
+ self,
+ db: AsyncSession,
+ ctx: ApiContext,
+ request: OwnedSearchRequest,
+ sync_ids: list[UUID],
+ scopes: dict[UUID, _SourceScope],
+ results: SearchResults,
+ ) -> tuple[
+ dict[UUID, OwnedSearchHit],
+ dict[UUID, tuple[float, ProjectionLocator]],
+ dict[UUID, str],
+ int,
+ int,
+ ]:
+ locators = [value for result in results.results if (value := self._locator(result))]
+ if not locators:
+ return {}, {}, {}, len(results.results), 0
+ hits, scores, matched_text, exclusions, postfiltered = {}, {}, {}, 0, 0
+ # Recheck current publication while obtaining native identity. Never enrich
+ # from an unvalidated cached hit after a concurrent capture/permission change.
+ rows = await db.execute(
+ select(
+ Entity.id,
+ Entity.sync_id,
+ Entity.record_revision,
+ Entity.indexed_generation,
+ Entity.indexed_pipeline_version,
+ Entity.entity_definition_short_name,
+ Entity.native_id,
+ Entity.container_id,
+ Entity.parent_record_type,
+ Entity.parent_native_id,
+ Entity.parent_container_id,
+ Entity.observed_at,
+ Entity.source_created_at,
+ Entity.source_updated_at,
+ Entity.completeness,
+ case(
+ (
+ func.jsonb_typeof(Entity.source_payload["threadId"]) == "string",
+ Entity.source_payload["threadId"].astext,
+ ),
+ else_=None,
+ ).label("email_thread_id"),
+ )
+ .join(Sync, Sync.id == Entity.sync_id)
+ .where(
+ Entity.organization_id == ctx.organization.id,
+ Sync.organization_id == ctx.organization.id,
+ Entity.sync_id.in_(sync_ids),
+ publications_match(locators),
+ )
+ )
+ records = [_EnrichmentRecord.model_validate(row) for row in rows.mappings()]
+ by_id = {record.id: record for record in records}
+ extraction_rows = await db.execute(
+ select(ProjectionGeneration.id, ProjectionGeneration.extraction_coverage).where(
+ ProjectionGeneration.id.in_([r.indexed_generation for r in records])
+ )
+ )
+ extraction = {
+ generation: ExtractionCoverage.model_validate(raw) if raw is not None else None
+ for generation, raw in extraction_rows
+ }
+ for rank, result in enumerate(results.results, 1):
+ locator = self._locator(result)
+ row = by_id.get(locator.record_id) if locator else None
+ if (
+ row is None
+ or row.record_revision != locator.revision
+ or row.indexed_generation != locator.generation
+ or row.indexed_pipeline_version != locator.pipeline_version
+ or str(row.sync_id) != result.airweave_system_metadata.sync_id
+ or scopes[row.sync_id].short_name != result.airweave_system_metadata.source_name
+ ):
+ exclusions += 1
+ continue
+ coverage = extraction.get(row.indexed_generation)
+ if coverage is not None and not any(
+ part.part_index == locator.part_index and part.outcome == "indexed"
+ for part in coverage.parts
+ ):
+ exclusions += 1
+ continue
+ if not self._matches(row, request):
+ postfiltered += 1
+ continue
+ if row.id not in hits:
+ hits[row.id] = OwnedSearchHit(
+ record_id=row.id,
+ revision=row.record_revision,
+ sync_id=row.sync_id,
+ source_connection_id=scopes[row.sync_id].id,
+ provider=scopes[row.sync_id].short_name,
+ identity=RecordIdentity(
+ record_type=row.entity_definition_short_name,
+ native_id=row.native_id,
+ container_id=row.container_id,
+ ),
+ title=result.name,
+ excerpts=(),
+ observed_at=row.observed_at,
+ source_created_at=row.source_created_at,
+ source_updated_at=row.source_updated_at,
+ completeness=row.completeness,
+ extraction=extraction.get(row.indexed_generation),
+ email_thread_id=(
+ thread_id
+ if scopes[row.sync_id].short_name == "gmail"
+ and row.entity_definition_short_name == "message"
+ and (thread_id := row.email_thread_id) is not None
+ and re.fullmatch(r"[A-Za-z0-9_-]{1,512}", thread_id)
+ else None
+ ),
+ )
+ hits[row.id].group = self._conversation(row, hits[row.id])
+ scores[row.id] = (1 / (60 + rank), locator)
+ matched_text[row.id] = result.textual_representation
+ # Dynamic summaries are keyword fragments, not semantic explanations.
+ # No-mark and older-schema responses keep the existing chunk fallback.
+ snippet = result.query_snippet
+ if (
+ request.mode.value != "semantic"
+ and snippet
+ and "" in snippet
+ and "" in snippet
+ ):
+ # Remove only Vespa's presentation delimiters. All other markup is
+ # ordinary untrusted text, never parsed/rendered as HTML here.
+ excerpt = (
+ snippet.replace("", "")
+ .replace("", "")
+ .replace("", " … ")
+ .strip()
+ or result.textual_representation
+ )[:2000]
+ else:
+ excerpt = result.textual_representation[:2000]
+ current = hits[row.id]
+ if excerpt and excerpt not in current.excerpts and len(current.excerpts) < 3:
+ hits[row.id] = current.model_copy(update={"excerpts": (*current.excerpts, excerpt)})
+ return hits, scores, matched_text, exclusions, postfiltered
+
+ @staticmethod
+ def _scope_identity(scopes: dict[UUID, _SourceScope]) -> set[tuple]:
+ return {
+ (
+ sync,
+ row.id,
+ row.short_name,
+ row.readable_collection_id,
+ row.collection_id,
+ row.index_pipeline_version,
+ )
+ for sync, row in scopes.items()
+ }
+
+ @staticmethod
+ async def _final_publications(
+ db: AsyncSession,
+ ctx: ApiContext,
+ request: OwnedSearchRequest,
+ scores: dict[UUID, tuple[float, ProjectionLocator]],
+ scope_snapshot: set[tuple],
+ ) -> set[UUID]:
+ # Later collection retrieval may outlive an edit/revocation of earlier hits.
+ # Recheck the exact publication, including parent visibility, at the final
+ # SQL read boundary. No network I/O occurs after this authorization check.
+ if not scores:
+ return set()
+ # One statement snapshots all candidates, including current authentication.
+ # The 20*200 bound keeps these tuple parameters below PostgreSQL's limit.
+ return set(
+ await db.scalars(
+ select(Entity.id)
+ .join(Sync, Sync.id == Entity.sync_id)
+ .join(
+ SourceConnection,
+ and_(
+ SourceConnection.sync_id == Sync.id,
+ SourceConnection.organization_id == ctx.organization.id,
+ SourceConnection.is_authenticated.is_(True),
+ ),
+ )
+ .join(
+ Collection,
+ and_(
+ Collection.readable_id == SourceConnection.readable_collection_id,
+ Collection.organization_id == ctx.organization.id,
+ ),
+ )
+ .where(
+ Entity.organization_id == ctx.organization.id,
+ Sync.organization_id == ctx.organization.id,
+ Entity.sync_id.in_(request.sync_ids),
+ tuple_(
+ SourceConnection.sync_id,
+ SourceConnection.id,
+ SourceConnection.short_name,
+ SourceConnection.readable_collection_id,
+ Collection.id,
+ Sync.index_pipeline_version,
+ ).in_(list(scope_snapshot)),
+ publications_match(locator for _, locator in scores.values()),
+ )
+ )
+ )
+
+ @staticmethod
+ def _locator(result: SearchResult) -> ProjectionLocator | None:
+ try:
+ return ProjectionLocator.parse(result.airweave_system_metadata.original_entity_id)
+ except ValueError:
+ return None
+
+ @staticmethod
+ def _matches(row: _EnrichmentRecord, request: OwnedSearchRequest) -> bool:
+ if request.record_types and row.entity_definition_short_name not in request.record_types:
+ return False
+ for value, after, before in (
+ (row.source_created_at, request.created_after, request.created_before),
+ (row.source_updated_at, request.updated_after, request.updated_before),
+ ):
+ if (after is not None and (value is None or value < after)) or (
+ before is not None and (value is None or value >= before)
+ ):
+ return False
+ return True
+
+ @staticmethod
+ async def _coverage(
+ db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest
+ ) -> tuple[OwnedSearchCoverage, ...]:
+ pending = or_(
+ Entity.indexed_revision.is_distinct_from(Entity.record_revision),
+ Entity.indexed_pipeline_version.is_distinct_from(Sync.index_pipeline_version),
+ Entity.indexed_generation.is_(None),
+ )
+ has_indexed = ProjectionGeneration.extraction_coverage.contains(
+ {"parts": [{"outcome": "indexed"}]}
+ )
+ has_omitted = ProjectionGeneration.extraction_coverage.contains(
+ {"parts": [{"outcome": "unsupported"}]}
+ ) | ProjectionGeneration.extraction_coverage.contains(
+ {"parts": [{"outcome": "unavailable_original"}]}
+ )
+ rows = await db.execute(
+ select(
+ Entity.sync_id,
+ func.count(),
+ func.count().filter(pending),
+ func.count().filter(~pending & has_indexed & has_omitted),
+ func.count().filter(~pending & ~has_indexed & has_omitted),
+ func.count().filter(~pending & ProjectionGeneration.extraction_coverage.is_(None)),
+ )
+ .join(Sync, Sync.id == Entity.sync_id)
+ .outerjoin(ProjectionGeneration, ProjectionGeneration.id == Entity.indexed_generation)
+ .where(
+ Entity.organization_id == ctx.organization.id,
+ Sync.organization_id == ctx.organization.id,
+ Entity.sync_id.in_(request.sync_ids),
+ Entity.record_revision > 0,
+ Entity.deleted_at.is_(None),
+ content_is_available(),
+ )
+ .group_by(Entity.sync_id)
+ )
+ counts = {row[0]: tuple(row[1:]) for row in rows}
+ captures = await capture_coverage(db, ctx.organization.id, tuple(request.sync_ids))
+ return tuple(
+ OwnedSearchCoverage(
+ sync_id=sync,
+ capture=captures.get(sync),
+ active_records=counts.get(sync, (0, 0))[0],
+ pending_records=counts.get(sync, (0, 0))[1],
+ partially_indexed_records=counts.get(sync, (0, 0, 0, 0, 0))[2],
+ extraction_unavailable_records=counts.get(sync, (0, 0, 0, 0, 0))[3],
+ extraction_unknown_records=counts.get(sync, (0, 0, 0, 0, 0))[4],
+ )
+ for sync in request.sync_ids
+ )
diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py
new file mode 100644
index 0000000000..5800a99c4f
--- /dev/null
+++ b/backend/airweave/domains/search/owned_models.py
@@ -0,0 +1,123 @@
+"""Bounded original-record retrieval; unsupported controls fail validation."""
+
+from typing import Literal
+from uuid import UUID
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator
+
+from airweave.domains.entities.canonical.coverage_models import CaptureCoverage
+from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.search.retrieval_strategy import RetrievalStrategy
+
+
+class OwnedSearchRequest(BaseModel):
+ """Source IDs are organization-scoped; the product separately owns account authorization."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ query: str = Field(min_length=1, max_length=4000)
+ sync_ids: tuple[UUID, ...] = Field(min_length=1, max_length=20)
+ mode: RetrievalStrategy = RetrievalStrategy.HYBRID
+ record_types: tuple[str, ...] = Field(default=(), max_length=20)
+ created_after: AwareDatetime | None = None
+ created_before: AwareDatetime | None = None
+ updated_after: AwareDatetime | None = None
+ updated_before: AwareDatetime | None = None
+ limit: int = Field(default=20, ge=1, le=200)
+
+ @model_validator(mode="after")
+ def validate_scope(self):
+ """Reject ambiguous requests rather than silently correcting them."""
+ if not self.query.strip() or len(set(self.sync_ids)) != len(self.sync_ids):
+ raise ValueError("Query must be nonblank and source IDs unique")
+ if any(not item.strip() for item in self.record_types):
+ raise ValueError("Record types must be nonblank")
+ for after, before in (
+ (self.created_after, self.created_before),
+ (self.updated_after, self.updated_before),
+ ):
+ if after is not None and before is not None and after >= before:
+ raise ValueError("Time interval must have after before before")
+ return self
+
+
+class OwnedSearchMatch(BaseModel):
+ """Canonical identity plus excerpts; no native payload or generation IDs."""
+
+ record_id: UUID
+ revision: int
+ sync_id: UUID
+ source_connection_id: UUID
+ provider: str
+ identity: RecordIdentity
+ title: str
+ excerpts: tuple[str, ...]
+ email_thread_id: str | None = Field(default=None, max_length=512, pattern=r"^[A-Za-z0-9_-]+$")
+ observed_at: AwareDatetime
+ source_created_at: AwareDatetime | None
+ source_updated_at: AwareDatetime | None
+ completeness: Literal["complete", "partial", "metadata_only"]
+ extraction: ExtractionCoverage | None = None
+
+
+class OwnedSearchGroup(BaseModel):
+ """Observed eligible shortlist members, never a complete conversation count."""
+
+ kind: Literal["session", "email_thread"]
+ native_id: str
+ matched_records: int = Field(ge=1)
+ additional_matches: tuple[OwnedSearchMatch, ...] = Field(default=(), max_length=3)
+
+
+class OwnedSearchHit(OwnedSearchMatch):
+ """An exact representative match, optionally with other conversation matches."""
+
+ group: OwnedSearchGroup | None = None
+
+
+class OwnedSearchCoverage(BaseModel):
+ """Counts are source-store state, never total matches for this query."""
+
+ sync_id: UUID
+ active_records: int
+ capture: CaptureCoverage | None = None
+ pending_records: int
+ partially_indexed_records: int = 0
+ extraction_unavailable_records: int = 0
+ extraction_unknown_records: int = 0
+
+
+class OwnedRanking(BaseModel):
+ """Ranking applies only to this request's bounded canonical-source candidates."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ method: Literal["shared_rerank", "retrieval_rank"] = "retrieval_rank"
+ fallback_reason: (
+ Literal["unconfigured", "timeout", "provider_error", "invalid_output"] | None
+ ) = "unconfigured"
+ candidates_considered: int = Field(default=0, ge=0)
+ candidates_reranked: int = Field(default=0, ge=0)
+ shortlisted_candidates: int = Field(default=0, ge=0, le=200)
+ shortlist_limit: Literal[200] = 200
+ document_token_limit: Literal[2048] = 2048
+ token_count_basis: Literal["local_tokenizer"] = "local_tokenizer"
+ timeout_seconds: Literal[10] = 10
+ input_truncated_documents: int = Field(default=0, ge=0)
+ shortlist_truncated: bool = False
+
+
+class OwnedSearchResponse(BaseModel):
+ """Top results from bounded live candidate windows, without traversal promises."""
+
+ items: tuple[OwnedSearchHit, ...]
+ ranking: OwnedRanking = Field(default_factory=OwnedRanking)
+ sources: tuple[OwnedSearchCoverage, ...]
+ consistency: Literal["live"] = "live"
+ order: Literal["relevance"] = "relevance"
+ coverage: Literal["bounded_candidates"] = "bounded_candidates"
+ candidate_limit_per_collection: int = 200
+ candidate_window_full: bool
+ engine_partial: bool
+ excluded_candidates: int
+ postfilter_excluded: int
+ retrieval_incomplete: bool
diff --git a/backend/airweave/domains/search/protocols.py b/backend/airweave/domains/search/protocols.py
index f8403372b4..7e6f558743 100644
--- a/backend/airweave/domains/search/protocols.py
+++ b/backend/airweave/domains/search/protocols.py
@@ -13,6 +13,7 @@
SearchPlan,
SearchResults,
)
+from airweave.domains.search.types.embeddings import PreparedQueryEmbeddings
if TYPE_CHECKING:
from airweave.schemas.search_v2 import (
@@ -33,6 +34,10 @@ class SearchPlanExecutorProtocol(Protocol):
Used by all three tiers (instant, classic, agentic).
"""
+ async def prepare_query(self, plan: SearchPlan) -> PreparedQueryEmbeddings:
+ """Compute pure query work once for serial collection searches in one request."""
+ ...
+
async def execute(
self,
plan: SearchPlan,
@@ -42,6 +47,8 @@ async def execute(
ctx: ApiContext,
collection_readable_id: str,
user_principal: Optional[str] = None,
+ indexed_only: bool = False,
+ prepared_query: PreparedQueryEmbeddings | None = None,
) -> SearchResults:
"""Execute a search plan and return results."""
...
diff --git a/backend/airweave/domains/search/retrieval_strategy.py b/backend/airweave/domains/search/retrieval_strategy.py
new file mode 100644
index 0000000000..4a08ea1b90
--- /dev/null
+++ b/backend/airweave/domains/search/retrieval_strategy.py
@@ -0,0 +1,11 @@
+"""Public retrieval mode values, independent of execution and backend settings."""
+
+from enum import Enum
+
+
+class RetrievalStrategy(str, Enum):
+ """Supported retrieval strategies."""
+
+ SEMANTIC = "semantic"
+ KEYWORD = "keyword"
+ HYBRID = "hybrid"
diff --git a/backend/airweave/domains/search/tests/test_collection_metadata_builder.py b/backend/airweave/domains/search/tests/test_collection_metadata_builder.py
index de6a307abb..81b56be22e 100644
--- a/backend/airweave/domains/search/tests/test_collection_metadata_builder.py
+++ b/backend/airweave/domains/search/tests/test_collection_metadata_builder.py
@@ -23,6 +23,7 @@ def _make_source_connection(short_name: str = "slack", sync_id=None):
"""Create a mock source connection."""
sc = MagicMock()
sc.short_name = short_name
+ sc.is_authenticated = True
sc.sync_id = sync_id or uuid4()
return sc
@@ -178,3 +179,33 @@ def test_get_source_description_unknown(self, builder: CollectionMetadataBuilder
"""Unknown sources raise ValueError."""
with pytest.raises(ValueError, match="No description found"):
builder._get_source_description("unknown_source_xyz")
+
+
+async def test_wispr_collection_builds_metadata_for_search_planning(builder):
+ """An owned Wispr connection must not abort classic/agentic prompt construction."""
+ from airweave.platform.entities.wispr import WisprMeetingEntity
+
+ builder._collection_repo.get_by_readable_id.return_value = _make_collection()
+ builder._sc_repo.get_by_collection_ids.return_value = [_make_source_connection("wispr")]
+ builder._entity_definition_registry.list_for_source.return_value = [
+ _make_entity_definition_entry(
+ "WisprMeetingEntity", "WisprMeetingEntity", WisprMeetingEntity.model_json_schema()
+ )
+ ]
+ builder._entity_count_repo.get_counts_per_sync_and_type.return_value = []
+ builder._source_registry.get.return_value.federated_search = False
+ result = await builder.build(AsyncMock(), MagicMock(), "test-collection")
+ assert result.sources[0].short_name == "wispr"
+ assert "partial" in result.sources[0].description
+ assert result.sources[0].entity_types[0].fields
+
+
+async def test_unauthenticated_source_is_omitted_not_reported_empty(builder):
+ builder._collection_repo.get_by_readable_id.return_value = _make_collection()
+ source = _make_source_connection("wispr")
+ source.is_authenticated = False
+ builder._sc_repo.get_by_collection_ids.return_value = [source]
+ metadata = await builder.build(AsyncMock(), MagicMock(), "test-collection")
+ assert metadata.sources == []
+ builder._entity_definition_registry.list_for_source.assert_not_called()
+ builder._entity_count_repo.get_counts_per_sync_and_type.assert_not_called()
diff --git a/backend/airweave/domains/search/tests/test_executor.py b/backend/airweave/domains/search/tests/test_executor.py
index 1768da0777..b9d521f747 100644
--- a/backend/airweave/domains/search/tests/test_executor.py
+++ b/backend/airweave/domains/search/tests/test_executor.py
@@ -48,6 +48,17 @@
from airweave.platform.entities._base import AirweaveSystemMetadata, BaseEntity, Breadcrumb
from airweave.platform.entities.slack import SlackMessageEntity
+
+@pytest.fixture(autouse=True)
+def isolated_publication_gate(monkeypatch):
+ """Pipeline unit tests use fake retrieval; real SQL eligibility has PostgreSQL tests."""
+
+ async def passthrough(db, organization_id, collection, results, registry):
+ return results
+
+ monkeypatch.setattr("airweave.domains.search.executor.visible_results", passthrough)
+
+
# ── Helpers ──────────────────────────────────────────────────────────
@@ -483,9 +494,7 @@ def test_source_name_filter_keeps_matching(self):
def test_entity_type_filter(self):
slack = _make_federated_result(entity_id="f1", entity_type="SlackMessageEntity")
github = _make_search_result(entity_id="g1", entity_type="GitHubPREntity")
- filters = [
- _make_filter("airweave_system_metadata.entity_type", "equals", "GitHubPREntity")
- ]
+ filters = [_make_filter("airweave_system_metadata.entity_type", "equals", "GitHubPREntity")]
filtered = SearchPlanExecutor._apply_filters_in_memory([slack, github], filters)
assert len(filtered) == 1
assert filtered[0].entity_id == "g1"
@@ -494,9 +503,7 @@ def test_sync_id_filter_excludes_federated(self):
"""Federated results have sync_id=None, so filtering on sync_id excludes them."""
fed = _make_federated_result(entity_id="f1")
synced = _make_search_result(entity_id="s1", sync_id="some-sync-id")
- filters = [
- _make_filter("airweave_system_metadata.sync_id", "equals", "some-sync-id")
- ]
+ filters = [_make_filter("airweave_system_metadata.sync_id", "equals", "some-sync-id")]
filtered = SearchPlanExecutor._apply_filters_in_memory([fed, synced], filters)
assert len(filtered) == 1
assert filtered[0].entity_id == "s1"
@@ -533,19 +540,13 @@ def test_or_across_groups(self):
_make_filter("airweave_system_metadata.source_name", "equals", "slack"),
_make_filter("airweave_system_metadata.source_name", "equals", "github"),
]
- filtered = SearchPlanExecutor._apply_filters_in_memory(
- [slack, github, notion], filters
- )
+ filtered = SearchPlanExecutor._apply_filters_in_memory([slack, github, notion], filters)
assert len(filtered) == 2
assert {r.entity_id for r in filtered} == {"f1", "g1"}
def test_date_filter(self):
- old = _make_federated_result(
- entity_id="f1", created_at=datetime(2025, 1, 1)
- )
- new = _make_federated_result(
- entity_id="f2", created_at=datetime(2026, 3, 1)
- )
+ old = _make_federated_result(entity_id="f1", created_at=datetime(2025, 1, 1))
+ new = _make_federated_result(entity_id="f2", created_at=datetime(2026, 3, 1))
filters = [_make_filter("created_at", "greater_than", "2026-01-01T00:00:00")]
filtered = SearchPlanExecutor._apply_filters_in_memory([old, new], filters)
assert len(filtered) == 1
@@ -760,9 +761,7 @@ async def test_filter_source_name_github_excludes_slack(self):
source_lifecycle=source_lifecycle,
)
- user_filter = [
- _make_filter("airweave_system_metadata.source_name", "equals", "github")
- ]
+ user_filter = [_make_filter("airweave_system_metadata.source_name", "equals", "github")]
results = await executor.execute(
plan=_make_plan(),
@@ -774,9 +773,7 @@ async def test_filter_source_name_github_excludes_slack(self):
)
# Only GitHub results should remain — Slack filtered out in-memory
- assert all(
- r.airweave_system_metadata.source_name == "github" for r in results.results
- )
+ assert all(r.airweave_system_metadata.source_name == "github" for r in results.results)
@pytest.mark.asyncio
async def test_filter_source_name_slack_returns_federated_only(self):
@@ -805,9 +802,7 @@ async def test_filter_source_name_slack_returns_federated_only(self):
source_lifecycle=source_lifecycle,
)
- user_filter = [
- _make_filter("airweave_system_metadata.source_name", "equals", "slack")
- ]
+ user_filter = [_make_filter("airweave_system_metadata.source_name", "equals", "slack")]
results = await executor.execute(
plan=_make_plan(),
@@ -1082,9 +1077,7 @@ class TestExecutorFilterEdgeCases:
def test_empty_filter_groups_passes_all(self):
"""filter_groups=[] → all results pass through unfiltered."""
- results = [
- _make_federated_result(entity_id=f"f{i}") for i in range(5)
- ]
+ results = [_make_federated_result(entity_id=f"f{i}") for i in range(5)]
filtered = SearchPlanExecutor._apply_filters_in_memory(results, [])
assert len(filtered) == 5
assert filtered == results
@@ -1151,9 +1144,7 @@ async def search(self, query: str, limit: int) -> list:
source = _TransientSource()
executor = _build_executor()
- results = await executor._search_single_source(
- source, "test", 10, _make_ctx()
- )
+ results = await executor._search_single_source(source, "test", 10, _make_ctx())
assert results == []
assert source._attempt == 2
@@ -1201,3 +1192,105 @@ async def search(self, query: str, limit: int) -> list:
# 1 initial + 2 retries = 3 total
assert source._attempt == 3
+
+
+@pytest.mark.asyncio
+async def test_executor_returns_only_publication_gate_survivors(monkeypatch):
+ """Classic/instant callers never receive a rejected candidate for reranking."""
+ vector_db = FakeVectorDB()
+ fresh, stale = _make_search_result("fresh"), _make_search_result("stale")
+ vector_db.seed_results(SearchResults(results=[fresh, stale]))
+ gate = AsyncMock(return_value=[fresh])
+ monkeypatch.setattr("airweave.domains.search.executor.visible_results", gate)
+ executor = _build_executor(vector_db=vector_db)
+ result = await executor.execute(
+ plan=_make_plan(),
+ user_filter=[],
+ collection_id="col-1",
+ db=AsyncMock(),
+ ctx=_make_ctx(),
+ collection_readable_id="my-collection",
+ )
+ assert result.results == [fresh]
+ assert result.retrieval_incomplete is True
+ assert result.excluded_candidates == 1
+ assert gate.await_args.args[3] == [fresh, stale]
+
+
+@pytest.mark.parametrize("mode,dense_calls", [("hybrid", 1), ("keyword", 0)])
+async def test_prepared_query_reuses_only_embeddings_across_collection_scopes(mode, dense_calls):
+ executor = _build_executor()
+ executor._dense_embedder.embed_many = AsyncMock(wraps=executor._dense_embedder.embed_many)
+ executor._sparse_embedder.embed = AsyncMock(wraps=executor._sparse_embedder.embed)
+ executor._vector_db.compile_query = AsyncMock(wraps=executor._vector_db.compile_query)
+ executor._resolve_acl_principals = AsyncMock(side_effect=[["alice"], ["bob"]])
+ plan = SearchPlan(
+ query=SearchQuery(primary="meeting", variations=["discussion"]),
+ limit=200,
+ offset=0,
+ retrieval_strategy=mode,
+ )
+ prepared = await executor.prepare_query(plan)
+ ctx = MagicMock()
+ for collection in ("first", "second"):
+ filters = [
+ FilterGroup(
+ conditions=[
+ FilterCondition(
+ field="airweave_system_metadata.sync_id",
+ operator="equals",
+ value=collection,
+ )
+ ]
+ )
+ ]
+ result = await executor.execute(
+ plan,
+ filters,
+ collection,
+ MagicMock(),
+ ctx,
+ collection,
+ indexed_only=True,
+ prepared_query=prepared,
+ )
+ assert result.results == []
+ assert executor._dense_embedder.embed_many.await_count == dense_calls
+ if dense_calls:
+ executor._dense_embedder.embed_many.assert_awaited_once_with(
+ ["meeting", "discussion"], purpose="query"
+ )
+ executor._sparse_embedder.embed.assert_awaited_once_with("meeting")
+ calls = executor._vector_db.compile_query.await_args_list
+ assert [call.kwargs["collection_id"] for call in calls] == ["first", "second"]
+ assert [call.kwargs["acl_principals"] for call in calls] == [["alice"], ["bob"]]
+ assert all(call.kwargs["plan"].limit == 200 for call in calls)
+ assert [call.kwargs["plan"].filter_groups[0].conditions[0].value for call in calls] == [
+ "first",
+ "second",
+ ]
+ # A new request has no shared cache, including an identical query.
+ fresh = await executor.prepare_query(plan)
+ assert fresh is not prepared
+ assert executor._sparse_embedder.embed.await_count == 2
+
+
+async def test_prepared_query_rejects_other_executor_query_variation_and_mode():
+ executor = _build_executor()
+ plan = SearchPlan(
+ query=SearchQuery(primary="meeting"), limit=200, offset=0, retrieval_strategy="keyword"
+ )
+ prepared = await executor.prepare_query(plan)
+ other = _build_executor()
+ for candidate, candidate_plan in [
+ (other, plan),
+ (executor, plan.model_copy(update={"query": SearchQuery(primary="other")})),
+ (
+ executor,
+ plan.model_copy(update={"query": SearchQuery(primary="meeting", variations=["x"])}),
+ ),
+ (executor, plan.model_copy(update={"retrieval_strategy": "semantic"})),
+ ]:
+ with pytest.raises(ValueError, match="Prepared query"):
+ await candidate._execute_vector_search(candidate_plan, "collection", None, prepared)
+ assert executor._vector_db._calls == other._vector_db._calls == []
diff --git a/backend/airweave/domains/search/tests/test_owned_grouping.py b/backend/airweave/domains/search/tests/test_owned_grouping.py
new file mode 100644
index 0000000000..69cc4e7a27
--- /dev/null
+++ b/backend/airweave/domains/search/tests/test_owned_grouping.py
@@ -0,0 +1,113 @@
+"""Conversation grouping retains exact, bounded independent originals."""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+from pydantic import ValidationError
+
+from airweave.domains.entities.canonical.extraction_models import (
+ ExtractionCoverage,
+ ExtractionOutcome,
+)
+from airweave.domains.search.owned import OwnedSearchService, _EnrichmentRecord
+from airweave.domains.search.owned_models import OwnedSearchGroup, OwnedSearchHit, OwnedSearchMatch
+
+
+def match(sync_id=None, provider="almanac", kind="message", native="m", container="s"):
+ return OwnedSearchHit(
+ record_id=uuid4(),
+ revision=3,
+ sync_id=sync_id or uuid4(),
+ source_connection_id=uuid4(),
+ provider=provider,
+ identity={"record_type": kind, "native_id": native, "container_id": container},
+ title=native,
+ excerpts=("best " + native, "other"),
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=None,
+ source_updated_at=None,
+ completeness="complete",
+ )
+
+
+def row(hit, **changes):
+ values = {
+ "id": hit.record_id,
+ "sync_id": hit.sync_id,
+ "record_revision": 3,
+ "indexed_generation": uuid4(),
+ "indexed_pipeline_version": 2,
+ "entity_definition_short_name": hit.identity.record_type,
+ "native_id": hit.identity.native_id,
+ "container_id": hit.identity.container_id,
+ "parent_record_type": "session",
+ "parent_native_id": "s",
+ "parent_container_id": None,
+ "observed_at": hit.observed_at,
+ "source_created_at": None,
+ "source_updated_at": None,
+ "completeness": "complete",
+ "email_thread_id": None,
+ }
+ values.update(changes)
+ return _EnrichmentRecord(**values)
+
+
+def test_source_scoped_group_keeps_exact_anchors_and_observed_count():
+ sync = uuid4()
+ hits = [match(sync, native=str(i)) for i in range(6)]
+ other = match(native="other")
+ for hit in hits:
+ hit.extraction = ExtractionCoverage(
+ parts=(ExtractionOutcome(part_index=0, key="body", kind="body", outcome="indexed"),)
+ )
+ for hit in [*hits, other]:
+ hit.group = OwnedSearchService._conversation(row(hit), hit)
+ grouped = OwnedSearchService._group_hits([hits[0], other, *hits[1:]])
+ assert [h.record_id for h in grouped] == [hits[0].record_id, other.record_id]
+ group = grouped[0].group
+ assert group.native_id == "s" and group.matched_records == 6
+ assert len(group.additional_matches) == 3
+ for original, extra in zip(hits[1:4], group.additional_matches, strict=True):
+ assert type(extra) is OwnedSearchMatch
+ assert (extra.record_id, extra.identity, extra.revision) == (
+ original.record_id,
+ original.identity,
+ original.revision,
+ )
+ assert extra.excerpts == original.excerpts[:1]
+ assert "group" not in extra.model_dump()
+ assert extra.extraction == original.extraction
+ assert extra.model_dump()["extraction"]["status"] == "complete"
+
+
+def test_only_admitted_sessions_or_validated_gmail_threads_group():
+ root = match(kind="session", native="s", container=None)
+ root.group = OwnedSearchService._conversation(
+ row(root, parent_record_type=None, parent_native_id=None), root
+ )
+ child = match(root.sync_id)
+ child.group = OwnedSearchService._conversation(row(child), child)
+ assert OwnedSearchService._group_hits([child, root])[0].group.matched_records == 2
+ for provider, kind in [
+ ("google_calendar", "event"),
+ ("slack", "message"),
+ ("almanac", "knowledge"),
+ ("gdrive", "file"),
+ ]:
+ hit = match(provider=provider, kind=kind)
+ assert OwnedSearchService._conversation(row(hit), hit) is None
+ assert OwnedSearchService._conversation(row(child, parent_native_id="wrong"), child) is None
+ assert OwnedSearchService._conversation(row(child, parent_container_id="nested"), child) is None
+ gmail = match(provider="gmail")
+ gmail.email_thread_id = "thread_1"
+ assert OwnedSearchService._conversation(row(gmail), gmail).kind == "email_thread"
+ with pytest.raises(ValidationError):
+ OwnedSearchGroup(kind="record", native_id="x", matched_records=1)
+ with pytest.raises(ValidationError):
+ OwnedSearchGroup(kind="session", native_id="x", matched_records=0)
+ with pytest.raises(ValidationError):
+ OwnedSearchGroup(
+ kind="session", native_id="x", matched_records=5, additional_matches=[gmail] * 4
+ )
diff --git a/backend/airweave/domains/search/tests/test_owned_ranking.py b/backend/airweave/domains/search/tests/test_owned_ranking.py
new file mode 100644
index 0000000000..bce98eae4c
--- /dev/null
+++ b/backend/airweave/domains/search/tests/test_owned_ranking.py
@@ -0,0 +1,196 @@
+"""Shared ranking orchestration with fake ports; SQL fencing is separately integrated."""
+
+import asyncio
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+from uuid import UUID, uuid4
+
+import pytest
+
+from airweave.core.protocols.reranker import RerankerResult
+from airweave.domains.search.owned import OwnedSearchService
+from airweave.domains.search.owned_models import OwnedSearchRequest
+from airweave.domains.search.types import SearchResults
+
+
+class CharacterTokenizer:
+ """Conservative fake: one character per token, including Unicode."""
+
+ def count_tokens(self, text):
+ return len(text)
+
+
+def service(rerank=None):
+ return OwnedSearchService(
+ MagicMock(),
+ MagicMock(),
+ reranker=SimpleNamespace(rerank=rerank) if rerank else None,
+ tokenizer=CharacterTokenizer(),
+ )
+
+
+async def test_one_shared_call_with_exact_budget_and_no_raw_payload():
+ keys = [uuid4() for _ in range(205)]
+ hits = {
+ key: SimpleNamespace(title="Title " + "x" * 300, raw_source_fields={"secret": "NEVER"})
+ for key in keys
+ }
+ matched = {key: "🎵" * 3000 for key in keys}
+ rerank = AsyncMock(return_value=[RerankerResult(i, float(i)) for i in range(200)])
+ ranked, meta = await service(rerank)._rank("query", keys, hits, matched)
+ assert ranked == keys[199::-1]
+ rerank.assert_awaited_once()
+ query, documents = rerank.call_args.args
+ assert query == "query" and len(documents) == 200
+ assert all(len(doc) <= 2048 and "NEVER" not in doc for doc in documents)
+ assert meta.candidates_considered == 205 and meta.candidates_reranked == 200
+ assert meta.shortlist_truncated and meta.input_truncated_documents == 200
+ assert meta.method == "shared_rerank" and meta.fallback_reason is None
+
+
+@pytest.mark.parametrize(
+ "output",
+ [
+ [RerankerResult(0, 0.5)],
+ [RerankerResult(0, 0.5), RerankerResult(0, 0.3)],
+ [RerankerResult(True, 0.5), RerankerResult(0, 0.3)],
+ [RerankerResult(-1, 0.5), RerankerResult(0, 0.3)],
+ [RerankerResult(0, float("nan")), RerankerResult(1, 0.3)],
+ [RerankerResult(0, float("inf")), RerankerResult(1, 0.3)],
+ [{"index": 0}, {"index": 1}],
+ None,
+ ],
+)
+async def test_invalid_output_has_explicit_deterministic_fallback(output):
+ keys = [uuid4(), uuid4()]
+ hits = {key: SimpleNamespace(title="Title") for key in keys}
+ ranked, meta = await service(AsyncMock(return_value=output))._rank(
+ "query",
+ keys,
+ hits,
+ {key: "Text" for key in keys},
+ )
+ assert ranked == keys and meta.fallback_reason == "invalid_output"
+ assert meta.method == "retrieval_rank" and meta.candidates_reranked == 0
+
+
+@pytest.mark.parametrize(
+ "error,reason", [(TimeoutError(), "timeout"), (RuntimeError("secret"), "provider_error")]
+)
+async def test_error_fallback_is_sanitized(error, reason):
+ key = uuid4()
+ ranked, meta = await service(AsyncMock(side_effect=error))._rank(
+ "query",
+ [key],
+ {key: SimpleNamespace(title="Title")},
+ {key: "Text"},
+ )
+ assert ranked == [key] and meta.fallback_reason == reason
+ assert "secret" not in meta.model_dump_json()
+
+
+async def test_absent_port_and_cancellation():
+ key = uuid4()
+ args = ("q", [key], {key: SimpleNamespace(title="Title")}, {key: "Text"})
+ assert (await service()._rank(*args))[1].fallback_reason == "unconfigured"
+ with pytest.raises(asyncio.CancelledError):
+ await service(AsyncMock(side_effect=asyncio.CancelledError()))._rank(*args)
+
+
+@pytest.mark.parametrize("withdraw", ["before", "during", None])
+async def test_global_orchestration_rechecks_before_and_after_without_holding_session(withdraw):
+ keys = [UUID(int=1), UUID(int=2)]
+ syncs = [uuid4(), uuid4()]
+ collections = [uuid4(), uuid4()]
+ hits = {key: SimpleNamespace(title="Title", group=None) for key in keys}
+ active = 0
+
+ @asynccontextmanager
+ async def sessions():
+ nonlocal active
+ active += 1
+ try:
+ yield object()
+ finally:
+ active -= 1
+
+ async def rerank(query, documents, top_n):
+ assert active == 0
+ assert len(documents) == (1 if withdraw == "before" else 2)
+ return [RerankerResult(index, float(index)) for index in range(top_n)]
+
+ search = service(AsyncMock(side_effect=rerank))
+ search._executor.prepare_query = AsyncMock(return_value=None)
+ search._executor.execute = AsyncMock(return_value=SearchResults(results=[]))
+ search._resolve_scopes = AsyncMock(
+ return_value=(
+ {},
+ {(collections[0], "one"): [syncs[0]], (collections[1], "two"): [syncs[1]]},
+ )
+ )
+ search._enrich = AsyncMock(
+ side_effect=[
+ ({keys[0]: hits[keys[0]]}, {keys[0]: (1, None)}, {keys[0]: "First collection"}, 0, 0),
+ (
+ {keys[1]: hits[keys[1]]},
+ {keys[1]: (0.5, None)},
+ {keys[1]: "Second collection"},
+ 0,
+ 0,
+ ),
+ ]
+ )
+ search._final_publications = AsyncMock(
+ side_effect=[
+ {keys[1]} if withdraw == "before" else set(keys),
+ {keys[1]} if withdraw else set(keys),
+ ]
+ )
+ search._coverage = AsyncMock(return_value=())
+ # Observe final ordering without replacing the actual ranking/session flow.
+ from unittest.mock import patch
+
+ import airweave.domains.search.owned as module
+
+ with patch.object(module, "OwnedSearchResponse", side_effect=lambda **fields: fields):
+ result = await search.search(
+ sessions,
+ SimpleNamespace(),
+ OwnedSearchRequest(query="q", sync_ids=tuple(syncs), limit=1),
+ )
+ assert result["items"] == (hits[keys[1]],)
+ assert search._final_publications.await_count == 2
+ assert search._resolve_scopes.await_count == 3
+ assert result["excluded_candidates"] == int(withdraw is not None)
+
+
+async def test_timeout_enforced_around_remote_await(monkeypatch):
+ import airweave.domains.search.owned as module
+
+ timeout = asyncio.timeout
+ limits = []
+
+ def short_timeout(seconds):
+ limits.append(seconds)
+ return timeout(0.001)
+
+ async def slow(*args, **kwargs):
+ await asyncio.sleep(1)
+
+ monkeypatch.setattr(module.asyncio, "timeout", short_timeout)
+ key = uuid4()
+ ranked, metadata = await service(AsyncMock(side_effect=slow))._rank(
+ "query",
+ [key],
+ {key: SimpleNamespace(title="Title")},
+ {key: "Text"},
+ )
+ assert limits == [10] and ranked == [key] and metadata.fallback_reason == "timeout"
+
+
+async def test_empty_candidates_need_no_ranking_fallback():
+ ranked, metadata = await service()._rank("q", [], {}, {})
+ assert ranked == [] and metadata.method == "retrieval_rank"
+ assert metadata.fallback_reason is None and metadata.candidates_considered == 0
+ assert metadata.token_count_basis == "local_tokenizer"
diff --git a/backend/airweave/domains/search/tests/test_visible_vector_db.py b/backend/airweave/domains/search/tests/test_visible_vector_db.py
new file mode 100644
index 0000000000..536e2e1fa1
--- /dev/null
+++ b/backend/airweave/domains/search/tests/test_visible_vector_db.py
@@ -0,0 +1,163 @@
+"""Every vector content retrieval must pass authoritative publication validation."""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+from uuid import uuid4
+
+import pytest
+
+from airweave.domains.search.adapters.vector_db.exceptions import VectorDBError
+from airweave.domains.search.agentic.state import AgentState
+from airweave.domains.search.agentic.tests.conftest import make_result
+from airweave.domains.search.agentic.tools.read import ReadTool
+from airweave.domains.search.types.results import SearchResults
+from airweave.domains.search.visible_vector_db import UnavailableExactCount, VisibleVectorDB
+
+
+def wrapper():
+ delegate = AsyncMock()
+ db = AsyncMock()
+ ctx = SimpleNamespace(organization=SimpleNamespace(id=uuid4()))
+ registry = Mock()
+ return (
+ VisibleVectorDB(delegate, db, ctx, "readable", "collection", registry),
+ delegate,
+ db,
+ registry,
+ )
+
+
+@pytest.mark.asyncio
+async def test_execute_rejects_stale_candidate_and_reports_underfilled_page(monkeypatch):
+ visible, delegate, _, _ = wrapper()
+ fresh, stale = make_result("fresh"), make_result("stale", content="do not expose")
+ delegate.execute_query.return_value = SearchResults(results=[fresh, stale])
+ gate = AsyncMock(return_value=[fresh])
+ monkeypatch.setattr("airweave.domains.search.visible_vector_db.visible_results", gate)
+ result = await visible.execute_query("compiled")
+ assert result.results == [fresh]
+ assert result.retrieval_incomplete is True
+ assert result.excluded_candidates == 1
+ assert gate.await_args.args[3] == [fresh, stale]
+
+
+@pytest.mark.asyncio
+async def test_filter_read_validates_before_returning_original_text(monkeypatch):
+ visible, delegate, _, _ = wrapper()
+ stale = make_result("stale", content="do not expose")
+ delegate.filter_search.return_value = [stale]
+ gate = AsyncMock(return_value=[])
+ monkeypatch.setattr("airweave.domains.search.visible_vector_db.visible_results", gate)
+ assert await visible.filter_search([], "collection") == []
+ gate.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_cached_read_does_not_fall_back_after_gate_removes_chunks(monkeypatch):
+ visible, delegate, _, _ = wrapper()
+ stale = make_result("stale", content="do not expose")
+ delegate.filter_search.return_value = [stale]
+ monkeypatch.setattr(
+ "airweave.domains.search.visible_vector_db.visible_results", AsyncMock(return_value=[])
+ )
+ state = AgentState()
+ state.results[stale.entity_id] = stale
+ result = await ReadTool(visible, "collection").execute({"entity_ids": [stale.entity_id]}, state)
+ assert result.entities == []
+ assert result.not_found == ["stale"]
+
+
+@pytest.mark.asyncio
+async def test_cross_collection_reads_fail_before_delegate_request():
+ visible, delegate, _, _ = wrapper()
+ with pytest.raises(VectorDBError, match="Collection"):
+ await visible.filter_search([], "other")
+ with pytest.raises(VectorDBError, match="Collection"):
+ await visible.compile_query(None, None, "other")
+ with pytest.raises(VectorDBError, match="Collection"):
+ await visible.count([], "other")
+ assert delegate.mock_calls == []
+
+
+@pytest.mark.asyncio
+async def test_canonical_count_is_unavailable_instead_of_stale_index_total():
+ visible, delegate, db, registry = wrapper()
+ db.scalars.return_value = [SimpleNamespace(short_name="gmail", sync_id=uuid4())]
+ registry.get.return_value = SimpleNamespace(
+ source_class_ref=SimpleNamespace(canonical_record_types=("message",))
+ )
+ with pytest.raises(UnavailableExactCount):
+ await visible.count([], "collection")
+ delegate.count.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_gate_failure_propagates_without_returning_unvalidated_data(monkeypatch):
+ visible, delegate, _, _ = wrapper()
+ delegate.filter_search.return_value = [make_result("unvalidated")]
+ monkeypatch.setattr(
+ "airweave.domains.search.visible_vector_db.visible_results",
+ AsyncMock(side_effect=RuntimeError("database unavailable")),
+ )
+ with pytest.raises(RuntimeError, match="database unavailable"):
+ await visible.filter_search([], "collection")
+
+
+@pytest.mark.asyncio
+async def test_wrapper_close_preserves_shared_adapter():
+ visible, delegate, _, _ = wrapper()
+ await visible.close()
+ delegate.close.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_cached_read_never_masks_vector_failure_with_old_text():
+ visible, delegate, _, _ = wrapper()
+ stale = make_result("stale", content="must not appear after failure")
+ delegate.filter_search.side_effect = VectorDBError("vector unavailable")
+ state = AgentState()
+ state.results[stale.entity_id] = stale
+ with pytest.raises(VectorDBError, match="vector unavailable"):
+ await ReadTool(visible, "collection").execute({"entity_ids": ["stale"]}, state)
+
+
+@pytest.mark.asyncio
+async def test_legacy_count_ands_authorized_syncs_into_each_or_group():
+ from airweave.domains.search.types.filters import (
+ FilterableField,
+ FilterCondition,
+ FilterGroup,
+ FilterOperator,
+ )
+
+ visible, delegate, db, registry = wrapper()
+ sync_id = uuid4()
+ db.scalars.return_value = [SimpleNamespace(short_name="legacy", sync_id=sync_id)]
+ registry.get.return_value = SimpleNamespace(source_class_ref=SimpleNamespace())
+ delegate.count.return_value = 7
+ groups = [
+ FilterGroup(
+ conditions=[
+ FilterCondition(
+ field=FilterableField.NAME, operator=FilterOperator.EQUALS, value=name
+ )
+ ]
+ )
+ for name in ("A", "B")
+ ]
+ assert await visible.count(groups, "collection") == 7
+ passed = delegate.count.await_args.args[0]
+ assert len(passed) == 2
+ for original, scoped in zip(groups, passed, strict=True):
+ assert scoped.conditions[:-1] == original.conditions
+ assert scoped.conditions[-1].field == FilterableField.SYSTEM_METADATA_SYNC_ID
+ assert scoped.conditions[-1].operator == FilterOperator.IN
+ assert scoped.conditions[-1].value == [str(sync_id)]
+
+
+@pytest.mark.asyncio
+async def test_no_authorized_source_count_does_not_use_index_total():
+ visible, delegate, db, _ = wrapper()
+ db.scalars.return_value = []
+ assert await visible.count([], "collection") == 0
+ delegate.count.assert_not_called()
diff --git a/backend/airweave/domains/search/types/__init__.py b/backend/airweave/domains/search/types/__init__.py
index d5c0d58952..06e524767f 100644
--- a/backend/airweave/domains/search/types/__init__.py
+++ b/backend/airweave/domains/search/types/__init__.py
@@ -4,6 +4,7 @@
from airweave.domains.search.types import SearchResult, FilterGroup, SearchPlan
"""
+from airweave.domains.search.retrieval_strategy import RetrievalStrategy
from airweave.domains.search.types.embeddings import QueryEmbeddings
from airweave.domains.search.types.filters import (
FilterableField,
@@ -18,7 +19,6 @@
SourceMetadata,
)
from airweave.domains.search.types.plan import (
- RetrievalStrategy,
SearchPlan,
SearchQuery,
)
diff --git a/backend/airweave/domains/search/types/embeddings.py b/backend/airweave/domains/search/types/embeddings.py
index 8cc403bdf4..48adb3cfd3 100644
--- a/backend/airweave/domains/search/types/embeddings.py
+++ b/backend/airweave/domains/search/types/embeddings.py
@@ -7,9 +7,10 @@
from typing import Optional
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
from airweave.domains.embedders.types import DenseEmbedding, SparseEmbedding
+from airweave.domains.search.types.plan import RetrievalStrategy, SearchPlan
class QueryEmbeddings(BaseModel):
@@ -21,3 +22,25 @@ class QueryEmbeddings(BaseModel):
sparse_embedding: Optional[SparseEmbedding] = Field(
default=None, description="Sparse embedding for the primary query only."
)
+
+
+class PreparedQueryEmbeddings(BaseModel):
+ """Internal request-local work, valid only for its creating executor and query."""
+
+ model_config = ConfigDict(frozen=True, extra="forbid")
+ primary: str
+ variations: tuple[str, ...]
+ strategy: RetrievalStrategy
+ embeddings: QueryEmbeddings
+ _owner: object | None = PrivateAttr(default=None)
+
+ def require_match(self, plan: SearchPlan, owner: object) -> QueryEmbeddings:
+ """Never reuse another model's embeddings or silently change a query."""
+ if (
+ self._owner is not owner
+ or self.primary != plan.query.primary
+ or self.variations != tuple(plan.query.variations)
+ or self.strategy != plan.retrieval_strategy
+ ):
+ raise ValueError("Prepared query does not belong to this executor and search")
+ return self.embeddings
diff --git a/backend/airweave/domains/search/types/filters.py b/backend/airweave/domains/search/types/filters.py
index a1f03435ac..14d462197b 100644
--- a/backend/airweave/domains/search/types/filters.py
+++ b/backend/airweave/domains/search/types/filters.py
@@ -38,6 +38,11 @@ class FilterableField(str, Enum):
SYSTEM_METADATA_CHUNK_INDEX = "airweave_system_metadata.chunk_index"
SYSTEM_METADATA_SYNC_ID = "airweave_system_metadata.sync_id"
SYSTEM_METADATA_SYNC_JOB_ID = "airweave_system_metadata.sync_job_id"
+ CANONICAL_RECORD_TYPE = "airweave_system_metadata.canonical_record_type"
+ CANONICAL_SOURCE_CREATED_US = "airweave_system_metadata.source_created_us"
+ CANONICAL_SOURCE_UPDATED_US = "airweave_system_metadata.source_updated_us"
+ CANONICAL_SOURCE_CREATED_KNOWN = "airweave_system_metadata.source_created_known"
+ CANONICAL_SOURCE_UPDATED_KNOWN = "airweave_system_metadata.source_updated_known"
class FilterOperator(str, Enum):
@@ -66,6 +71,7 @@ class FilterOperator(str, Enum):
FilterableField.BREADCRUMBS_ENTITY_TYPE,
FilterableField.SYSTEM_METADATA_ENTITY_TYPE,
FilterableField.SYSTEM_METADATA_SOURCE_NAME,
+ FilterableField.CANONICAL_RECORD_TYPE,
FilterableField.SYSTEM_METADATA_ORIGINAL_ENTITY_ID,
FilterableField.SYSTEM_METADATA_SYNC_ID,
FilterableField.SYSTEM_METADATA_SYNC_JOB_ID,
@@ -82,6 +88,10 @@ class FilterOperator(str, Enum):
_NUMERIC_FIELDS: frozenset[FilterableField] = frozenset(
{
FilterableField.SYSTEM_METADATA_CHUNK_INDEX,
+ FilterableField.CANONICAL_SOURCE_CREATED_US,
+ FilterableField.CANONICAL_SOURCE_UPDATED_US,
+ FilterableField.CANONICAL_SOURCE_CREATED_KNOWN,
+ FilterableField.CANONICAL_SOURCE_UPDATED_KNOWN,
}
)
diff --git a/backend/airweave/domains/search/types/metadata.py b/backend/airweave/domains/search/types/metadata.py
index a574e007c5..5793543cea 100644
--- a/backend/airweave/domains/search/types/metadata.py
+++ b/backend/airweave/domains/search/types/metadata.py
@@ -12,7 +12,9 @@ class EntityTypeMetadata(BaseModel):
"""Entity type metadata with fields."""
name: str = Field(..., description="Entity type name (e.g., 'NotionPageEntity').")
- count: int = Field(..., description="Number of entities of this type.")
+ count: int | None = Field(
+ ..., description="Known projection entity count, or null when not measured."
+ )
fields: dict[str, str] = Field(..., description="Field names mapped to their descriptions.")
def to_md(self) -> str:
@@ -23,7 +25,8 @@ def to_md(self) -> str:
self-explanatory enough for the LLM to craft queries and filters.
"""
field_names = ", ".join(sorted(self.fields.keys()))
- return f"- **{self.name}** ({self.count}) — fields: {field_names}"
+ count = self.count if self.count is not None else "indexed count unknown"
+ return f"- **{self.name}** ({count}) — fields: {field_names}"
class SourceMetadata(BaseModel):
@@ -34,6 +37,10 @@ class SourceMetadata(BaseModel):
entity_types: list[EntityTypeMetadata] = Field(
..., description="Entity types with their fields and counts."
)
+ captured_record_counts: dict[str, int] | None = Field(
+ default=None,
+ description="Visible captured records by native record type; not index counts.",
+ )
federated: bool = Field(
default=False,
description="Whether this source uses federated (real-time) search instead of syncing.",
@@ -53,12 +60,30 @@ def to_md(self) -> str:
f"Filters work normally. No entity counts available.*"
)
- # Only show entity types that actually have entities
- non_empty = [et for et in self.entity_types if et.count > 0]
+ if self.captured_record_counts is not None:
+ counts = ", ".join(
+ f"{kind}: {count}" for kind, count in sorted(self.captured_record_counts.items())
+ )
+ lines = [
+ f"### {self.short_name}",
+ self.description,
+ f"Captured records ({sum(self.captured_record_counts.values())} total): {counts}.",
+ "These are captured records; search indexing may lag or be incomplete.",
+ "Searchable projection schemas (counts are not measured):",
+ *(entity.to_md() for entity in self.entity_types),
+ ]
+ return "\n".join(lines)
+
+ # Only show known nonempty or unmeasured projection entity types.
+ non_empty = [et for et in self.entity_types if et.count is None or et.count > 0]
if not non_empty:
return f"### {self.short_name}\n{self.description}\n\n*(no entities)*"
- total = sum(et.count for et in non_empty)
+ total = (
+ sum(et.count for et in non_empty)
+ if all(et.count is not None for et in non_empty)
+ else "unknown"
+ )
lines = [
f"### {self.short_name} ({total} entities total)",
f"{self.description}",
diff --git a/backend/airweave/domains/search/types/plan.py b/backend/airweave/domains/search/types/plan.py
index c9af96f9bc..64077b31ba 100644
--- a/backend/airweave/domains/search/types/plan.py
+++ b/backend/airweave/domains/search/types/plan.py
@@ -5,22 +5,14 @@
from __future__ import annotations
-from enum import Enum
from typing import List
from pydantic import BaseModel, Field
+from airweave.domains.search.retrieval_strategy import RetrievalStrategy
from airweave.domains.search.types.filters import FilterGroup
-class RetrievalStrategy(str, Enum):
- """Supported retrieval strategies."""
-
- SEMANTIC = "semantic"
- KEYWORD = "keyword"
- HYBRID = "hybrid"
-
-
class SearchQuery(BaseModel):
"""Search query."""
diff --git a/backend/airweave/domains/search/types/results.py b/backend/airweave/domains/search/types/results.py
index 367b94efd4..c1fe9459f2 100644
--- a/backend/airweave/domains/search/types/results.py
+++ b/backend/airweave/domains/search/types/results.py
@@ -100,13 +100,16 @@ class SearchResult(BaseModel):
)
textual_representation: str = Field(..., description="Semantically searchable text content")
+ query_snippet: str | None = Field(
+ default=None, exclude=True, description="Optional Vespa lexical fragment; display only."
+ )
airweave_system_metadata: SearchSystemMetadata = Field(..., description="System metadata")
access: SearchAccessControl = Field(..., description="Access control")
web_url: str = Field(
...,
- description="URL to view the entity in its source application (e.g., Notion, Asana).",
+ description="Source application URL, or empty when the record uses a canonical locator.",
)
url: Optional[str] = Field(
@@ -239,6 +242,19 @@ class SearchResults(BaseModel):
description="Search results ordered by relevance (highest first).",
)
+ engine_partial: bool = False
+ engine_coverage_percent: float | None = None
+
+ retrieval_incomplete: bool = Field(
+ default=False,
+ description="Some retrieved candidates were unavailable; this page may underfill.",
+ )
+ excluded_candidates: int = Field(
+ default=0,
+ ge=0,
+ description="Stale or unauthorized index candidates omitted from this page.",
+ )
+
def __len__(self) -> int:
"""Return the number of results."""
return len(self.results)
diff --git a/backend/airweave/domains/search/visible_vector_db.py b/backend/airweave/domains/search/visible_vector_db.py
new file mode 100644
index 0000000000..cde8b80358
--- /dev/null
+++ b/backend/airweave/domains/search/visible_vector_db.py
@@ -0,0 +1,138 @@
+"""Request-scoped vector reads validated against the authoritative publication store."""
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.api.context import ApiContext
+from airweave.domains.search.adapters.vector_db.exceptions import VectorDBError
+from airweave.domains.search.adapters.vector_db.protocol import VectorDBProtocol
+from airweave.domains.search.canonical_visibility import visible_results
+from airweave.domains.search.types import QueryEmbeddings, SearchPlan, SearchResults
+from airweave.domains.search.types.filters import (
+ FilterableField,
+ FilterCondition,
+ FilterGroup,
+ FilterOperator,
+)
+from airweave.domains.search.types.results import CompiledQuery, SearchResult
+from airweave.domains.sources.protocols import SourceRegistryProtocol
+from airweave.models.source_connection import SourceConnection
+
+
+class UnavailableExactCount(VectorDBError):
+ """The index cannot prove an exact currently visible record count."""
+
+
+class VisibleVectorDB:
+ """Same adapter contract with current publication checks on every content retrieval."""
+
+ def __init__(
+ self,
+ delegate: VectorDBProtocol,
+ db: AsyncSession,
+ ctx: ApiContext,
+ readable_id: str,
+ collection_id: str,
+ registry: SourceRegistryProtocol,
+ ) -> None:
+ """Bind one authenticated collection; callers cannot widen that scope."""
+ self._delegate = delegate
+ self._db = db
+ self._ctx = ctx
+ self._readable_id = readable_id
+ self._collection_id = collection_id
+ self._registry = registry
+
+ def _check_scope(self, collection_id: str) -> None:
+ if collection_id != self._collection_id:
+ raise VectorDBError("Collection does not match this search request")
+
+ async def validate(self, results: list[SearchResult]) -> list[SearchResult]:
+ """Revalidate even previously cached results before further consumption."""
+ return await visible_results(
+ self._db, self._ctx.organization.id, self._readable_id, results, self._registry
+ )
+
+ async def compile_query(
+ self,
+ plan: SearchPlan,
+ embeddings: QueryEmbeddings,
+ collection_id: str,
+ acl_principals: list[str] | None = None,
+ ) -> CompiledQuery:
+ """Compile only within the bound collection."""
+ self._check_scope(collection_id)
+ return await self._delegate.compile_query(plan, embeddings, collection_id, acl_principals)
+
+ async def execute_query(self, compiled_query: CompiledQuery) -> SearchResults:
+ """Validate returned content regardless of query fields."""
+ result = await self._delegate.execute_query(compiled_query)
+ visible = await self.validate(result.results)
+ excluded = len(result.results) - len(visible)
+ return result.model_copy(
+ update={
+ "results": visible,
+ "retrieval_incomplete": result.retrieval_incomplete or excluded > 0,
+ "excluded_candidates": result.excluded_candidates + excluded,
+ }
+ )
+
+ async def filter_search(
+ self,
+ filter_groups: list[FilterGroup],
+ collection_id: str,
+ limit: int = 50,
+ offset: int = 0,
+ name_substring: str | None = None,
+ ) -> list[SearchResult]:
+ """Validate read and navigation candidates before returning any text."""
+ self._check_scope(collection_id)
+ results = await self._delegate.filter_search(
+ filter_groups, collection_id, limit, offset, name_substring
+ )
+ return await self.validate(results)
+
+ async def count(
+ self,
+ filter_groups: list[FilterGroup],
+ collection_id: str,
+ name_substring: str | None = None,
+ ) -> int:
+ """Legacy index counts remain available; canonical counts need database semantics."""
+ self._check_scope(collection_id)
+ sources = list(
+ await self._db.scalars(
+ select(SourceConnection).where(
+ SourceConnection.organization_id == self._ctx.organization.id,
+ SourceConnection.readable_collection_id == self._readable_id,
+ SourceConnection.is_authenticated.is_(True),
+ SourceConnection.sync_id.is_not(None),
+ )
+ )
+ )
+ for source in sources:
+ if getattr(
+ self._registry.get(source.short_name).source_class_ref,
+ "canonical_record_types",
+ (),
+ ):
+ raise UnavailableExactCount(
+ "Exact current count is unavailable for this indexed collection; "
+ "enumerate committed source records instead."
+ )
+ if not sources:
+ return 0
+ allowed = FilterCondition(
+ field=FilterableField.SYSTEM_METADATA_SYNC_ID,
+ operator=FilterOperator.IN,
+ value=[str(source.sync_id) for source in sources],
+ )
+ scoped = (
+ [FilterGroup(conditions=[*group.conditions, allowed]) for group in filter_groups]
+ if filter_groups
+ else [FilterGroup(conditions=[allowed])]
+ )
+ return await self._delegate.count(scoped, collection_id, name_substring)
+
+ async def close(self) -> None:
+ """The shared adapter is owned by the application, not this request."""
diff --git a/backend/airweave/domains/source_connections/create.py b/backend/airweave/domains/source_connections/create.py
index 63bc07acae..cacfbc9c90 100644
--- a/backend/airweave/domains/source_connections/create.py
+++ b/backend/airweave/domains/source_connections/create.py
@@ -2,6 +2,7 @@
import hashlib
import secrets
+from contextlib import nullcontext
from typing import Any, Optional
from uuid import UUID
@@ -22,6 +23,7 @@
from airweave.domains.connections.protocols import ConnectionRepositoryProtocol
from airweave.domains.credentials.protocols import IntegrationCredentialServiceProtocol
from airweave.domains.oauth.protocols import OAuthFlowServiceProtocol
+from airweave.domains.owned_provisioning.guard import require_provider_source
from airweave.domains.source_connections.protocols import (
ResponseBuilderProtocol,
SourceConnectionCreateServiceProtocol,
@@ -163,6 +165,25 @@ async def create(
)
return result
+ async def create_deferred(
+ self,
+ db: AsyncSession,
+ *,
+ obj_in: SourceConnectionCreate,
+ ctx: ApiContext,
+ uow: UnitOfWork,
+ ) -> SourceConnectionSchema:
+ """Persist managed source in caller transaction without Temporal side effects."""
+ if obj_in.sync_immediately:
+ raise ValueError("Deferred creation cannot start a job before native verification")
+ entry = self._get_source_entry(obj_in.short_name)
+ self._validate_auth_compatibility(
+ entry.source_class_ref, entry.short_name, AuthenticationMethod.AUTH_PROVIDER
+ )
+ return await self._create_with_auth_provider(
+ db, obj_in=obj_in, entry=entry, ctx=ctx, transaction=uow
+ )
+
async def reinitiate_oauth(
self,
db: AsyncSession,
@@ -178,6 +199,7 @@ async def reinitiate_oauth(
source_conn = await self._sc_repo.get(db, id=id, ctx=ctx)
if not source_conn:
raise NotFoundException("Source connection not found")
+ require_provider_source(source_conn.short_name)
if source_conn.is_authenticated and not await self._has_credential_error(db, source_conn):
raise HTTPException(
status_code=400,
@@ -388,7 +410,13 @@ async def _create_with_oauth_token(
)
async def _create_with_auth_provider(
- self, db: AsyncSession, *, obj_in: SourceConnectionCreate, entry, ctx: ApiContext
+ self,
+ db: AsyncSession,
+ *,
+ obj_in: SourceConnectionCreate,
+ entry,
+ ctx: ApiContext,
+ transaction: UnitOfWork | None = None,
) -> SourceConnectionSchema:
if not obj_in.authentication or not isinstance(
obj_in.authentication, AuthProviderAuthentication
@@ -430,7 +458,7 @@ async def _create_with_auth_provider(
connection_schema: Optional[schemas.Connection] = None
collection_schema: Optional[schemas.CollectionRecord] = None
- async with UnitOfWork(db) as uow:
+ async with nullcontext(transaction) if transaction else UnitOfWork(db) as uow:
collection = await self._get_collection(uow.session, obj_in.readable_collection_id, ctx)
collection_schema = schemas.CollectionRecord.model_validate(
collection, from_attributes=True
@@ -460,10 +488,12 @@ async def _create_with_auth_provider(
collection_id=collection.id,
collection_readable_id=collection.readable_id,
source_entry=entry,
+ source_config=validated_config,
schedule_config=obj_in.schedule,
run_immediately=bool(obj_in.sync_immediately),
ctx=ctx,
uow=uow,
+ **({"defer_execution": True} if transaction else {}),
)
await uow.session.flush()
@@ -485,11 +515,12 @@ async def _create_with_auth_provider(
uow=uow,
)
await uow.session.flush()
- await uow.commit()
+ if transaction is None:
+ await uow.commit()
await uow.session.refresh(source_conn)
response = await self._response_builder.build_response(db, source_conn, ctx)
- if sync_result and sync_result.sync_job and obj_in.sync_immediately:
+ if transaction is None and sync_result and sync_result.sync_job and obj_in.sync_immediately:
if connection_schema is None or collection_schema is None:
raise RuntimeError("Connection or collection schema not materialized")
await self._trigger_sync_workflow(
@@ -672,6 +703,7 @@ async def _create_authenticated_connection(
collection_id=collection.id,
collection_readable_id=collection.readable_id,
source_entry=entry,
+ source_config=validated_config,
schedule_config=obj_in.schedule,
run_immediately=bool(obj_in.sync_immediately),
ctx=ctx,
diff --git a/backend/airweave/domains/source_connections/delete.py b/backend/airweave/domains/source_connections/delete.py
index 8a25a14169..5e36e61501 100644
--- a/backend/airweave/domains/source_connections/delete.py
+++ b/backend/airweave/domains/source_connections/delete.py
@@ -8,6 +8,7 @@
from airweave.api.context import ApiContext
from airweave.core.exceptions import NotFoundException
from airweave.domains.collections.protocols import CollectionRepositoryProtocol
+from airweave.domains.owned_provisioning.guard import require_unmanaged_source
from airweave.domains.source_connections.protocols import (
ResponseBuilderProtocol,
SourceConnectionDeletionServiceProtocol,
@@ -49,6 +50,10 @@ async def delete(
if not source_conn:
raise NotFoundException("Source connection not found")
+ await require_unmanaged_source(
+ db, id, ctx.organization.id, short_name=source_conn.short_name
+ )
+
sync_id = source_conn.sync_id
collection_orm = await self._collection_repo.get_by_readable_id(
db, readable_id=source_conn.readable_collection_id, ctx=ctx
diff --git a/backend/airweave/domains/source_connections/protocols.py b/backend/airweave/domains/source_connections/protocols.py
index ae385cc697..a69071b7f9 100644
--- a/backend/airweave/domains/source_connections/protocols.py
+++ b/backend/airweave/domains/source_connections/protocols.py
@@ -188,6 +188,17 @@ async def create(
"""Create a source connection."""
...
+ async def create_deferred(
+ self,
+ db: AsyncSession,
+ *,
+ obj_in: SourceConnectionCreate,
+ ctx: ApiContext,
+ uow: UnitOfWork,
+ ) -> SourceConnectionSchema:
+ """Create a managed source in the caller transaction without execution."""
+ ...
+
async def reinitiate_oauth(
self,
db: AsyncSession,
diff --git a/backend/airweave/domains/source_connections/service.py b/backend/airweave/domains/source_connections/service.py
index 33cc05ef6f..987da12fef 100644
--- a/backend/airweave/domains/source_connections/service.py
+++ b/backend/airweave/domains/source_connections/service.py
@@ -16,6 +16,7 @@
from airweave.domains.collections.protocols import CollectionRepositoryProtocol
from airweave.domains.connections.protocols import ConnectionRepositoryProtocol
from airweave.domains.oauth.protocols import OAuthRedirectSessionRepositoryProtocol
+from airweave.domains.owned_provisioning.guard import require_provider_source
from airweave.domains.source_connections.protocols import (
ResponseBuilderProtocol,
SourceConnectionCreateServiceProtocol,
@@ -170,6 +171,7 @@ async def run(
) -> SourceConnectionJob:
"""Trigger a sync run for this source connection."""
source_conn = await self._resolve_source_connection(db, id, ctx)
+ require_provider_source(source_conn.short_name)
sync_id = source_conn.sync_id
assert sync_id is not None
diff --git a/backend/airweave/domains/source_connections/tests/test_create.py b/backend/airweave/domains/source_connections/tests/test_create.py
index b57f596f27..227b027c40 100644
--- a/backend/airweave/domains/source_connections/tests/test_create.py
+++ b/backend/airweave/domains/source_connections/tests/test_create.py
@@ -1117,3 +1117,13 @@ async def test_reinitiate_oauth_preserves_byoc_credentials(monkeypatch):
assert init_kwargs["template_configs"] == {"subdomain": "acme"}
# Verify payload was reused from old init session
assert init_kwargs["payload"]["name"] == "My GitHub"
+
+
+async def test_native_source_cannot_start_provider_oauth():
+ svc = _service(_entry(oauth_type="access_only"))
+ sc = _shell_source_conn()
+ sc.short_name = "almanac"
+ svc._sc_repo.seed(sc.id, sc)
+ with pytest.raises(HTTPException) as error:
+ await svc.reinitiate_oauth(AsyncMock(), id=sc.id, ctx=_ctx())
+ assert error.value.status_code == 409
diff --git a/backend/airweave/domains/source_connections/tests/test_delete.py b/backend/airweave/domains/source_connections/tests/test_delete.py
index a70bf44288..fc632ca091 100644
--- a/backend/airweave/domains/source_connections/tests/test_delete.py
+++ b/backend/airweave/domains/source_connections/tests/test_delete.py
@@ -11,6 +11,7 @@
from uuid import uuid4
import pytest
+from fastapi import HTTPException
from airweave.api.context import ApiContext
from airweave.core.exceptions import NotFoundException
@@ -32,6 +33,11 @@
COLLECTION_ID = uuid4()
+def _unmanaged_db():
+ """Legacy scenarios have no owned account relationship."""
+ return AsyncMock(scalar=AsyncMock(return_value=None))
+
+
def _make_ctx() -> ApiContext:
org = Organization(id=str(ORG_ID), name="Test Org", created_at=NOW, modified_at=NOW)
return ApiContext(
@@ -42,7 +48,9 @@ def _make_ctx() -> ApiContext:
)
-def _make_sc(*, id=None, sync_id=None, readable_collection_id="test-col", name="Test SC", short_name="github"):
+def _make_sc(
+ *, id=None, sync_id=None, readable_collection_id="test-col", name="Test SC", short_name="github"
+):
sc = MagicMock(spec=SourceConnection)
sc.id = id or uuid4()
sc.sync_id = sync_id
@@ -123,7 +131,7 @@ async def test_delete_happy_path(case: DeleteCase):
sync_service=sync_service,
)
- result = await svc.delete(AsyncMock(), id=sc.id, ctx=_make_ctx())
+ result = await svc.delete(_unmanaged_db(), id=sc.id, ctx=_make_ctx())
assert result.id == sc.id
assert sc_repo._store.get(sc.id) is None
@@ -149,8 +157,20 @@ class DeleteErrorCase:
DELETE_ERROR_CASES = [
- DeleteErrorCase("not_found", seed_sc=False, seed_collection=False, expect_exception=NotFoundException, expect_match="Source connection not found"),
- DeleteErrorCase("collection_not_found", seed_sc=True, seed_collection=False, expect_exception=NotFoundException, expect_match="Collection not found"),
+ DeleteErrorCase(
+ "not_found",
+ seed_sc=False,
+ seed_collection=False,
+ expect_exception=NotFoundException,
+ expect_match="Source connection not found",
+ ),
+ DeleteErrorCase(
+ "collection_not_found",
+ seed_sc=True,
+ seed_collection=False,
+ expect_exception=NotFoundException,
+ expect_match="Collection not found",
+ ),
]
@@ -168,7 +188,7 @@ async def test_delete_error(case: DeleteErrorCase):
svc = _build_service(sc_repo=sc_repo, collection_repo=col_repo)
with pytest.raises(case.expect_exception, match=case.expect_match):
- await svc.delete(AsyncMock(), id=sc.id, ctx=_make_ctx())
+ await svc.delete(_unmanaged_db(), id=sc.id, ctx=_make_ctx())
async def test_delete_sync_service_failure_propagates():
@@ -192,4 +212,17 @@ async def test_delete_sync_service_failure_propagates():
)
with pytest.raises(RuntimeError, match="sync delete boom"):
- await svc.delete(AsyncMock(), id=sc.id, ctx=_make_ctx())
+ await svc.delete(_unmanaged_db(), id=sc.id, ctx=_make_ctx())
+
+
+async def test_native_delete_rejected_before_cleanup():
+ sc = _make_sc(short_name="almanac")
+ repo = FakeSourceConnectionRepository()
+ repo.seed(sc.id, sc)
+ sync_service = AsyncMock()
+ with pytest.raises(HTTPException) as error:
+ await _build_service(sc_repo=repo, sync_service=sync_service).delete(
+ _unmanaged_db(), id=sc.id, ctx=_make_ctx()
+ )
+ assert error.value.status_code == 409
+ sync_service.delete.assert_not_awaited()
diff --git a/backend/airweave/domains/source_connections/tests/test_service.py b/backend/airweave/domains/source_connections/tests/test_service.py
index def0b11696..aaaa52aca3 100644
--- a/backend/airweave/domains/source_connections/tests/test_service.py
+++ b/backend/airweave/domains/source_connections/tests/test_service.py
@@ -13,6 +13,7 @@
from uuid import UUID, uuid4
import pytest
+from fastapi import HTTPException
from airweave.api.context import ApiContext
from airweave.core.datetime_utils import utc_now
@@ -689,3 +690,16 @@ async def test_count_by_organization():
svc = _build_run_service(sc_repo=sc_repo)
count = await svc.count_by_organization(AsyncMock(), organization_id=ORG_ID)
assert count == 0
+
+
+async def test_native_source_cannot_start_provider_pull_run():
+ sc = _make_source_conn()
+ sc.short_name = "almanac"
+ repo = FakeSourceConnectionRepository()
+ repo.seed(SC_ID, sc)
+ sync_service = _RecordingFakeSyncService()
+ svc = _build_run_service(sc_repo=repo, sync_service=sync_service)
+ with pytest.raises(HTTPException) as error:
+ await svc.run(AsyncMock(), id=SC_ID, ctx=_make_ctx(), force_full_sync=True)
+ assert error.value.status_code == 409
+ assert sync_service.last_trigger_run is None
diff --git a/backend/airweave/domains/source_connections/tests/test_update.py b/backend/airweave/domains/source_connections/tests/test_update.py
index 581cec880c..990840aa97 100644
--- a/backend/airweave/domains/source_connections/tests/test_update.py
+++ b/backend/airweave/domains/source_connections/tests/test_update.py
@@ -5,7 +5,7 @@
from dataclasses import dataclass, field
from datetime import datetime, timezone
-from typing import Any, Optional
+from typing import Optional
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
@@ -19,7 +19,6 @@
from airweave.core.exceptions import NotFoundException
from airweave.core.logging import logger
from airweave.core.shared_models import AuthMethod
-from airweave.domains.syncs.types import InvalidSyncTransitionError
from airweave.domains.collections.fakes.repository import FakeCollectionRepository
from airweave.domains.connections.fakes.repository import FakeConnectionRepository
from airweave.domains.credentials.fakes.repository import (
@@ -34,8 +33,8 @@
from airweave.domains.sources.fakes.service import FakeSourceService
from airweave.domains.sources.fakes.validation import FakeSourceValidationService
from airweave.domains.sources.types import SourceRegistryEntry
-from airweave.domains.syncs.fakes.service import FakeSyncService
from airweave.domains.syncs.fakes.repository import FakeSyncRepository
+from airweave.domains.syncs.fakes.service import FakeSyncService
from airweave.domains.syncs.types import SyncProvisionResult
from airweave.domains.temporal.fakes.schedule_service import FakeTemporalScheduleService
from airweave.models.collection import Collection
@@ -54,6 +53,11 @@ class _AuthPayload(BaseModel):
token: str
+def _unmanaged_db():
+ """Legacy scenarios have no owned account relationship."""
+ return AsyncMock(scalar=AsyncMock(return_value=None))
+
+
def _make_ctx() -> ApiContext:
org = Organization(id=str(ORG_ID), name="Test Org", created_at=NOW, modified_at=NOW)
return ApiContext(
@@ -141,7 +145,7 @@ async def test_update_not_found():
svc = _build_service()
obj_in = SourceConnectionUpdate(name="New Name")
with pytest.raises(NotFoundException, match="Source connection not found"):
- await svc.update(AsyncMock(), id=uuid4(), obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=uuid4(), obj_in=obj_in, ctx=_make_ctx())
# ---------------------------------------------------------------------------
@@ -159,7 +163,9 @@ class FieldUpdateCase:
FIELD_UPDATE_CASES = [
FieldUpdateCase("name_only", SourceConnectionUpdate(name="Renamed")),
FieldUpdateCase("description_only", SourceConnectionUpdate(description="New desc")),
- FieldUpdateCase("name_and_description", SourceConnectionUpdate(name="Renamed", description="Desc")),
+ FieldUpdateCase(
+ "name_and_description", SourceConnectionUpdate(name="Renamed", description="Desc")
+ ),
]
@@ -170,7 +176,7 @@ async def test_field_update(case: FieldUpdateCase):
sc_repo.seed(sc.id, sc)
svc = _build_service(sc_repo=sc_repo)
- result = await svc.update(AsyncMock(), id=sc.id, obj_in=case.obj_in, ctx=_make_ctx())
+ result = await svc.update(_unmanaged_db(), id=sc.id, obj_in=case.obj_in, ctx=_make_ctx())
assert result.id == sc.id
@@ -190,11 +196,84 @@ async def test_update_config_valid():
svc = _build_service(sc_repo=sc_repo, source_validation=validation)
obj_in = SourceConnectionUpdate(config={"key": "value"})
- result = await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ result = await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert result.id == sc.id
assert any(c[0] == "validate_config" for c in validation._calls)
+@pytest.mark.parametrize("provider", ["outlook_mail", "outlook_calendar"])
+@pytest.mark.parametrize("originals", [False, True])
+async def test_outlook_capture_mode_change_requires_new_source(originals: bool, provider: str):
+ sc = _make_sc(short_name=provider)
+ sc.config_fields = {"expected_principal_id": "native-user", "capture_originals": originals}
+ sc_repo = FakeSourceConnectionRepository()
+ sc_repo.seed(sc.id, sc)
+ config = {"expected_principal_id": "native-user", "capture_originals": not originals}
+ validation = FakeSourceValidationService()
+ validation.seed_config_result(provider, config)
+ svc = _build_service(sc_repo=sc_repo, source_validation=validation)
+
+ with pytest.raises(HTTPException, match="Create a new source") as exc:
+ await svc.update(
+ _unmanaged_db(), id=sc.id, obj_in=SourceConnectionUpdate(config=config), ctx=_make_ctx()
+ )
+
+ assert exc.value.status_code == 400
+ assert sc.config_fields["capture_originals"] is originals
+ assert not any(call[0] == "update" for call in sc_repo._calls)
+
+
+@pytest.mark.parametrize("provider", ["outlook_mail", "outlook_calendar"])
+@pytest.mark.parametrize("originals", [False, True])
+async def test_outlook_principal_change_preserves_capture_identity(originals: bool, provider: str):
+ sc = _make_sc(short_name=provider)
+ sc.config_fields = {"expected_principal_id": "native-user", "capture_originals": originals}
+ sc_repo = FakeSourceConnectionRepository()
+ sc_repo.seed(sc.id, sc)
+ config = {**sc.config_fields, "expected_principal_id": "another-user"}
+ validation = FakeSourceValidationService()
+ validation.seed_config_result(provider, config)
+ svc = _build_service(sc_repo=sc_repo, source_validation=validation)
+
+ if originals:
+ with pytest.raises(HTTPException, match="Create a new source") as exc:
+ await svc.update(
+ _unmanaged_db(),
+ id=sc.id,
+ obj_in=SourceConnectionUpdate(config=config),
+ ctx=_make_ctx(),
+ )
+ assert exc.value.status_code == 400
+ assert sc.config_fields["expected_principal_id"] == "native-user"
+ assert not any(call[0] == "update" for call in sc_repo._calls)
+ else:
+ await svc.update(
+ _unmanaged_db(), id=sc.id, obj_in=SourceConnectionUpdate(config=config), ctx=_make_ctx()
+ )
+ assert sc.config_fields["expected_principal_id"] == "another-user"
+
+
+@pytest.mark.parametrize("originals", [None, False, True])
+async def test_outlook_unchanged_capture_mode_allows_config_update(originals: bool | None):
+ sc = _make_sc(short_name="outlook_mail")
+ sc.config_fields = {"expected_principal_id": "native-user"}
+ if originals is not None:
+ sc.config_fields["capture_originals"] = originals
+ sc_repo = FakeSourceConnectionRepository()
+ sc_repo.seed(sc.id, sc)
+ config = {**sc.config_fields, "included_folders": ["inbox"]}
+ validation = FakeSourceValidationService()
+ validation.seed_config_result("outlook_mail", config)
+ svc = _build_service(sc_repo=sc_repo, source_validation=validation)
+
+ result = await svc.update(
+ _unmanaged_db(), id=sc.id, obj_in=SourceConnectionUpdate(config=config), ctx=_make_ctx()
+ )
+
+ assert result.id == sc.id
+ assert sc.config_fields == config
+
+
# ---------------------------------------------------------------------------
# Schedule updates -- table-driven
# ---------------------------------------------------------------------------
@@ -212,10 +291,36 @@ class ScheduleCase:
SCHEDULE_CASES = [
- ScheduleCase("update_existing", has_sync=True, new_cron="0 * * * *", expect_temporal_create=True, expect_temporal_delete=False),
- ScheduleCase("remove_schedule", has_sync=True, new_cron=None, expect_temporal_create=False, expect_temporal_delete=True),
- ScheduleCase("add_no_sync", has_sync=False, new_cron="0 * * * *", expect_temporal_create=False, expect_temporal_delete=False, expect_sync_record_create=True),
- ScheduleCase("no_connection_id_warning", has_sync=False, new_cron="0 * * * *", has_connection_id=False, expect_temporal_create=False, expect_temporal_delete=False),
+ ScheduleCase(
+ "update_existing",
+ has_sync=True,
+ new_cron="0 * * * *",
+ expect_temporal_create=True,
+ expect_temporal_delete=False,
+ ),
+ ScheduleCase(
+ "remove_schedule",
+ has_sync=True,
+ new_cron=None,
+ expect_temporal_create=False,
+ expect_temporal_delete=True,
+ ),
+ ScheduleCase(
+ "add_no_sync",
+ has_sync=False,
+ new_cron="0 * * * *",
+ expect_temporal_create=False,
+ expect_temporal_delete=False,
+ expect_sync_record_create=True,
+ ),
+ ScheduleCase(
+ "no_connection_id_warning",
+ has_sync=False,
+ new_cron="0 * * * *",
+ has_connection_id=False,
+ expect_temporal_create=False,
+ expect_temporal_delete=False,
+ ),
]
@@ -288,7 +393,7 @@ async def test_schedule_update(case: ScheduleCase):
collection_repo=col_repo,
)
- result = await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ result = await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert result.id == sc.id
if case.expect_temporal_create:
@@ -311,7 +416,7 @@ async def test_schedule_add_collection_not_found():
obj_in = SourceConnectionUpdate(schedule={"cron": "0 * * * *"})
with pytest.raises(NotFoundException, match="Collection not found"):
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
async def test_schedule_add_rejects_federated_source():
@@ -348,7 +453,7 @@ async def test_schedule_add_rejects_federated_source():
obj_in = SourceConnectionUpdate(schedule={"cron": "0 * * * *"})
with pytest.raises(HTTPException) as exc_info:
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert exc_info.value.status_code == 400
assert "federated search" in str(exc_info.value.detail)
@@ -415,13 +520,13 @@ async def test_credential_update(case: CredentialCase):
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
if case.expect_success:
- result = await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ result = await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert result.id == sc.id
assert len(encryptor._encrypt_calls) == 1
assert any(c[0] == "update" for c in cred_repo._calls)
else:
with pytest.raises(HTTPException) as exc_info:
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert exc_info.value.status_code == 400
assert "direct authentication" in str(exc_info.value.detail)
@@ -441,7 +546,7 @@ async def test_credential_update_direct_auth_missing_connection_raises():
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
with pytest.raises(NotFoundException, match="Connection not found"):
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
async def test_credential_update_direct_auth_missing_integration_credential_raises():
@@ -469,7 +574,7 @@ async def test_credential_update_direct_auth_missing_integration_credential_rais
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
with pytest.raises(NotFoundException, match="Integration credential not configured"):
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
async def test_credential_update_direct_auth_missing_credential_record_raises():
@@ -498,7 +603,7 @@ async def test_credential_update_direct_auth_missing_credential_record_raises():
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
with pytest.raises(NotFoundException, match="Integration credential not found"):
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
# ---------------------------------------------------------------------------
@@ -579,7 +684,7 @@ async def test_credential_update_triggers_unpause():
)
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
- await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert any(c[0] == "resume" for c in sync_svc._calls)
@@ -622,5 +727,53 @@ async def test_credential_update_unpause_failure_is_nonfatal():
)
obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}})
- result = await svc.update(AsyncMock(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
+ result = await svc.update(_unmanaged_db(), id=sc.id, obj_in=obj_in, ctx=_make_ctx())
assert result.id == sc.id
+
+
+async def test_native_update_cannot_rebind_identity():
+ sc = _make_sc(short_name="almanac")
+ repo = FakeSourceConnectionRepository()
+ repo.seed(sc.id, sc)
+ db = _unmanaged_db()
+ with pytest.raises(HTTPException) as error:
+ await _build_service(sc_repo=repo).update(
+ db,
+ id=sc.id,
+ obj_in=SourceConnectionUpdate(config={"owner_id": "other"}),
+ ctx=_make_ctx(),
+ )
+ assert error.value.status_code == 409
+ db.commit.assert_not_awaited()
+ db.scalar.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+ "change",
+ [
+ None,
+ {"livemode": True},
+ {"expected_account_id": "acct_other"},
+ {"api_version": "2026-01-28.clover"},
+ ],
+)
+async def test_stripe_capture_binding_cannot_change_in_place(change):
+ binding = {
+ "expected_account_id": "acct_selected",
+ "livemode": False,
+ "api_version": "2025-06-30.basil",
+ }
+ sc = _make_sc(short_name="stripe")
+ sc.config_fields = {"original_capture": binding}
+ sc_repo = FakeSourceConnectionRepository()
+ sc_repo.seed(sc.id, sc)
+ config = {"original_capture": None if change is None else {**binding, **change}}
+ validation = FakeSourceValidationService()
+ validation.seed_config_result("stripe", config)
+ svc = _build_service(sc_repo=sc_repo, source_validation=validation)
+ with pytest.raises(HTTPException, match="Create a new source"):
+ await svc.update(
+ _unmanaged_db(), id=sc.id, obj_in=SourceConnectionUpdate(config=config), ctx=_make_ctx()
+ )
+ assert sc.config_fields == {"original_capture": binding}
+ assert not any(call[0] == "update" for call in sc_repo._calls)
diff --git a/backend/airweave/domains/source_connections/update.py b/backend/airweave/domains/source_connections/update.py
index 670220f2ab..f1b199037a 100644
--- a/backend/airweave/domains/source_connections/update.py
+++ b/backend/airweave/domains/source_connections/update.py
@@ -15,6 +15,7 @@
from airweave.domains.collections.protocols import CollectionRepositoryProtocol
from airweave.domains.connections.protocols import ConnectionRepositoryProtocol
from airweave.domains.credentials.protocols import IntegrationCredentialRepositoryProtocol
+from airweave.domains.owned_provisioning.guard import require_unmanaged_source
from airweave.domains.source_connections.protocols import (
ResponseBuilderProtocol,
SourceConnectionRepositoryProtocol,
@@ -30,6 +31,7 @@
from airweave.domains.syncs.types import InvalidSyncTransitionError, OptimisticLockError
from airweave.domains.temporal.protocols import TemporalScheduleServiceProtocol
from airweave.models.source_connection import SourceConnection
+from airweave.platform.configs.config import OutlookCalendarConfig, OutlookMailConfig, StripeConfig
from airweave.schemas.source_connection import (
AuthenticationMethod,
ScheduleConfig,
@@ -40,6 +42,43 @@
)
+def _validate_outlook_capture_mode(source: SourceConnection, config: dict[str, Any]) -> None:
+ """Mutable legacy IDs and immutable originals cannot share an existing sync."""
+ if source.short_name not in {"outlook_mail", "outlook_calendar"}:
+ return
+ config_type = (
+ OutlookMailConfig if source.short_name == "outlook_mail" else OutlookCalendarConfig
+ )
+ previous = config_type.model_validate(source.config_fields or {})
+ proposed = config_type.model_validate(config)
+ if previous.capture_originals != proposed.capture_originals:
+ raise HTTPException(
+ status_code=400,
+ detail="Create a new source to change Outlook original capture mode",
+ )
+ if (
+ previous.capture_originals
+ and previous.expected_principal_id != proposed.expected_principal_id
+ ):
+ raise HTTPException(
+ status_code=400,
+ detail="Create a new source to change the Outlook original capture principal",
+ )
+
+
+def _validate_stripe_capture_binding(source: SourceConnection, config: dict[str, Any]) -> None:
+ """Account, API version, mode and capture representation require fresh identities."""
+ if source.short_name != "stripe":
+ return
+ previous = StripeConfig.model_validate(source.config_fields or {})
+ proposed = StripeConfig.model_validate(config)
+ if previous.original_capture != proposed.original_capture:
+ raise HTTPException(
+ status_code=400,
+ detail="Create a new source to change Stripe original capture binding",
+ )
+
+
class SourceConnectionUpdateService(SourceConnectionUpdateServiceProtocol):
"""Updates a source connection.
@@ -93,6 +132,10 @@ async def update(
if not source_conn:
raise NotFoundException("Source connection not found")
+ await require_unmanaged_source(
+ db, id, ctx.organization.id, short_name=source_conn.short_name
+ )
+
# Update fields
update_data = obj_in.model_dump(exclude_unset=True)
@@ -110,6 +153,8 @@ async def update(
validated_config = self._source_validation.validate_config(
source_conn.short_name, update_data["config"], ctx
)
+ _validate_outlook_capture_mode(source_conn, validated_config)
+ _validate_stripe_capture_binding(source_conn, validated_config)
update_data["config_fields"] = validated_config
del update_data["config"]
@@ -233,6 +278,7 @@ async def _handle_schedule_update(
collection_id=collection.id,
collection_readable_id=collection.readable_id,
source_entry=source_entry,
+ source_config=source_conn.config_fields,
schedule_config=ScheduleConfig(cron=new_cron),
run_immediately=False,
ctx=ctx,
diff --git a/backend/airweave/domains/sources/exceptions/__init__.py b/backend/airweave/domains/sources/exceptions/__init__.py
index 04469e0c85..099f9c676f 100644
--- a/backend/airweave/domains/sources/exceptions/__init__.py
+++ b/backend/airweave/domains/sources/exceptions/__init__.py
@@ -13,6 +13,7 @@
SourceEntitySkippedError,
SourceError,
SourceFileDownloadError,
+ SourceGoneError,
SourceNotFoundError,
SourcePermanentError,
SourceRateLimitError,
diff --git a/backend/airweave/domains/sources/exceptions/_exceptions.py b/backend/airweave/domains/sources/exceptions/_exceptions.py
index 9c30308351..d9d0c5b3f8 100644
--- a/backend/airweave/domains/sources/exceptions/_exceptions.py
+++ b/backend/airweave/domains/sources/exceptions/_exceptions.py
@@ -83,6 +83,10 @@ def __init__(self, message: str, *, source_short_name: str = ""):
# ---------------------------------------------------------------------------
+class SourceGoneError(SourceError):
+ """HTTP 410; the source decides whether a resource or continuation expired."""
+
+
class SourceCreationError(SourceError):
"""source_class.create() failed (bad credentials, missing fields, etc.)."""
diff --git a/backend/airweave/domains/sources/exceptions/classifier.py b/backend/airweave/domains/sources/exceptions/classifier.py
index a9bb853110..ed7da011ec 100644
--- a/backend/airweave/domains/sources/exceptions/classifier.py
+++ b/backend/airweave/domains/sources/exceptions/classifier.py
@@ -18,6 +18,8 @@
AuthProviderAccountNotFoundError,
AuthProviderAuthError,
AuthProviderError,
+ AuthProviderRateLimitError,
+ AuthProviderServerError,
)
from airweave.domains.source_connections.types import ErrorClassification
from airweave.domains.sources.exceptions import (
@@ -78,6 +80,13 @@ def classify_error(exc: Exception) -> ErrorClassification: # noqa: C901
message=str(exc),
)
+ if isinstance(exc, AuthProviderRateLimitError):
+ return ErrorClassification(
+ category=SourceConnectionErrorCategory.RATE_LIMITED, message=str(exc)
+ )
+ if isinstance(exc, AuthProviderServerError):
+ return ErrorClassification(category=None, message=None)
+
# Catch-all for remaining AuthProviderError subtypes (e.g. MissingFieldsError,
# ConfigError) — these are auth provider issues the user needs to address.
if isinstance(exc, AuthProviderError):
diff --git a/backend/airweave/domains/sources/lifecycle.py b/backend/airweave/domains/sources/lifecycle.py
index cc568f47db..f548d902d0 100644
--- a/backend/airweave/domains/sources/lifecycle.py
+++ b/backend/airweave/domains/sources/lifecycle.py
@@ -54,6 +54,10 @@
TokenProviderAccountGoneError,
)
from airweave.domains.sources.token_providers.oauth import OAuthTokenProvider
+from airweave.domains.sources.token_providers.protocol import (
+ ManagedAuthProvider,
+ ManagedToolAuthProvider,
+)
from airweave.domains.sources.token_providers.static import StaticTokenProvider
from airweave.domains.sources.types import AuthConfig, SourceConnectionData, SourceRegistryEntry
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
@@ -159,40 +163,45 @@ async def create(
source_connection_id=source_connection_data.source_connection_id,
ctx=ctx,
logger=logger,
+ managed_auth=auth_config.managed_auth,
)
- # 5. Parse config_fields into typed config
- entry = self._source_registry.get(source_connection_data.short_name)
- config = self._build_typed_config(entry, source_connection_data.config_fields)
-
- # 6. Create source instance with all deps injected
- source = await source_connection_data.source_class.create(
- auth=token_provider,
- logger=logger,
- http_client=http_client,
- config=config,
- )
-
- # 7. Validate credentials early so failures surface as NEEDS_REAUTH.
- # Only catch auth-related exceptions — transient errors (server 5xx,
- # rate limits, network timeouts) must propagate so they don't
- # incorrectly pause schedules or set NEEDS_REAUTH status.
try:
- await source.validate()
- except (
- SourceAuthError,
- AuthProviderAuthError,
- AuthProviderAccountNotFoundError,
- TokenCredentialsInvalidError,
- TokenExpiredError,
- TokenProviderAccountGoneError,
- ) as exc:
- raise SourceValidationError(
- short_name=source_connection_data.short_name,
- reason=f"credential validation failed: {exc}",
- ) from exc
+ # 5. Parse config_fields into typed config
+ entry = self._source_registry.get(source_connection_data.short_name)
+ config = self._build_typed_config(entry, source_connection_data.config_fields)
- return source
+ # 6. Create source instance with all deps injected
+ source = await source_connection_data.source_class.create(
+ auth=token_provider,
+ logger=logger,
+ http_client=http_client,
+ config=config,
+ )
+
+ # 7. Validate credentials early so failures surface as NEEDS_REAUTH.
+ # Only catch auth-related exceptions — transient errors (server 5xx,
+ # rate limits, network timeouts) must propagate so they don't
+ # incorrectly pause schedules or set NEEDS_REAUTH status.
+ try:
+ await source.validate()
+ except (
+ SourceAuthError,
+ AuthProviderAuthError,
+ AuthProviderAccountNotFoundError,
+ TokenCredentialsInvalidError,
+ TokenExpiredError,
+ TokenProviderAccountGoneError,
+ ) as exc:
+ raise SourceValidationError(
+ short_name=source_connection_data.short_name,
+ reason=f"credential validation failed: {exc}",
+ ) from exc
+
+ return source
+ except BaseException:
+ await http_client.aclose()
+ raise
async def validate(
self,
@@ -398,6 +407,7 @@ async def _get_auth_provider_configuration(
return AuthConfig(
credentials=auth_result.credentials,
+ managed_auth=auth_result.managed_auth,
auth_provider_instance=auth_provider_instance,
)
@@ -617,7 +627,7 @@ def _process_creds_dict(
async def _resolve_token_provider(
self,
source_connection_data: SourceConnectionData,
- source_credentials: SourceCredentials,
+ source_credentials: SourceCredentials | None,
ctx: ApiContext,
logger: ContextualLogger,
access_token: Optional[str],
@@ -634,6 +644,9 @@ async def _resolve_token_provider(
auth_provider_instance: Optional[BaseAuthProvider] = auth_config.auth_provider_instance
short_name = source_connection_data.short_name
+ if auth_config.managed_auth is not None:
+ return auth_config.managed_auth
+
if access_token is not None:
return StaticTokenProvider(access_token, source_short_name=short_name)
@@ -708,12 +721,24 @@ def _build_http_client(
source_connection_id: UUID,
ctx: ApiContext,
logger: ContextualLogger,
+ managed_auth: ManagedAuthProvider | ManagedToolAuthProvider | None = None,
) -> AirweaveHttpClient:
"""Build an AirweaveHttpClient with rate limiting for this source."""
feature_enabled = ctx.has_feature(FeatureFlag.SOURCE_RATE_LIMITING)
+ from airweave.platform.http_client.composio_transport import ComposioTransport
+
+ transport = (
+ ComposioTransport(
+ api_key=managed_auth.api_key.get_secret_value(),
+ connected_account_id=managed_auth.connected_account_id,
+ allowed_hosts=managed_auth.allowed_hosts,
+ )
+ if isinstance(managed_auth, ManagedAuthProvider)
+ else None
+ )
client = AirweaveHttpClient(
- wrapped_client=httpx.AsyncClient(),
+ wrapped_client=httpx.AsyncClient(transport=transport),
org_id=ctx.organization.id,
source_short_name=source_short_name,
rate_limiter=self._rate_limiter,
diff --git a/backend/airweave/domains/sources/tests/test_validation.py b/backend/airweave/domains/sources/tests/test_validation.py
index 85979a2cdd..fb7c9f8f18 100644
--- a/backend/airweave/domains/sources/tests/test_validation.py
+++ b/backend/airweave/domains/sources/tests/test_validation.py
@@ -65,3 +65,31 @@ def test_validate_config_passthrough_when_no_schema():
svc = SourceValidationService(source_registry=registry)
payload: dict[str, Any] = {"x": 1}
assert svc.validate_config("slack", payload, _make_ctx()) == payload
+
+
+def test_calendar_config_is_json_serializable_and_round_trips():
+ """Validated nested dates must survive source connection JSON persistence."""
+ import json
+
+ from airweave.platform.configs.config import GoogleCalendarConfig
+
+ registry = FakeSourceRegistry()
+ registry.seed(
+ _make_entry("google_calendar", "Google Calendar").model_copy(
+ update={"config_ref": GoogleCalendarConfig}
+ )
+ )
+ config = {
+ "calendar_ids": ["owner@example.test"],
+ "occurrence_window": {
+ "start": "2026-10-01T00:00:00+05:30",
+ "end": "2026-10-08T00:00:00+05:30",
+ },
+ }
+ validated = SourceValidationService(registry).validate_config(
+ "google_calendar", config, _make_ctx()
+ )
+ stored = json.loads(json.dumps(validated))
+ assert GoogleCalendarConfig.model_validate(stored) == GoogleCalendarConfig.model_validate(
+ config
+ )
diff --git a/backend/airweave/domains/sources/token_providers/protocol.py b/backend/airweave/domains/sources/token_providers/protocol.py
index 26252ed78c..b813756bbc 100644
--- a/backend/airweave/domains/sources/token_providers/protocol.py
+++ b/backend/airweave/domains/sources/token_providers/protocol.py
@@ -16,7 +16,9 @@
"""
from enum import Enum
-from typing import Protocol, runtime_checkable
+from typing import Literal, Protocol, runtime_checkable
+
+from pydantic import BaseModel, ConfigDict, Field, SecretStr
class AuthProviderKind(str, Enum):
@@ -81,3 +83,39 @@ async def force_refresh(self) -> str:
TokenProviderServerError: If upstream fails during refresh.
"""
...
+
+
+class ManagedAuthProvider(BaseModel):
+ """Bound credential custody; deliberately does not expose a bearer token."""
+
+ model_config = ConfigDict(frozen=True, extra="forbid")
+ api_key: SecretStr = Field(repr=False)
+ connected_account_id: str = Field(min_length=1)
+ allowed_hosts: frozenset[str] = Field(min_length=1)
+ provider_kind: Literal[AuthProviderKind.AUTH_PROVIDER] = AuthProviderKind.AUTH_PROVIDER
+ supports_refresh: Literal[False] = False
+
+
+class ManagedToolAuthProvider(BaseModel):
+ """Explicit account-bound tool sessions for a provider without a REST proxy."""
+
+ model_config = ConfigDict(frozen=True, extra="forbid")
+ api_key: SecretStr = Field(repr=False)
+ connected_account_id: str = Field(min_length=1)
+ user_id: str = Field(min_length=1)
+ provider_kind: Literal[AuthProviderKind.AUTH_PROVIDER] = AuthProviderKind.AUTH_PROVIDER
+ supports_refresh: Literal[False] = False
+
+
+async def authorization_headers(
+ auth: SourceAuthProvider, *, refresh: bool = False
+) -> dict[str, str]:
+ """Resolve direct credentials or explicitly defer authentication to transport."""
+ if isinstance(auth, ManagedAuthProvider):
+ if refresh:
+ raise ValueError("Managed authentication refresh belongs to the auth provider")
+ return {}
+ if not isinstance(auth, TokenProviderProtocol):
+ raise TypeError("This request requires token or managed authentication")
+ token = await auth.force_refresh() if refresh else await auth.get_token()
+ return {"Authorization": f"Bearer {token}"}
diff --git a/backend/airweave/domains/sources/types.py b/backend/airweave/domains/sources/types.py
index d4013500b9..8f53d8947d 100644
--- a/backend/airweave/domains/sources/types.py
+++ b/backend/airweave/domains/sources/types.py
@@ -9,6 +9,10 @@
from airweave.core.protocols.registry import BaseRegistryEntry
from airweave.domains.auth_provider._base import BaseAuthProvider
from airweave.domains.credentials.types import DecryptedCredential
+from airweave.domains.sources.token_providers.protocol import (
+ ManagedAuthProvider,
+ ManagedToolAuthProvider,
+)
from airweave.models.connection import Connection
from airweave.models.source_connection import SourceConnection
from airweave.platform.configs._base import BaseConfig, Fields
@@ -91,3 +95,4 @@ class AuthConfig:
credentials: Any
auth_provider_instance: Optional[BaseAuthProvider]
decrypted_credential: Optional[DecryptedCredential] = field(default=None)
+ managed_auth: ManagedAuthProvider | ManagedToolAuthProvider | None = None
diff --git a/backend/airweave/domains/sources/validation.py b/backend/airweave/domains/sources/validation.py
index 25edf4ecb3..fef969b934 100644
--- a/backend/airweave/domains/sources/validation.py
+++ b/backend/airweave/domains/sources/validation.py
@@ -31,7 +31,7 @@ def __init__(self, source_registry: SourceRegistryProtocol) -> None:
def validate_config(
self, short_name: str, config_fields: Mapping[str, Any] | None, ctx: ApiContext
) -> dict[str, Any]:
- """Validate configuration fields against source schema, returning a plain dict.
+ """Validate source configuration and return JSON-compatible persistence fields.
Also strips fields that have feature flags not enabled for the organization.
"""
@@ -48,7 +48,7 @@ def validate_config(
try:
model = config_class.model_validate(payload)
- return model.model_dump()
+ return model.model_dump(mode="json")
except ValidationError as e:
errors = "; ".join([f"{self._loc(err)}: {err.get('msg')}" for err in e.errors()])
raise HTTPException(status_code=422, detail=f"Invalid config fields: {errors}") from e
diff --git a/backend/airweave/domains/storage/exceptions.py b/backend/airweave/domains/storage/exceptions.py
index dfffed4272..5fe09cac2c 100644
--- a/backend/airweave/domains/storage/exceptions.py
+++ b/backend/airweave/domains/storage/exceptions.py
@@ -55,3 +55,7 @@ def __init__(self, reason: str, filename: str):
self.reason = reason
self.filename = filename
super().__init__(f"File '{filename}' skipped: {reason}")
+
+
+class StorageReadLimitExceeded(StorageException):
+ """Actual stored object exceeds the explicitly requested read bound."""
diff --git a/backend/airweave/domains/storage/file_service.py b/backend/airweave/domains/storage/file_service.py
index 945ebf40d6..5500d0974c 100644
--- a/backend/airweave/domains/storage/file_service.py
+++ b/backend/airweave/domains/storage/file_service.py
@@ -10,6 +10,7 @@
Raw httpx errors propagate to the caller.
"""
+import hashlib
import os
import shutil
from typing import Optional
@@ -20,14 +21,20 @@
from tenacity import retry, stop_after_attempt
from airweave.core.logging import ContextualLogger
-from airweave.domains.sources.token_providers.protocol import SourceAuthProvider
+from airweave.domains.entities.canonical.requests import BlobReference
+from airweave.domains.sources.token_providers.protocol import (
+ ManagedAuthProvider,
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES
from airweave.domains.storage.paths import paths
from airweave.domains.storage.protocols import StorageBackend
from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS
from airweave.platform.entities._base import FileEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
@@ -37,19 +44,87 @@
class FileService:
"""Unified file service for downloading and restoring files."""
- MAX_FILE_SIZE_BYTES = 209715200
+ MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_BYTES
def __init__(
self,
sync_job_id: UUID,
storage_backend: StorageBackend,
+ *,
+ sync_id: UUID | None = None,
) -> None:
"""Initialize file service."""
self.sync_job_id = sync_job_id
+ self.sync_id = sync_id
self.storage = storage_backend
self.base_temp_dir = paths.temp_sync_dir(sync_job_id)
self._ensure_base_dir()
+ async def store_canonical_blob(
+ self, content: bytes, *, media_type: str | None = None
+ ) -> BlobReference:
+ """Persist immutable source bytes before publishing their record reference."""
+ if self.sync_id is None:
+ raise ValueError("Canonical blob storage requires a sync identity")
+ if len(content) > self.MAX_FILE_SIZE_BYTES:
+ raise ValueError("Canonical blob exceeds the file size limit")
+ digest = hashlib.sha256(content).hexdigest()
+ key = f"canonical/{self.sync_id}/blobs/sha256/{digest}"
+ await self.storage.write_file(key, content)
+ return BlobReference(key=key, sha256=digest, size_bytes=len(content), media_type=media_type)
+
+ async def capture_canonical_url(
+ self,
+ url: str,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ logger: ContextualLogger,
+ media_type: str | None = None,
+ *,
+ follow_redirects: bool = True,
+ expected_media_type: str | None = None,
+ ) -> BlobReference:
+ """Capture original bytes with size/auth checks, independent of search file types."""
+ if self.sync_id is None:
+ raise ValueError("Canonical blob storage requires a sync identity")
+ validate_url(url)
+ headers = await self._resolve_headers(auth, url)
+ temp_path = f"{self.base_temp_dir}/{uuid4()}-canonical"
+ try:
+ try:
+ await self._stream_download(
+ client,
+ url,
+ headers,
+ temp_path,
+ logger,
+ follow_redirects=follow_redirects,
+ expected_media_type=expected_media_type,
+ )
+ except httpx.HTTPStatusError as exc:
+ if (
+ exc.response.status_code != 401
+ or not headers.get("Authorization")
+ or not auth.supports_refresh
+ ):
+ raise
+ # Never promote an unauthenticated signed URL to provider credentials.
+ headers = await authorization_headers(auth, refresh=True)
+ await self._stream_download(
+ client,
+ url,
+ headers,
+ temp_path,
+ logger,
+ follow_redirects=follow_redirects,
+ expected_media_type=expected_media_type,
+ )
+ async with aiofiles.open(temp_path, "rb") as downloaded:
+ content = await downloaded.read(self.MAX_FILE_SIZE_BYTES + 1)
+ return await self.store_canonical_blob(content, media_type=media_type)
+ finally:
+ self._cleanup_temp(temp_path)
+
def _ensure_base_dir(self) -> None:
"""Ensure temp directory exists."""
os.makedirs(self.base_temp_dir, exist_ok=True)
@@ -61,12 +136,11 @@ def _ensure_base_dir(self) -> None:
@staticmethod
async def _resolve_headers(auth: SourceAuthProvider, url: str) -> dict:
"""Build auth headers. Pre-signed URLs skip the bearer token."""
+ if isinstance(auth, ManagedAuthProvider):
+ return {}
if "X-Amz-Algorithm" in url or "tempauth=" in url:
return {}
- token = await auth.get_token() if hasattr(auth, "get_token") else None
- if not token:
- raise ValueError(f"No access token available for downloading {url}")
- return {"Authorization": f"Bearer {token}"}
+ return await authorization_headers(auth)
# =========================================================================
# URL Download
@@ -115,16 +189,31 @@ async def _stream_download(
headers: dict,
temp_path: str,
logger: ContextualLogger,
+ *,
+ follow_redirects: bool = True,
+ expected_media_type: str | None = None,
) -> None:
"""Stream-download a file to disk with retry on 429/5xx/timeout."""
async with client.stream(
"GET",
url,
headers=headers,
- follow_redirects=True,
+ follow_redirects=follow_redirects,
timeout=httpx.Timeout(180.0, read=540.0),
) as response:
response.raise_for_status()
+ if expected_media_type is not None:
+ actual_type = (
+ response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower()
+ )
+ expected_type = expected_media_type.split(";", 1)[0].strip().lower()
+ if not actual_type or actual_type not in {
+ expected_type,
+ "application/octet-stream",
+ }:
+ raise ValueError(
+ "Downloaded file content type does not match provider metadata"
+ )
content_length = response.headers.get("Content-Length")
if content_length and int(content_length) > self.MAX_FILE_SIZE_BYTES:
@@ -188,10 +277,13 @@ async def download_from_url(
self._cleanup_temp(temp_path)
raise
except httpx.HTTPStatusError as first_error:
- if first_error.response.status_code == 401 and auth.supports_refresh:
+ if (
+ first_error.response.status_code == 401
+ and headers.get("Authorization")
+ and auth.supports_refresh
+ ):
logger.info("Download got 401, refreshing token and retrying")
- new_token = await auth.force_refresh()
- headers = {"Authorization": f"Bearer {new_token}"}
+ headers = await authorization_headers(auth, refresh=True)
try:
await self._stream_download(client, entity.url, headers, temp_path, logger)
except Exception:
@@ -201,6 +293,10 @@ async def download_from_url(
self._cleanup_temp(temp_path)
raise
+ except BaseException:
+ self._cleanup_temp(temp_path)
+ raise
+
logger.debug(f"Downloaded file to: {temp_path}")
entity.local_path = temp_path
return entity
diff --git a/backend/airweave/domains/storage/limits.py b/backend/airweave/domains/storage/limits.py
new file mode 100644
index 0000000000..d21d76d1db
--- /dev/null
+++ b/backend/airweave/domains/storage/limits.py
@@ -0,0 +1,3 @@
+"""Shared byte ceilings for captured files and their projection inputs."""
+
+MAX_FILE_SIZE_BYTES = 209715200
diff --git a/backend/airweave/domains/storage/protocols.py b/backend/airweave/domains/storage/protocols.py
index f36c7d2c84..c2d1973707 100644
--- a/backend/airweave/domains/storage/protocols.py
+++ b/backend/airweave/domains/storage/protocols.py
@@ -66,11 +66,12 @@ async def write_file(self, path: str, content: bytes) -> None:
"""
...
- async def read_file(self, path: str) -> bytes:
- """Read binary content from storage.
+ async def read_file(self, path: str, *, max_bytes: int | None = None) -> bytes:
+ """Read binary content from storage; optional bound stops reads at max_bytes + 1.
Args:
path: Relative path
+ max_bytes: Optional maximum actual bytes; oversized objects raise a storage error.
Returns:
Binary content
@@ -91,6 +92,14 @@ async def exists(self, path: str) -> bool:
"""
...
+ async def delete_file(self, path: str) -> None:
+ """Delete exactly one object; absence succeeds, other failures raise.
+
+ Never fall back to prefix/directory deletion. Backend retention/versioning
+ policies still govern physical historical-byte reclamation.
+ """
+ ...
+
async def delete(self, path: str) -> bool:
"""Delete a file or directory.
diff --git a/backend/airweave/domains/storage/tests/test_exact_delete.py b/backend/airweave/domains/storage/tests/test_exact_delete.py
new file mode 100644
index 0000000000..d92098c783
--- /dev/null
+++ b/backend/airweave/domains/storage/tests/test_exact_delete.py
@@ -0,0 +1,79 @@
+"""Strict exact-object deletion using installed SDK errors, without cloud calls."""
+
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
+from botocore.exceptions import ClientError
+from google.cloud.exceptions import Forbidden, NotFound
+
+from airweave.adapters.storage.aws_s3 import S3Backend
+from airweave.adapters.storage.azure_blob import AzureBlobBackend
+from airweave.adapters.storage.filesystem import FilesystemBackend
+from airweave.adapters.storage.gcp_gcs import GCSBackend
+from airweave.domains.storage.exceptions import StorageException
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_filesystem_exact_delete_is_idempotent_and_never_recursive(tmp_path, monkeypatch):
+ storage = FilesystemBackend(tmp_path)
+ await storage.write_file("one", b"body")
+ await storage.delete_file("one")
+ await storage.delete_file("one")
+ await storage.write_file("folder/child", b"keep")
+ with pytest.raises(StorageException):
+ await storage.delete_file("folder")
+ assert (tmp_path / "folder/child").read_bytes() == b"keep"
+ monkeypatch.setattr("aiofiles.os.remove", AsyncMock(side_effect=PermissionError("private")))
+ with pytest.raises(StorageException, match="^Exact file deletion failed$"):
+ await storage.delete_file("folder/child")
+
+
+async def test_s3_delete_uses_only_exact_key_and_preserves_errors():
+ storage = S3Backend("bucket", "us-east-1", prefix="owned")
+ client = Mock(delete_object=AsyncMock())
+ storage._get_client = AsyncMock(return_value=client)
+ await storage.delete_file("key")
+ client.delete_object.assert_awaited_once_with(Bucket="bucket", Key="owned/key")
+ client.delete_object.side_effect = ClientError({"Error": {"Code": "NoSuchKey"}}, "DeleteObject")
+ await storage.delete_file("key")
+ client.delete_object.side_effect = ClientError(
+ {"Error": {"Code": "AccessDenied"}}, "DeleteObject"
+ )
+ with pytest.raises(StorageException):
+ await storage.delete_file("key")
+ client.head_object.assert_not_called()
+ client.get_paginator.assert_not_called()
+
+
+async def test_azure_delete_uses_only_exact_blob_and_preserves_errors():
+ storage = AzureBlobBackend("account", "container", prefix="owned")
+ blob = Mock(delete_blob=AsyncMock())
+ container = Mock(get_blob_client=Mock(return_value=blob))
+ storage._get_container_client = AsyncMock(return_value=container)
+ await storage.delete_file("key")
+ container.get_blob_client.assert_called_once_with("owned/key")
+ blob.delete_blob.assert_awaited_once_with()
+ blob.delete_blob.side_effect = ResourceNotFoundError("missing")
+ await storage.delete_file("key")
+ blob.delete_blob.side_effect = HttpResponseError("private")
+ with pytest.raises(StorageException):
+ await storage.delete_file("key")
+ container.list_blobs.assert_not_called()
+
+
+async def test_gcs_delete_uses_only_exact_blob_and_preserves_errors():
+ storage = GCSBackend("bucket", prefix="owned")
+ blob = Mock(delete=Mock())
+ bucket = Mock(blob=Mock(return_value=blob))
+ storage._get_bucket = Mock(return_value=bucket)
+ await storage.delete_file("key")
+ bucket.blob.assert_called_once_with("owned/key")
+ blob.delete.assert_called_once_with()
+ blob.delete.side_effect = NotFound("missing")
+ await storage.delete_file("key")
+ blob.delete.side_effect = Forbidden("private")
+ with pytest.raises(StorageException):
+ await storage.delete_file("key")
+ bucket.list_blobs.assert_not_called()
diff --git a/backend/airweave/domains/storage/tests/test_file_service.py b/backend/airweave/domains/storage/tests/test_file_service.py
index 65637d0309..3cbbfa6878 100644
--- a/backend/airweave/domains/storage/tests/test_file_service.py
+++ b/backend/airweave/domains/storage/tests/test_file_service.py
@@ -5,6 +5,7 @@
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4
+import httpx
import pytest
from airweave.domains.storage.exceptions import FileSkippedException
@@ -19,9 +20,7 @@ def _make_service(tmpdir: str) -> FileService:
storage.write_file = AsyncMock()
storage.delete_directory = AsyncMock()
- with patch(
- "airweave.domains.storage.file_service.paths.temp_sync_dir", return_value=tmpdir
- ):
+ with patch("airweave.domains.storage.file_service.paths.temp_sync_dir", return_value=tmpdir):
svc = FileService(sync_job_id=sync_job_id, storage_backend=storage)
return svc, storage
@@ -128,3 +127,170 @@ async def test_does_not_raise_when_dir_already_missing(self):
# tmpdir has been deleted by the context manager exit
await svc.cleanup_sync_directory(logger=MagicMock())
+
+
+class TestCanonicalBlobs:
+ @pytest.mark.asyncio
+ async def test_canonical_capture_keeps_original_bytes_and_cleans_temp(self, tmp_path):
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+
+ svc, storage = _make_service(str(tmp_path))
+ svc.sync_id = uuid4()
+ content = b"original unsupported format"
+
+ async def download(
+ client,
+ url,
+ headers,
+ destination,
+ logger,
+ *,
+ follow_redirects=True,
+ expected_media_type=None,
+ ):
+ with open(destination, "wb") as output:
+ output.write(content)
+
+ svc._stream_download = download
+ blob = await svc.capture_canonical_url(
+ "https://example.com/file.unknown",
+ MagicMock(),
+ StaticTokenProvider("token"),
+ MagicMock(),
+ )
+ assert blob.key.startswith(f"canonical/{svc.sync_id}/blobs/sha256/")
+ storage.write_file.assert_awaited_once_with(blob.key, content)
+ assert not list(tmp_path.iterdir())
+
+ @pytest.mark.asyncio
+ async def test_storage_failure_does_not_publish_reference(self, tmp_path):
+ svc, storage = _make_service(str(tmp_path))
+ svc.sync_id = uuid4()
+ storage.write_file.side_effect = OSError("storage down")
+ with pytest.raises(OSError, match="storage down"):
+ await svc.store_canonical_blob(b"bytes")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("canonical", [True, False])
+@pytest.mark.parametrize("signed", [True, False])
+async def test_401_refresh_requires_original_bearer_request(tmp_path, canonical, signed):
+ from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+ from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+
+ auth = MagicMock(spec=TokenProviderProtocol)
+ auth.supports_refresh = True
+ auth.get_token = AsyncMock(return_value="fixture-old")
+ auth.force_refresh = AsyncMock(return_value="fixture-new")
+ requests = []
+
+ async def download(request):
+ if request.method == "HEAD":
+ return httpx.Response(200)
+ requests.append(request.headers.get("Authorization"))
+ return httpx.Response(401 if len(requests) == 1 else 200, content=b"original")
+
+ service, storage = _make_service(str(tmp_path))
+ service.sync_id = uuid4()
+ url = "https://files.example/original.pdf"
+ if signed:
+ url += "?X-Amz-Algorithm=fixture"
+ async with httpx.AsyncClient(transport=httpx.MockTransport(download)) as raw:
+ client = AirweaveHttpClient(raw, uuid4(), "fixture", feature_flag_enabled=False)
+ if canonical:
+ operation = service.capture_canonical_url(url, client, auth, MagicMock())
+ else:
+ entity = MagicMock(name="entity")
+ entity.name, entity.url = "original.pdf", url
+ operation = service.download_from_url(entity, client, auth, MagicMock())
+ if signed:
+ with pytest.raises(httpx.HTTPStatusError):
+ await operation
+ auth.force_refresh.assert_not_awaited()
+ assert requests == [None]
+ storage.write_file.assert_not_awaited()
+ assert not list(tmp_path.iterdir())
+ else:
+ await operation
+ auth.force_refresh.assert_awaited_once()
+ assert requests == ["Bearer fixture-old", "Bearer fixture-new"]
+
+
+@pytest.mark.asyncio
+async def test_canonical_can_reject_redirect_without_storing_its_body(tmp_path):
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+
+ service, storage = _make_service(str(tmp_path))
+ service.sync_id = uuid4()
+ calls = []
+
+ async def download(request):
+ if request.method == "HEAD":
+ return httpx.Response(200)
+ calls.append(request.url.path)
+ if request.url.path == "/original.pdf":
+ return httpx.Response(302, headers={"location": "/target.pdf"}, content=b"redirect")
+ return httpx.Response(200, content=b"original")
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(download)) as raw:
+ client = AirweaveHttpClient(raw, uuid4(), "fixture", feature_flag_enabled=False)
+ with pytest.raises(httpx.HTTPStatusError) as error:
+ await service.capture_canonical_url(
+ "https://files.example/original.pdf",
+ client,
+ StaticTokenProvider("fixture"),
+ MagicMock(),
+ follow_redirects=False,
+ )
+ assert error.value.response.status_code == 302
+ assert calls == ["/original.pdf"]
+ storage.write_file.assert_not_awaited()
+ assert not list(tmp_path.iterdir())
+ # Existing callers retain their redirect behavior without a new argument.
+ blob = await service.capture_canonical_url(
+ "https://files.example/original.pdf",
+ client,
+ StaticTokenProvider("fixture"),
+ MagicMock(),
+ )
+ assert calls == ["/original.pdf", "/original.pdf", "/target.pdf"]
+ storage.write_file.assert_awaited_once_with(blob.key, b"original")
+
+ entity = MagicMock(name="entity")
+ entity.name, entity.url = "original.pdf", "https://files.example/original.pdf"
+ downloaded = await service.download_from_url(
+ entity, client, StaticTokenProvider("fixture"), MagicMock()
+ )
+ assert calls[-2:] == ["/original.pdf", "/target.pdf"]
+ with open(downloaded.local_path, "rb") as original:
+ assert original.read() == b"original"
+
+
+@pytest.mark.asyncio
+async def test_expected_file_type_rejects_login_html_but_accepts_binary(tmp_path):
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+
+ service, storage = _make_service(str(tmp_path))
+ service.sync_id = uuid4()
+ responses = [
+ httpx.Response(200, headers={"Content-Type": "text/html; charset=utf-8"}, content=b"login"),
+ httpx.Response(
+ 200, headers={"Content-Type": "application/octet-stream"}, content=b"original"
+ ),
+ ]
+ async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: responses.pop(0))) as raw:
+ client = AirweaveHttpClient(raw, uuid4(), "fixture", feature_flag_enabled=False)
+ args = (
+ "https://files.example/original.pdf",
+ client,
+ StaticTokenProvider("fixture"),
+ MagicMock(),
+ )
+ with pytest.raises(ValueError, match="content type does not match"):
+ await service.capture_canonical_url(*args, expected_media_type="application/pdf")
+ storage.write_file.assert_not_awaited()
+ assert not list(tmp_path.iterdir())
+ blob = await service.capture_canonical_url(*args, expected_media_type="application/pdf")
+ storage.write_file.assert_awaited_once_with(blob.key, b"original")
diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py
new file mode 100644
index 0000000000..bf8c34faf5
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py
@@ -0,0 +1,365 @@
+"""Capture pipeline: native records commit before independent search projection."""
+
+from collections.abc import Awaitable, Callable
+from contextlib import AbstractAsyncContextManager
+from datetime import datetime, timezone
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.core.events.sync import EntityBatchProcessedEvent, TypeActionCounts
+from airweave.core.protocols.event_bus import EventBus
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CompleteCycle,
+ CycleConfiguration,
+ TerminalCheckpoint,
+)
+from airweave.domains.entities.canonical.models import CaptureResult
+from airweave.domains.entities.canonical.page_source import (
+ CanonicalPageSource,
+ CheckpointedPageSource,
+ ScopedPageSource,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureBatch,
+ CaptureRecord,
+ CompletedScope,
+ ReconcileScope,
+ RecordIdentity,
+ RemovedScope,
+ StartedScope,
+ WriterFence,
+)
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.source import SourceObservation
+from airweave.domains.storage.file_service import FileService
+from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
+from airweave.domains.sync_pipeline.contexts import SyncContext
+from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
+from airweave.domains.sync_pipeline.exceptions import SyncFailureError
+from airweave.domains.sync_pipeline.pipeline.cleanup_service import cleanup_service
+
+
+class CanonicalCapturePipeline:
+ """One source writer, one record authority; no legacy sink dispatcher participates."""
+
+ def __init__(
+ self,
+ service: CanonicalCaptureService,
+ sessions: Callable[[], AbstractAsyncContextManager[AsyncSession]],
+ event_bus: EventBus,
+ record_types: tuple[str, ...],
+ attempt: CaptureAttempt,
+ container_parents: dict[str, str | tuple[str | None, ...]] | None = None,
+ page_source: CanonicalPageSource | None = None,
+ files: FileService | None = None,
+ ):
+ """Inject transactional persistence and existing progress-event infrastructure."""
+ if not record_types:
+ raise ValueError("Canonical source must declare at least one audited record type")
+ self._service = service
+ self._sessions = sessions
+ self._event_bus = event_bus
+ self._record_types = frozenset(record_types)
+ self._container_parents = dict(container_parents or {})
+ self.page_source = page_source
+ self.files = files
+ self._cycle: CaptureCycle | None = None
+ if page_source is not None:
+ config = page_source.capture_cycle_configuration
+ declared = CycleConfiguration.from_source(
+ fingerprint=config.fingerprint,
+ record_types=record_types,
+ container_parents=self._container_parents,
+ completion_policies=config.completion_policies,
+ known_object_validation=config.known_object_validation,
+ scope_changes=config.scope_changes,
+ exact_parent_validation=config.exact_parent_validation,
+ )
+ if config != declared:
+ raise ValueError("Page source cycle must match its declared container topology")
+ else:
+ for child, parent in self._container_parents.items():
+ if (
+ child not in self._record_types
+ or not isinstance(parent, str)
+ or parent not in self._record_types
+ ):
+ raise ValueError(
+ "Container relationships must name declared canonical record types"
+ )
+ if parent in self._container_parents:
+ raise ValueError("Nested sources require the durable page capability")
+ self._attempt = attempt
+ self._fence: WriterFence | None = None
+ self._completed_scopes: set[CompletedScope] = set()
+ self._batch_sequence = 0
+ self._removed_scopes: set[CompletedScope] = set()
+
+ async def start(self, sync_context: SyncContext) -> None:
+ """Fence the activity after RUNNING transition and before fetching provider data."""
+ async with self._sessions() as db:
+ self._fence = await self._service.activate_writer(
+ db,
+ sync_context.organization_id,
+ sync_context.sync_id,
+ sync_context.sync_job_id,
+ attempt_id=self._attempt.id,
+ attempt_number=self._attempt.number,
+ )
+
+ async def run_scans(
+ self,
+ sync_context: SyncContext,
+ runtime: SyncRuntime,
+ check_limits: Callable[[], Awaitable[None]],
+ ) -> None:
+ """Run opted-in pages with a commit barrier before every next provider call."""
+ if self.files is None:
+ raise SyncFailureError("Page capture requires the source file service")
+ if self.page_source is None:
+ raise SyncFailureError("Source has no durable page capability")
+ if sync_context.execution_config.cursor.skip_updates:
+ raise SyncFailureError("Durable page capture cannot disable checkpoint updates")
+
+ async def progress(result: CaptureResult, records: tuple[CaptureRecord, ...]) -> None:
+ for record in records:
+ await runtime.entity_tracker.track_entity(
+ record.identity.record_type, record.identity.entity_key
+ )
+ await self._record_progress(result, sync_context, runtime)
+
+ self._cycle = await CanonicalScanDriver(
+ self._service,
+ self._sessions,
+ self._writer(),
+ self.page_source,
+ progress,
+ check_limits,
+ self.files,
+ force_full=(
+ sync_context.force_full_sync or sync_context.execution_config.cursor.skip_load
+ )
+ if isinstance(self.page_source, (CheckpointedPageSource, ScopedPageSource))
+ else False,
+ ).run()
+
+ def _writer(self) -> WriterFence:
+ if self._fence is None:
+ raise SyncFailureError("Canonical capture was not activated")
+ return self._fence
+
+ async def process(
+ self, entities: list[SourceObservation], sync_context: SyncContext, runtime: SyncRuntime
+ ) -> None:
+ """Persist original source observations in stream order; never infer raw data."""
+ records: list[CaptureRecord] = []
+ for observation in entities:
+ if isinstance(observation, CaptureRecord):
+ record_type = observation.identity.record_type
+ scope = CompletedScope(
+ record_type=record_type, container_id=observation.identity.container_id
+ )
+ if scope in self._removed_scopes:
+ raise SyncFailureError("Source yielded a record after confirmed scope loss")
+ records.append(self._with_parent(observation))
+ elif isinstance(observation, (StartedScope, RemovedScope, CompletedScope)):
+ record_type = observation.record_type
+ else:
+ raise SyncFailureError("Canonical source yielded a non-observation value")
+ if record_type not in self._record_types:
+ raise SyncFailureError(f"Undeclared canonical record type: {record_type}")
+ if isinstance(observation, CaptureRecord):
+ continue
+ # A lifecycle marker takes effect after all prior records and before following ones.
+ await self._capture_records(records, sync_context, runtime)
+ records = []
+ await self._apply_scope(observation, sync_context, runtime)
+ await self._capture_records(records, sync_context, runtime)
+
+ def _with_parent(self, record: CaptureRecord) -> CaptureRecord:
+ parent_type = self._container_parents.get(record.identity.record_type)
+ if parent_type is None:
+ return record
+ if record.identity.container_id is None:
+ raise SyncFailureError("Container-scoped canonical record is missing its container ID")
+ parent = RecordIdentity(record_type=parent_type, native_id=record.identity.container_id)
+ if record.parent is not None and record.parent != parent:
+ raise SyncFailureError("Record parent conflicts with its declared container identity")
+ return record.model_copy(update={"parent": parent})
+
+ async def _apply_scope(
+ self,
+ observation: StartedScope | RemovedScope | CompletedScope,
+ sync_context: SyncContext,
+ runtime: SyncRuntime,
+ ) -> None:
+ """Apply a scope boundary after earlier records have committed."""
+ scope = CompletedScope(
+ record_type=observation.record_type, container_id=observation.container_id
+ )
+ if isinstance(observation, StartedScope):
+ self._completed_scopes.discard(scope)
+ self._removed_scopes.discard(scope)
+ async with self._sessions() as db:
+ await self._service.start_scope(db, self._writer(), observation)
+ elif isinstance(observation, RemovedScope):
+ self._completed_scopes.discard(scope)
+ self._removed_scopes.add(scope)
+ while True:
+ async with self._sessions() as db:
+ result = await self._service.remove_scope(db, self._writer(), observation)
+ if result.capture.changes:
+ await self._record_progress(result.capture, sync_context, runtime)
+ if not result.has_more:
+ break
+ else:
+ if scope in self._removed_scopes:
+ raise SyncFailureError("Removed scope cannot be marked complete without a new scan")
+ self._completed_scopes.add(scope)
+
+ async def _capture_records(
+ self, records: list[CaptureRecord], sync_context: SyncContext, runtime: SyncRuntime
+ ) -> None:
+ # Keep bounded DB transactions independently of stream micro-batch tuning.
+ for offset in range(0, len(records), 500):
+ chunk = tuple(records[offset : offset + 500])
+ async with self._sessions() as db:
+ result = await self._service.capture(
+ db, CaptureBatch(fence=self._writer(), records=chunk)
+ )
+ for record in chunk:
+ await runtime.entity_tracker.track_entity(
+ record.identity.record_type, record.identity.entity_key
+ )
+ await self._record_progress(result, sync_context, runtime)
+
+ async def _record_progress(
+ self, result: CaptureResult, ctx: SyncContext, runtime: SyncRuntime
+ ) -> None:
+ counts: dict[str, TypeActionCounts] = {}
+ inserted = updated = deleted = 0
+ for change in result.changes:
+ record_type = change.record.identity.record_type
+ old = counts.get(record_type, TypeActionCounts())
+ if change.kind == "delete":
+ deleted += 1
+ counts[record_type] = old.model_copy(update={"deleted": old.deleted + 1})
+ await runtime.entity_tracker.record_deletes(record_type)
+ elif change.record.revision == 1:
+ inserted += 1
+ counts[record_type] = old.model_copy(update={"inserted": old.inserted + 1})
+ await runtime.entity_tracker.record_inserts(record_type)
+ else:
+ updated += 1
+ counts[record_type] = old.model_copy(update={"updated": old.updated + 1})
+ await runtime.entity_tracker.record_updates(record_type)
+ await runtime.entity_tracker.record_kept(result.unchanged)
+ self._batch_sequence += 1
+ await self._event_bus.publish(
+ EntityBatchProcessedEvent(
+ organization_id=ctx.organization_id,
+ sync_id=ctx.sync_id,
+ sync_job_id=ctx.sync_job_id,
+ collection_id=ctx.collection_id,
+ source_connection_id=ctx.source_connection_id,
+ source_type=ctx.source_short_name,
+ inserted=inserted,
+ updated=updated,
+ deleted=deleted,
+ kept=result.unchanged,
+ type_breakdown=counts,
+ batch_seq=self._batch_sequence,
+ )
+ )
+
+ async def cleanup_orphaned_entities(
+ self, sync_context: SyncContext, runtime: SyncRuntime
+ ) -> None:
+ """Reconcile only explicit successful scopes after the entire capture barrier."""
+ for scope in sorted(
+ self._completed_scopes, key=lambda s: (s.record_type, s.container_id or "")
+ ):
+ while True:
+ async with self._sessions() as db:
+ result = await self._service.reconcile_scope(
+ db,
+ ReconcileScope(
+ fence=self._writer(),
+ scope=scope,
+ # Missing containers withdraw content scope; absence does
+ # not prove provider deletion or credential revocation.
+ removal_reason=(
+ "scope_removed"
+ if scope.record_type in self._container_parents.values()
+ else "absent"
+ ),
+ observed_at=datetime.now(timezone.utc),
+ ),
+ )
+ if result.capture.changes:
+ await self._record_progress(result.capture, sync_context, runtime)
+ if not result.has_more:
+ break
+ await self._reconcile_parents(sync_context, runtime)
+
+ async def _reconcile_parents(self, sync_context: SyncContext, runtime: SyncRuntime) -> None:
+ while True:
+ async with self._sessions() as db:
+ result = await self._service.reconcile_parents(db, self._writer())
+ if result.capture.changes:
+ await self._record_progress(result.capture, sync_context, runtime)
+ if not result.has_more:
+ break
+
+ async def save_checkpoint(self, sync_context: SyncContext, runtime: SyncRuntime) -> None:
+ """Commit cursor only after source success, durable captures and scoped reconciliation."""
+ if self.page_source is not None:
+ if self._cycle is None:
+ raise SyncFailureError("Page capture did not reach its completion barrier")
+ async with self._sessions() as db:
+ current = await self._service.read_cycle(db, self._writer())
+ if current != self._cycle:
+ raise SyncFailureError("Cycle changed before source finalization")
+ if current.phase == "complete":
+ if current.completed_job_id != self._writer().job_id:
+ raise SyncFailureError("Completed cycle belongs to another job")
+ return
+ terminal = await self._terminal_checkpoint(current)
+ async with self._sessions() as db:
+ self._cycle = await self._service.complete_cycle(
+ db,
+ CompleteCycle(
+ fence=self._writer(), expected=current.version, terminal_checkpoint=terminal
+ ),
+ )
+ return
+ if runtime.cursor is None or not runtime.cursor.cursor_data:
+ return
+ if sync_context.execution_config and sync_context.execution_config.cursor.skip_updates:
+ return
+ async with self._sessions() as db:
+ await self._service.save_checkpoint(db, self._writer(), runtime.cursor.cursor_data)
+
+ async def _terminal_checkpoint(self, current: CaptureCycle) -> TerminalCheckpoint | None:
+ """Read only the committed terminal root boundary; the store reattests it at publication."""
+ terminal = None
+ if current.starting_checkpoint is not None or current.mode == "changes":
+ if len(current.configuration.root_record_types) != 1:
+ raise SyncFailureError("Checkpoint capture requires one terminal root scope")
+ kind = current.configuration.root_record_types[0]
+ async with self._sessions() as db:
+ scan = await self._service.read_scan(
+ db, self._writer(), CompletedScope(record_type=kind)
+ )
+ if scan is None or scan.provider_checkpoint is None:
+ raise SyncFailureError("Capture has no committed terminal provider checkpoint")
+ terminal = TerminalCheckpoint(
+ record_type=kind, expected=scan.version, checkpoint=scan.provider_checkpoint
+ )
+ return terminal
+
+ async def cleanup_temp_files(self, sync_context: SyncContext, runtime: SyncRuntime) -> None:
+ """Reuse source-job temporary-file cleanup; immutable blobs are not temporary."""
+ await cleanup_service.cleanup_temp_files(sync_context, runtime)
diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py
new file mode 100644
index 0000000000..42573091d6
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py
@@ -0,0 +1,516 @@
+"""Sequential page driver inside the existing capture job, without provider prefetch."""
+
+from collections.abc import Awaitable, Callable
+from contextlib import AbstractAsyncContextManager
+from datetime import datetime, timezone
+from uuid import UUID
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from airweave.domains.entities.canonical.cycle_models import (
+ BeginCycle,
+ CaptureCycle,
+ RestartCycle,
+ ScopeWork,
+)
+from airweave.domains.entities.canonical.cycle_store import CycleConflict
+from airweave.domains.entities.canonical.models import CaptureResult, SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CanonicalPageSource,
+ CapturePlan,
+ CheckpointedPageSource,
+ InvalidCaptureCheckpoint,
+ InvalidScanContinuation,
+ InvalidScopeCheckpoint,
+ KnownObjectSource,
+ RequiredScopeAccessLost,
+ ScopeAccessLost,
+ ScopedPageSource,
+ ScopeRemovalReason,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ WriterFence,
+)
+from airweave.domains.entities.canonical.scan_models import (
+ BeginScan,
+ CommitOmission,
+ CommitScanPage,
+ ReconcileScan,
+ ScanAdmission,
+ ScanContinuation,
+ ScanState,
+)
+from airweave.domains.entities.canonical.scope_execution import ScopePlan
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.storage.file_service import FileService
+
+
+class _ChildScopeAccessLost(Exception):
+ """A required child failure whose owner has already been withdrawn."""
+
+ def __init__(self, error: RequiredScopeAccessLost) -> None:
+ super().__init__(str(error))
+ self.error = error
+
+
+class CanonicalScanDriver:
+ """Provider calls occur between existing fenced service transactions, never within one."""
+
+ def __init__(
+ self,
+ service: CanonicalCaptureService,
+ sessions: Callable[[], AbstractAsyncContextManager[AsyncSession]],
+ fence: WriterFence,
+ source: CanonicalPageSource,
+ progress: Callable[[CaptureResult, tuple[CaptureRecord, ...]], Awaitable[None]],
+ check_limits: Callable[[], Awaitable[None]],
+ files: FileService,
+ *,
+ force_full: bool = False,
+ ):
+ """Reuse the capture service and pipeline progress/guard callbacks."""
+ self.service, self.sessions, self.fence, self.source = service, sessions, fence, source
+ self.files = files
+ self.force_full = force_full
+ self.progress, self.check_limits = progress, check_limits
+ self._restarted_scopes: set[tuple[UUID, CompletedScope]] = set()
+
+ async def run(self) -> CaptureCycle:
+ """Refresh membership first, then resume its currently visible child scopes."""
+ configuration = self.source.capture_cycle_configuration
+ async with self.sessions() as db:
+ current = await self.service.read_cycle(db, self.fence)
+ if current is not None and current.phase == "complete":
+ if current.completed_job_id == self.fence.job_id:
+ return current
+ expected = current.version
+ else:
+ expected = None
+ if (
+ self.force_full
+ and isinstance(self.source, (CheckpointedPageSource, ScopedPageSource))
+ and current is not None
+ and current.phase == "active"
+ and (
+ current.mode == "changes"
+ or (current.mode == "mixed" and not current.force_full_scopes)
+ or current.configuration != configuration
+ )
+ ):
+ cycle = await self.fresh_full(current)
+ else:
+ plan = (
+ CapturePlan(
+ mode=current.mode,
+ starting_checkpoint=current.starting_checkpoint,
+ source_plan=current.source_plan,
+ )
+ if current is not None and current.phase == "active"
+ else (
+ await self.source.prepare_cycle(None if self.force_full else current)
+ if isinstance(self.source, (CheckpointedPageSource, ScopedPageSource))
+ else CapturePlan()
+ )
+ )
+ async with self.sessions() as db:
+ cycle = await self.service.begin_cycle(
+ db,
+ BeginCycle(
+ fence=self.fence,
+ configuration=configuration,
+ expected=expected,
+ mode=plan.mode,
+ starting_checkpoint=plan.starting_checkpoint,
+ source_plan=plan.source_plan,
+ force_full_scopes=(
+ current.force_full_scopes
+ if current and current.phase == "active"
+ else self.force_full
+ )
+ if plan.mode == "mixed"
+ else False,
+ ),
+ )
+ for restart in range(2):
+ try:
+ return await self.run_cycle(cycle)
+ except InvalidCaptureCheckpoint:
+ if restart or not isinstance(
+ self.source, (CheckpointedPageSource, ScopedPageSource)
+ ):
+ raise
+ await self.check_limits()
+ cycle = await self.fresh_full(cycle)
+ raise AssertionError("Unreachable checkpoint recovery")
+
+ async def fresh_full(self, previous: CaptureCycle) -> CaptureCycle:
+ """Explicit full requests and native expiry share the same exact-cycle restart."""
+ if not isinstance(self.source, (CheckpointedPageSource, ScopedPageSource)):
+ raise CycleConflict("Source has no checkpoint recovery contract")
+ plan = await self.source.prepare_cycle(None)
+ if plan.mode != "full" and not (
+ plan.mode == "mixed" and isinstance(self.source, ScopedPageSource)
+ ):
+ raise CycleConflict("Checkpoint recovery requires a fresh full capture")
+ async with self.sessions() as db:
+ current = await self.service.read_cycle(db, self.fence)
+ if current is None or current.version.cycle_id != previous.version.cycle_id:
+ raise CycleConflict("Cycle changed during checkpoint recovery")
+ return await self.service.restart_cycle(
+ db,
+ RestartCycle(
+ fence=self.fence,
+ expected=current.version,
+ configuration=self.source.capture_cycle_configuration,
+ mode=plan.mode,
+ starting_checkpoint=plan.starting_checkpoint,
+ source_plan=plan.source_plan,
+ force_full_scopes=plan.mode == "mixed",
+ ),
+ )
+
+ async def run_cycle(self, cycle: CaptureCycle) -> CaptureCycle:
+ """Process the current frontier; every acknowledgement remains in the existing SQL store."""
+ while True:
+ async with self.sessions() as db:
+ work = await self.service.next_scope_work(db, self.fence, cycle.version.cycle_id)
+ if work is None:
+ break
+ try:
+ await self.scan_work(cycle, work)
+ except _ChildScopeAccessLost as failure:
+ # Preserve the provider failure without withdrawing its enclosing scope.
+ raise failure.error from failure
+ async with self.sessions() as db:
+ current = await self.service.read_cycle(db, self.fence)
+ if current is None:
+ raise CycleConflict("Capture cycle disappeared before finalization")
+ return current
+
+ async def scan_work(
+ self, cycle: CaptureCycle, work: ScopeWork, *, max_pages: int | None = None
+ ) -> None:
+ """Ordinary and interleaved work share exact-owner and withdrawal handling."""
+ scope = (
+ CompletedScope(record_type=work.record_type)
+ if work.parent is None
+ else self.source.child_scope(work.parent, work.record_type)
+ )
+ if scope.record_type != work.record_type or scope.parent != (
+ work.parent.identity if work.parent else None
+ ):
+ raise CycleConflict("Source returned a scope with the wrong exact owner")
+ try:
+ await self.scan(
+ cycle,
+ scope,
+ parent=work.parent,
+ parent_epoch=work.parent_visibility_epoch,
+ refresh_membership=bool(cycle.configuration.children_of(work.record_type)),
+ max_pages=max_pages,
+ )
+ except ScopeAccessLost as error:
+ if work.parent is None:
+ raise
+ if work.parent_visibility_epoch is None:
+ raise CycleConflict("Scope owner epoch is missing") from error
+ await self.withdraw_parent(
+ work.parent, work.parent_visibility_epoch, error.removal_reason
+ )
+ if isinstance(error, RequiredScopeAccessLost):
+ raise
+
+ async def advance_children(self, cycle: CaptureCycle, scope: CompletedScope) -> None:
+ """Give freshly observed, exactly verified owners one child page of progress."""
+ if scope.record_type not in cycle.configuration.exact_parent_validation:
+ return
+ async with self.sessions() as db:
+ work = await self.service.next_scope_work(
+ db, self.fence, cycle.version.cycle_id, within=scope
+ )
+ if work is not None:
+ try:
+ await self.scan_work(cycle, work, max_pages=1)
+ except RequiredScopeAccessLost as error:
+ raise _ChildScopeAccessLost(error) from error
+
+ async def scan(
+ self,
+ cycle: CaptureCycle,
+ scope: CompletedScope,
+ *,
+ refresh_membership: bool = False,
+ parent: SourceRecord | None = None,
+ parent_epoch: int | None = None,
+ max_pages: int | None = None,
+ ) -> None:
+ """Only a provider invalid-cursor response permits one explicit sweep restart."""
+ admission = await self.begin_scope(cycle, scope, parent, parent_epoch)
+ if admission.state is None:
+ return
+ state = admission.state
+ parent = admission.parent
+ restart_key = (cycle.version.cycle_id, scope)
+ pages = 0
+ while state.phase == "collecting":
+ await self.check_limits()
+ try:
+ page = await self.source.capture_page(
+ scope, state.continuation, files=self.files, parent=parent
+ )
+ except (InvalidScanContinuation, InvalidScopeCheckpoint) as error:
+ if restart_key in self._restarted_scopes:
+ raise
+ self._restarted_scopes.add(restart_key)
+ state = await self.restart_scope(cycle, scope, state, parent, error)
+ continue
+ # Never split a page or retry an uncertain commit with this old version.
+ async with self.sessions() as db:
+ result = await self.service.commit_scan_page(
+ db,
+ CommitScanPage(
+ fence=self.fence,
+ scope=scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ records=tuple(self._parented(r, scope) for r in page.records),
+ discovered_records=page.discovered_records,
+ child_scope_observations=page.child_scope_observations,
+ continuation=page.continuation,
+ final=page.final,
+ provider_checkpoint=page.provider_checkpoint,
+ ),
+ )
+ state = result.state
+ await self.progress(result.capture, (*page.records, *page.discovered_records))
+ pages += 1
+ await self.advance_children(cycle, scope)
+ if max_pages is not None and pages >= max_pages and state.phase == "collecting":
+ return
+ if state.phase == "reconciling":
+ if state.completion_policy == "discovery_with_validation" or (
+ state.scope.record_type in cycle.configuration.known_object_validation
+ ):
+ state = await self.refresh_known(state)
+ elif state.completion_policy == "exhaustive" and refresh_membership:
+ await self.confirm_omissions(state)
+ await self.reconcile(state)
+
+ async def begin_scope(
+ self,
+ cycle: CaptureCycle,
+ scope: CompletedScope,
+ parent: SourceRecord | None,
+ parent_epoch: int | None,
+ ) -> ScanAdmission:
+ """Select a plan outside SQL, then attest prior version and owner under the fence."""
+ async with self.sessions() as db:
+ previous = await self.service.read_scan(db, self.fence, scope)
+ restart = bool(
+ previous
+ and previous.cycle_id == cycle.version.cycle_id
+ and (
+ (
+ cycle.configuration.fresh_inventory(scope.record_type)
+ and previous.membership_attempt_id != self.fence.attempt_id
+ )
+ or previous.parent_visibility_epoch != parent_epoch
+ )
+ )
+ scoped = isinstance(self.source, ScopedPageSource)
+ if cycle.mode == "mixed" and cycle.configuration.exact_parent_validation:
+ raise CycleConflict("Mixed planning does not support exact parent validation")
+ if cycle.mode == "mixed" and not scoped:
+ raise CycleConflict("Mixed capture requires source scope planning")
+ if scoped and cycle.force_full_scopes and previous and previous.mode == "changes":
+ restart = True
+ plan = None
+ if scoped:
+ if previous and previous.cycle_id == cycle.version.cycle_id and not restart:
+ if previous.execution is None:
+ raise CycleConflict("Mixed scope lost its persisted plan")
+ plan = previous.execution.plan
+ else:
+ plan = await self.source.prepare_scope(
+ scope,
+ cycle,
+ previous,
+ parent=parent,
+ force_full=cycle.force_full_scopes
+ or bool(previous and previous.parent_visibility_epoch != parent_epoch),
+ )
+ observation = None
+ if parent and parent.identity.record_type in cycle.configuration.exact_parent_validation:
+ verified = bool(
+ previous
+ and previous.cycle_id == cycle.version.cycle_id
+ and previous.parent_verified_attempt_id == self.fence.attempt_id
+ and previous.parent_verified_revision == parent.revision
+ and previous.parent_visibility_epoch == parent_epoch
+ )
+ if not verified:
+ if not isinstance(self.source, KnownObjectSource):
+ raise CycleConflict("Source requires exact owner refresh support")
+ await self.check_limits()
+ observation = await self.source.refresh_known(parent, files=self.files)
+ initial = self.scope_initial(scope, cycle, plan)
+ async with self.sessions() as db:
+ admission = await self.service.admit_scan(
+ db,
+ BeginScan(
+ fence=self.fence,
+ scope=scope,
+ cycle_id=cycle.version.cycle_id,
+ fingerprint=cycle.configuration.fingerprint,
+ expected=previous.version if previous else None,
+ restart=restart,
+ continuation=initial,
+ plan=plan,
+ expected_parent_epoch=parent_epoch,
+ expected_parent_revision=parent.revision if parent else None,
+ exact_parent_observation=observation,
+ ),
+ )
+
+ await self.progress(admission.capture, (observation,) if observation else ())
+ return admission
+
+ async def restart_scope(
+ self,
+ cycle: CaptureCycle,
+ scope: CompletedScope,
+ state: ScanState,
+ parent: SourceRecord | None,
+ error: Exception,
+ ) -> ScanState:
+ """Native cursor expiry can reset one sweep; native sync expiry forces full scope mode."""
+ plan = state.execution.plan if state.execution else None
+ if isinstance(error, InvalidScopeCheckpoint):
+ if not isinstance(self.source, ScopedPageSource):
+ raise error
+ plan = await self.source.prepare_scope(
+ scope, cycle, state, parent=parent, force_full=True
+ )
+ if plan.mode != "full":
+ raise CycleConflict("Invalid scope checkpoint requires a full scope restart")
+ initial = self.scope_initial(scope, cycle, plan)
+ async with self.sessions() as db:
+ return await self.service.begin_scan(
+ db,
+ BeginScan(
+ fence=self.fence,
+ scope=scope,
+ cycle_id=state.cycle_id,
+ fingerprint=state.fingerprint,
+ expected=state.version,
+ restart=True,
+ continuation=initial,
+ plan=plan,
+ expected_parent_epoch=state.parent_visibility_epoch,
+ expected_parent_revision=parent.revision if parent else None,
+ ),
+ )
+
+ def scope_initial(
+ self, scope: CompletedScope, cycle: CaptureCycle, plan: ScopePlan | None
+ ) -> ScanContinuation:
+ """Sources receive exactly the plan that will be committed with the new sweep."""
+ if isinstance(self.source, ScopedPageSource):
+ if plan is None:
+ raise CycleConflict("Mixed source requires a scope plan")
+ return self.source.initial_scope_continuation(scope, cycle, plan)
+ return self.initial_continuation(cycle)
+
+ def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation:
+ """Sources without a native changes contract retain the empty initial cursor."""
+ return (
+ self.source.initial_continuation(cycle)
+ if isinstance(self.source, CheckpointedPageSource)
+ else ScanContinuation()
+ )
+
+ @staticmethod
+ def _parented(record: CaptureRecord, scope: CompletedScope) -> CaptureRecord:
+ """Legacy flat adapters may omit the parent; never overwrite a conflicting declaration."""
+ if record.parent is not None and record.parent != scope.parent:
+ raise CycleConflict("Captured record declares a different scope parent")
+ return record.model_copy(update={"parent": scope.parent})
+
+ async def refresh_known(self, state: ScanState) -> ScanState:
+ """SQL sightings are the durable frontier; uncertain commits end this attempt."""
+ if not isinstance(self.source, KnownObjectSource):
+ raise CycleConflict("Source does not implement exact known-object validation")
+ while True:
+ async with self.sessions() as db:
+ records = await self.service.scan_missing(db, self.fence, state)
+ if not records:
+ return state
+ for record in records:
+ await self.check_limits()
+ observation = await self.source.refresh_known(record, files=self.files)
+ async with self.sessions() as db:
+ result = await self.service.commit_omission(
+ db,
+ CommitOmission(
+ fence=self.fence,
+ state=state,
+ expected_record=record,
+ observation=observation,
+ ),
+ )
+ state = result.state
+ await self.progress(result.capture, (observation,))
+
+ async def confirm_omissions(self, state: ScanState) -> None:
+ """Accessible omissions or provider errors fail before any absence removal."""
+ after: UUID | None = None
+ while True:
+ async with self.sessions() as db:
+ records = await self.service.scan_missing(db, self.fence, state, after=after)
+ if not records:
+ return
+ for record in records:
+ await self.check_limits()
+ await self.source.confirm_absent(record)
+ after = records[-1].id
+
+ async def reconcile(self, state: ScanState) -> None:
+ """Each bounded removal batch resumes from its committed scan version."""
+ while state.phase == "reconciling":
+ await self.check_limits()
+ async with self.sessions() as db:
+ result = await self.service.reconcile_scan(
+ db,
+ ReconcileScan(
+ fence=self.fence,
+ scope=state.scope,
+ cycle_id=state.cycle_id,
+ expected=state.version,
+ observed_at=datetime.now(timezone.utc),
+ ),
+ )
+ state = result.state
+ await self.progress(result.capture, ())
+
+ async def withdraw_parent(
+ self, parent: SourceRecord, expected_epoch: int, removal_reason: ScopeRemovalReason
+ ) -> None:
+ """Withdraw exact owner first; recursive visibility closes before bounded cleanup."""
+ async with self.sessions() as db:
+ result = await self.service.withdraw_scan_parent(
+ db,
+ self.fence,
+ parent,
+ expected_epoch=expected_epoch,
+ removal_reason=removal_reason,
+ )
+ await self.progress(result, ())
+ while True:
+ await self.check_limits()
+ async with self.sessions() as db:
+ result = await self.service.reconcile_parents(db, self.fence)
+ await self.progress(result.capture, ())
+ if not result.has_more:
+ return
diff --git a/backend/airweave/domains/sync_pipeline/capture_attempt.py b/backend/airweave/domains/sync_pipeline/capture_attempt.py
new file mode 100644
index 0000000000..acd8d0c8d6
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/capture_attempt.py
@@ -0,0 +1,28 @@
+"""Source activity attempt identity, including explicit non-Temporal invocation."""
+
+from uuid import NAMESPACE_URL, UUID, uuid5
+
+from pydantic import BaseModel, ConfigDict, Field
+from temporalio import activity
+
+
+class CaptureAttempt(BaseModel):
+ """Attempt number is supplied by the execution owner, never invented on restart."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ id: UUID
+ number: int = Field(ge=1)
+
+
+def resolve_capture_attempt(explicit: CaptureAttempt | None = None) -> CaptureAttempt:
+ """Use real Temporal metadata or require a deliberate local invocation identity."""
+ if explicit is not None:
+ return explicit
+ if not activity.in_activity():
+ raise ValueError("Canonical capture outside Temporal requires an explicit CaptureAttempt")
+ info = activity.info()
+ identity = (
+ f"temporal:{info.workflow_namespace}:{info.workflow_id}:"
+ f"{info.workflow_run_id}:{info.activity_id}:{info.attempt}"
+ )
+ return CaptureAttempt(id=uuid5(NAMESPACE_URL, identity), number=info.attempt)
diff --git a/backend/airweave/domains/sync_pipeline/contexts/runtime.py b/backend/airweave/domains/sync_pipeline/contexts/runtime.py
index fc3ba60199..383086ba62 100644
--- a/backend/airweave/domains/sync_pipeline/contexts/runtime.py
+++ b/backend/airweave/domains/sync_pipeline/contexts/runtime.py
@@ -11,6 +11,7 @@
from typing import TYPE_CHECKING, List, Optional
if TYPE_CHECKING:
+ from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
from airweave.domains.sync_pipeline.pipeline.entity_tracker import EntityTracker
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.destinations._base import BaseDestination
@@ -28,5 +29,6 @@ class SyncRuntime:
source: "BaseSource"
entity_tracker: "EntityTracker"
+ canonical_capture: Optional["CanonicalCapturePipeline"] = None
cursor: Optional["SyncCursor"] = None
destinations: List["BaseDestination"] = field(default_factory=list)
diff --git a/backend/airweave/domains/sync_pipeline/entity/pipeline.py b/backend/airweave/domains/sync_pipeline/entity/pipeline.py
index 5198bbdbe2..5d01bc2ad2 100644
--- a/backend/airweave/domains/sync_pipeline/entity/pipeline.py
+++ b/backend/airweave/domains/sync_pipeline/entity/pipeline.py
@@ -13,10 +13,9 @@
import time
from collections import defaultdict
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import TYPE_CHECKING, Dict, List
from airweave.core.events.sync import EntityBatchProcessedEvent, TypeActionCounts
-from airweave.core.shared_models import AirweaveFieldFlag
from airweave.domains.entities.protocols import EntityRepositoryProtocol
from airweave.domains.sync_pipeline.contexts import SyncContext
from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
@@ -25,6 +24,7 @@
from airweave.domains.sync_pipeline.pipeline.cleanup_service import cleanup_service
from airweave.domains.sync_pipeline.pipeline.entity_tracker import EntityTracker
from airweave.domains.sync_pipeline.pipeline.hash_computer import hash_computer
+from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields
from airweave.domains.sync_pipeline.protocols import (
EntityActionDispatcherProtocol,
EntityActionResolverProtocol,
@@ -333,39 +333,9 @@ async def _cleanup_temp_files_for_batch(
# Helper Methods
# -------------------------------------------------------------------------
- def _get_flagged_field_value(self, entity: BaseEntity, flag_name: str) -> Any:
- """Extract value of field marked with specified flag."""
- flag_key = flag_name.value if hasattr(flag_name, "value") else flag_name
-
- for field_name, field_info in entity.model_fields.items():
- json_extra = field_info.json_schema_extra
- if json_extra and isinstance(json_extra, dict):
- if json_extra.get(flag_key):
- return getattr(entity, field_name, None)
-
- return None
-
def _populate_base_entity_fields_from_flags(self, entity: BaseEntity) -> None:
- """Populate BaseEntity fields from flagged fields."""
- if not entity.entity_id:
- flagged_value = self._get_flagged_field_value(entity, AirweaveFieldFlag.IS_ENTITY_ID)
- if flagged_value:
- entity.entity_id = str(flagged_value)
-
- if not entity.name:
- flagged_value = self._get_flagged_field_value(entity, AirweaveFieldFlag.IS_NAME)
- if flagged_value:
- entity.name = str(flagged_value)
-
- if not entity.created_at:
- flagged_value = self._get_flagged_field_value(entity, AirweaveFieldFlag.IS_CREATED_AT)
- if flagged_value is not None:
- entity.created_at = flagged_value
-
- if not entity.updated_at:
- flagged_value = self._get_flagged_field_value(entity, AirweaveFieldFlag.IS_UPDATED_AT)
- if flagged_value is not None:
- entity.updated_at = flagged_value
+ """Use shared flag normalization for legacy and canonical projection."""
+ populate_base_fields(entity)
async def _enrich_early_metadata(
self,
diff --git a/backend/airweave/domains/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py
index cdf225b02f..692ce64ef3 100644
--- a/backend/airweave/domains/sync_pipeline/factory.py
+++ b/backend/airweave/domains/sync_pipeline/factory.py
@@ -15,6 +15,7 @@
from typing import List, Optional
from uuid import UUID
+from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from airweave import crud, schemas
@@ -22,6 +23,7 @@
from airweave.core.exceptions import NotFoundException
from airweave.core.logging import ContextualLogger, LoggerConfigurator, logger
from airweave.core.protocols.event_bus import EventBus
+from airweave.db.session import get_db_context
from airweave.domains.access_control.dispatcher import ACActionDispatcher
from airweave.domains.access_control.membership_tracker import ACLMembershipTracker
from airweave.domains.access_control.pipeline import AccessControlPipeline
@@ -31,6 +33,10 @@
from airweave.domains.arf.protocols import ArfServiceProtocol
from airweave.domains.browse_tree.protocols import NodeSelectionRepositoryProtocol
from airweave.domains.browse_tree.types import NodeSelectionData
+from airweave.domains.entities.canonical.page_source import CanonicalPageSource
+from airweave.domains.entities.canonical.service import CanonicalCaptureService
+from airweave.domains.entities.canonical.source import CanonicalSource, ContainerScopedSource
+from airweave.domains.entities.canonical.store import CanonicalRecordStore, StaleWriter
from airweave.domains.entities.protocols import (
EntityCountRepositoryProtocol,
EntityRepositoryProtocol,
@@ -43,6 +49,8 @@
from airweave.domains.storage.protocols import StorageBackend
from airweave.domains.sync_pipeline.builders import SyncContextBuilder
from airweave.domains.sync_pipeline.builders.destinations import DestinationsContextBuilder
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt, resolve_capture_attempt
from airweave.domains.sync_pipeline.config import SyncConfig
from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
from airweave.domains.sync_pipeline.contexts.sync import SyncContext
@@ -60,7 +68,11 @@
from airweave.domains.syncs.jobs.protocols import SyncJobStateMachineProtocol
from airweave.domains.syncs.protocols import SyncStateMachineProtocol
from airweave.domains.usage.protocols import UsageLedgerProtocol, UsageLimitCheckerProtocol
+from airweave.models.capture_scan import CaptureScan
+from airweave.models.entity import Entity
from airweave.models.source_connection import SourceConnection
+from airweave.models.sync_cursor import SyncCursor as StoredSyncCursor
+from airweave.platform.configs.config import OutlookMailConfig, StripeConfig
from airweave.platform.sources._base import BaseSource
from .entity.pipeline import EntityPipeline
@@ -145,6 +157,7 @@ async def create_orchestrator(
force_full_sync: bool = False,
execution_config: Optional[SyncConfig] = None,
access_token: Optional[str] = None,
+ capture_attempt: CaptureAttempt | None = None,
) -> SyncOrchestrator:
"""Create a dedicated orchestrator instance for a sync run."""
init_start = time.time()
@@ -161,7 +174,11 @@ async def create_orchestrator(
f"destinations={resolved_config.destinations.model_dump()}"
)
+ admission = CanonicalRecordStore()
+ generation = await admission.admit_job(db, ctx.organization.id, sync.id, sync_job.id)
sc = await self._resolve_source_connection(db, sync, ctx)
+ # The session may already hold config from before a reconnect.
+ await db.refresh(sc)
sc_id = sc.id # extract before _build_source can expire the ORM instance via OAuth refresh
# 2. Build source, destinations, tracker
@@ -184,14 +201,23 @@ async def create_orchestrator(
execution_config=resolved_config,
access_token=access_token,
)
+ if await admission.admit_job(db, ctx.organization.id, sync.id, sync_job.id) != generation:
+ raise StaleWriter("Source authorization changed during initialization")
source_entry = self._source_registry.get(sc.short_name)
- destinations = await self._build_destinations(
- db=db,
- sync=sync,
- collection=collection,
- ctx=ctx,
- execution_config=resolved_config,
- source_supports_acl=source_entry.supports_access_control,
+ canonical_source = source_result.source.capture_page_source
+ if canonical_source is None and isinstance(source_result.source, CanonicalSource):
+ canonical_source = source_result.source
+ destinations = (
+ []
+ if canonical_source is not None
+ else await self._build_destinations(
+ db=db,
+ sync=sync,
+ collection=collection,
+ ctx=ctx,
+ execution_config=resolved_config,
+ source_supports_acl=source_entry.supports_access_control,
+ )
)
entity_tracker = await self._build_entity_tracker(
db=db,
@@ -224,9 +250,28 @@ async def create_orchestrator(
logger.debug(f"Context + runtime built in {time.time() - init_start:.2f}s")
# 4. Wire pipelines
- entity_pipeline = self._build_entity_pipeline(
- sync_context, runtime, destinations, resolved_config
- )
+ if canonical_source is not None:
+ runtime.canonical_capture = CanonicalCapturePipeline(
+ service=CanonicalCaptureService(CanonicalRecordStore()),
+ sessions=get_db_context,
+ event_bus=self._event_bus,
+ record_types=canonical_source.canonical_record_types,
+ attempt=resolve_capture_attempt(capture_attempt),
+ files=source_result.files,
+ page_source=(
+ canonical_source if isinstance(canonical_source, CanonicalPageSource) else None
+ ),
+ container_parents=(
+ canonical_source.canonical_container_parents
+ if isinstance(canonical_source, ContainerScopedSource)
+ else None
+ ),
+ )
+ entity_pipeline = runtime.canonical_capture
+ else:
+ entity_pipeline = self._build_entity_pipeline(
+ sync_context, runtime, destinations, resolved_config
+ )
access_control_pipeline = self._build_access_control_pipeline(sync_context)
stream = self._build_stream(runtime, source_result, sync_context)
@@ -314,14 +359,26 @@ def _build_stream(
runtime: SyncRuntime,
source_result: SourceBuildResult,
sync_context: SyncContext,
- ) -> AsyncSourceStream:
- """Build the async source stream from the source generator."""
- return AsyncSourceStream(
- source_generator=runtime.source.generate_entities(
+ ) -> AsyncSourceStream | None:
+ """Page sources are driven sequentially; legacy generators retain their stream."""
+ if runtime.source.capture_page_source is not None:
+ if source_result.node_selections:
+ raise ValueError("Whole-scope page capture does not support selected nodes")
+ return None
+ if isinstance(runtime.source, CanonicalSource):
+ generator = runtime.source.generate_observations(
cursor=runtime.cursor,
files=source_result.files,
node_selections=source_result.node_selections,
- ),
+ )
+ else:
+ generator = runtime.source.generate_entities(
+ cursor=runtime.cursor,
+ files=source_result.files,
+ node_selections=source_result.node_selections,
+ )
+ return AsyncSourceStream(
+ source_generator=generator,
queue_size=10000,
logger=sync_context.logger,
)
@@ -343,7 +400,24 @@ async def _build_source(
access_token: Optional[str] = None,
) -> SourceBuildResult:
"""Build source instance, cursor, file service, and node selections."""
+ composed_capture = False
+ if source_connection.short_name == "outlook_mail":
+ composed_capture = OutlookMailConfig.model_validate(
+ source_connection.config_fields or {}
+ ).capture_originals
+ elif source_connection.short_name == "stripe":
+ composed_capture = (
+ StripeConfig.model_validate(source_connection.config_fields or {}).original_capture
+ is not None
+ )
+ if source_connection.short_name in {"outlook_mail", "stripe"}:
+ await self._validate_composed_capture_mode(db, sync.id, ctx, composed_capture)
if execution_config and execution_config.behavior.replay_from_arf:
+ source_class = self._source_registry.get(source_connection.short_name).source_class_ref
+ if composed_capture or isinstance(source_class, (CanonicalSource, CanonicalPageSource)):
+ raise ValueError(
+ "Canonical records must be reindexed from Postgres, not mutable ARF"
+ )
return await self._build_arf_replay_source(db=db, sync=sync, ctx=ctx, logger=logger)
self._validate_not_completed_snapshot(source_connection)
@@ -357,7 +431,9 @@ async def _build_source(
access_token=access_token,
)
- files = FileService(sync_job_id=sync_job.id, storage_backend=self._storage_backend)
+ files = FileService(
+ sync_job_id=sync_job.id, sync_id=sync.id, storage_backend=self._storage_backend
+ )
cursor = await self._create_cursor(
db=db,
@@ -377,6 +453,46 @@ async def _build_source(
source=source, cursor=cursor, files=files, node_selections=node_selections
)
+ async def _validate_composed_capture_mode(
+ self, db: AsyncSession, sync_id: UUID, ctx: BaseContext, capture_originals: bool
+ ) -> None:
+ """Never reinterpret persisted IDs or cursors when selecting an original capture adapter."""
+ scope = (Entity.sync_id == sync_id, Entity.organization_id == ctx.organization.id)
+ legacy, canonical, scanned = (
+ await db.execute(
+ select(
+ select(Entity.id).where(*scope, Entity.record_revision == 0).exists(),
+ select(Entity.id).where(*scope, Entity.record_revision > 0).exists(),
+ select(CaptureScan.id)
+ .where(
+ CaptureScan.sync_id == sync_id,
+ CaptureScan.organization_id == ctx.organization.id,
+ )
+ .exists(),
+ )
+ )
+ ).one()
+ cursor = await db.scalar(
+ select(StoredSyncCursor.cursor_data).where(
+ StoredSyncCursor.sync_id == sync_id,
+ StoredSyncCursor.organization_id == ctx.organization.id,
+ )
+ )
+ cursor = cursor or {}
+ canonical_keys = {"canonical_cycle", "canonical_checkpoint"}
+ has_canonical_cursor = bool(canonical_keys.intersection(cursor))
+ has_legacy_cursor = any(
+ value is not None and value != {} and value != []
+ for key, value in cursor.items()
+ if key not in canonical_keys
+ )
+ if (capture_originals and (legacy or has_legacy_cursor)) or (
+ not capture_originals and (canonical or scanned or has_canonical_cursor)
+ ):
+ raise ValueError(
+ "Original capture mode differs from retained state; create a new source"
+ )
+
async def _build_arf_replay_source(
self,
db: AsyncSession,
diff --git a/backend/airweave/domains/sync_pipeline/file_types.py b/backend/airweave/domains/sync_pipeline/file_types.py
index 1f3f44af70..43b0fbb0e1 100644
--- a/backend/airweave/domains/sync_pipeline/file_types.py
+++ b/backend/airweave/domains/sync_pipeline/file_types.py
@@ -3,7 +3,7 @@
Single source of truth for supported file types across downloader and pipeline.
"""
-_DOCUMENT_EXTENSIONS = {".pdf", ".doc", ".docx", ".pptx"}
+_DOCUMENT_EXTENSIONS = {".pdf", ".doc", ".docx", ".pptx", ".xlsx"}
_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
_HTML_EXTENSIONS = {".html", ".htm"}
_TEXT_EXTENSIONS = {".txt", ".json", ".xml", ".md", ".yaml", ".yml", ".toml"}
diff --git a/backend/airweave/domains/sync_pipeline/orchestrator.py b/backend/airweave/domains/sync_pipeline/orchestrator.py
index d8e30cc928..888b4064c4 100644
--- a/backend/airweave/domains/sync_pipeline/orchestrator.py
+++ b/backend/airweave/domains/sync_pipeline/orchestrator.py
@@ -19,6 +19,7 @@
from airweave.db.session import get_db_context
from airweave.domains.access_control.pipeline import AccessControlPipeline
from airweave.domains.sources.exceptions.classifier import classify_error
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
from airweave.domains.sync_pipeline.contexts import SyncContext
from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
from airweave.domains.sync_pipeline.entity.pipeline import EntityPipeline
@@ -53,9 +54,9 @@ class SyncOrchestrator:
def __init__(
self,
- entity_pipeline: EntityPipeline,
+ entity_pipeline: EntityPipeline | CanonicalCapturePipeline,
worker_pool: AsyncWorkerPool,
- stream: AsyncSourceStream,
+ stream: AsyncSourceStream | None,
sync_context: SyncContext,
runtime: SyncRuntime,
access_control_pipeline: AccessControlPipeline,
@@ -196,17 +197,29 @@ async def _start_sync(self) -> None:
"""Initialize sync job and start all components."""
self.sync_context.logger.info("Starting sync job")
- await self.stream.start()
-
await self._state_machine.transition(
sync_job_id=self.sync_context.sync_job.id,
target=SyncJobStatus.RUNNING,
ctx=self.sync_context,
lifecycle_data=self._lifecycle_data,
)
+ if self.runtime.canonical_capture is not None:
+ await self.runtime.canonical_capture.start(self.sync_context)
+ if self.stream is not None:
+ await self.stream.start()
async def _process_entities(self) -> None: # noqa: C901
"""Process entities using micro-batching with bounded inner concurrency."""
+ if (
+ self.runtime.canonical_capture is not None
+ and self.runtime.canonical_capture.page_source is not None
+ ):
+ await self.runtime.canonical_capture.run_scans(
+ self.sync_context, self.runtime, self._check_capture_limits
+ )
+ return
+ if self.stream is None:
+ raise SyncFailureError("Observation source has no stream")
source_name = self.runtime.source.source_name
self.sync_context.logger.info(
f"Starting pull-based processing from source {source_name} "
@@ -297,6 +310,14 @@ async def _process_entities(self) -> None: # noqa: C901
if stream_error:
raise stream_error
+ async def _check_capture_limits(self) -> None:
+ """Keep the existing usage guard before each durable provider page."""
+ if not self.sync_context.execution_config.behavior.skip_guardrails:
+ async with get_db_context() as db:
+ await self._usage_checker.is_allowed(
+ db, self.sync_context.organization.id, ActionType.ENTITIES
+ )
+
async def _submit_batch_and_trim(
self,
batch: list,
@@ -306,6 +327,15 @@ async def _submit_batch_and_trim(
if not batch:
return pending_tasks
+ if self.runtime.canonical_capture is not None:
+ # Preserve per-identity observation order. Source I/O remains concurrent.
+ await self.runtime.canonical_capture.process(
+ entities=list(batch),
+ sync_context=self.sync_context,
+ runtime=self.runtime,
+ )
+ return pending_tasks
+
task = await self.worker_pool.submit(
self.entity_pipeline.process,
entities=list(batch),
@@ -449,6 +479,12 @@ async def _finalize_stream_and_tasks(
async def _cleanup_orphaned_entities_if_needed(self) -> None:
"""Cleanup orphaned entities based on sync type."""
+ if self.runtime.canonical_capture is not None:
+ await self.runtime.canonical_capture.cleanup_orphaned_entities(
+ self.sync_context,
+ self.runtime,
+ )
+ return
cursor = self.runtime.cursor
is_incremental = cursor is not None and cursor.loaded_from_db
@@ -508,6 +544,8 @@ async def _process_access_control_memberships(self) -> None:
f"ACL sync error: {get_error_message(e)}",
exc_info=True,
)
+ if self.runtime.canonical_capture is not None:
+ raise
await self._publish_acl_heartbeat()
# Don't fail the entire sync for ACL errors
@@ -630,6 +668,9 @@ async def _update_snapshot_short_name(self) -> None:
async def _save_cursor_data(self) -> None:
"""Save cursor data to database if it exists."""
+ if self.runtime.canonical_capture is not None:
+ await self.runtime.canonical_capture.save_checkpoint(self.sync_context, self.runtime)
+ return
if self.runtime.cursor is None:
return
@@ -664,6 +705,8 @@ async def _save_cursor_data(self) -> None:
f"Failed to save cursor data for sync {self.sync_context.sync.id}: {e}",
exc_info=True,
)
+ # A run without its durable continuation boundary is not complete.
+ raise
async def _handle_sync_failure(
self, error: Exception
@@ -755,7 +798,8 @@ async def _handle_cancellation(self) -> None:
await self.worker_pool.cancel_all()
# Cancel stream to stop producer
- await self.stream.cancel()
+ if self.stream is not None:
+ await self.stream.cancel()
# Transition through CANCELLING → CANCELLED.
# RUNNING → CANCELLING is required by the state machine.
diff --git a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py
index 75a6d46901..dc801b0bea 100644
--- a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py
+++ b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py
@@ -9,11 +9,18 @@
from airweave.domains.converters.protocols import ConverterRegistryProtocol
from airweave.domains.sync_pipeline.exceptions import EntityProcessingError, SyncFailureError
from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS
+from airweave.domains.sync_pipeline.pipeline.text_models import (
+ BuiltText,
+ BuiltTextBatch,
+ NativeTextBody,
+)
from airweave.platform.entities._base import BaseEntity, CodeFileEntity, FileEntity, WebEntity
if TYPE_CHECKING:
- from airweave.domains.sync_pipeline.contexts import SyncContext
- from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
+ from airweave.domains.sync_pipeline.processors.context import (
+ ProcessingContext,
+ ProcessingRuntime,
+ )
class TextualRepresentationBuilder:
@@ -43,8 +50,8 @@ def __init__(self, converter_registry: Optional[ConverterRegistryProtocol] = Non
async def build_for_batch(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
) -> List[BaseEntity]:
"""Build textual_representation for all entities in batch.
@@ -60,47 +67,109 @@ async def build_for_batch(
Modifies entities in-place, setting textual_representation.
Failed entities are removed and counted as skipped.
"""
+ return (await self.build_with_text(entities, sync_context, runtime)).entities
+
+ async def build_with_text(
+ self,
+ entities: List[BaseEntity],
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
+ *,
+ native_bodies: dict[str, NativeTextBody] | None = None,
+ ) -> BuiltTextBatch:
+ """Retain exact converted text and content boundaries before chunking clears them."""
+ native_bodies = native_bodies or {}
+ if set(native_bodies) - {entity.entity_id for entity in entities}:
+ raise EntityProcessingError("Native body does not belong to this text batch")
+ if any(
+ entity.entity_id in native_bodies and isinstance(entity, (FileEntity, WebEntity))
+ for entity in entities
+ ):
+ raise EntityProcessingError("Native body cannot override converted content")
+ content_starts: dict[str, int] = {}
source_name = sync_context.source_short_name
# Step 1: Build metadata section for all entities
- await self._build_metadata_for_all(entities, source_name)
+ await self._build_metadata_for_all(entities, source_name, native_bodies)
+ for entity in entities:
+ body = native_bodies.get(entity.entity_id)
+ if body is not None:
+ prefix = f"{entity.textual_representation}\n\n# Content\n\n"
+ content_starts[entity.entity_id] = len(prefix)
+ entity.textual_representation = prefix + body.text
# Step 2: Partition entities by converter
converter_groups, failed_entities = self._partition_by_converter(entities, sync_context)
# Step 3: Convert each partition
- additional_failures = await self._convert_partitions(converter_groups, sync_context)
+ additional_failures = await self._convert_partitions(
+ converter_groups, sync_context, content_starts
+ )
failed_entities.extend(additional_failures)
# Step 4: Handle failures
await self._handle_conversion_failures(entities, failed_entities, sync_context, runtime)
- return entities
+ return BuiltTextBatch(
+ entities=entities,
+ representations=tuple(
+ BuiltText(
+ entity_id=entity.entity_id,
+ text=entity.textual_representation or "",
+ content_start=content_starts.get(entity.entity_id),
+ kind=(
+ "native_text"
+ if entity.entity_id in native_bodies
+ else "extracted_text"
+ if entity.entity_id in content_starts
+ else "generated_text"
+ ),
+ )
+ for entity in entities
+ ),
+ )
# ------------------------------------------------------------------------------------
# Metadata Building
# ------------------------------------------------------------------------------------
- async def _build_metadata_for_all(self, entities: List[BaseEntity], source_name: str) -> None:
+ async def _build_metadata_for_all(
+ self,
+ entities: List[BaseEntity],
+ source_name: str,
+ native_bodies: dict[str, NativeTextBody],
+ ) -> None:
"""Build metadata section for all entities.
Args:
entities: Entities to build metadata for
source_name: Name of the source connector
+ native_bodies: Explicit source bodies whose fields leave the metadata section
Note:
CodeFileEntity is exempt because code is self-documenting.
"""
async def build_metadata(entity: BaseEntity):
- metadata = self.build_metadata_section(entity, source_name)
+ body = native_bodies.get(entity.entity_id)
+ metadata = self.build_metadata_section(
+ entity,
+ source_name,
+ exclude_fields=body.metadata_fields if body else (),
+ )
if not metadata and not isinstance(entity, CodeFileEntity):
raise EntityProcessingError(f"Empty metadata for {entity.entity_id}")
entity.textual_representation = metadata
await asyncio.gather(*[build_metadata(e) for e in entities])
- def build_metadata_section(self, entity: BaseEntity, source_name: str) -> str:
+ def build_metadata_section(
+ self,
+ entity: BaseEntity,
+ source_name: str,
+ *,
+ exclude_fields: tuple[str, ...] = (),
+ ) -> str:
"""Build metadata section for any entity type.
This method is public to allow federated search sources to build
@@ -111,6 +180,7 @@ def build_metadata_section(self, entity: BaseEntity, source_name: str) -> str:
Args:
entity: Entity to build metadata for
source_name: Name of the source (e.g., "slack", "github")
+ exclude_fields: Explicit fields represented in the separate native body
Returns:
Markdown formatted metadata section
@@ -134,7 +204,11 @@ def build_metadata_section(self, entity: BaseEntity, source_name: str) -> str:
lines.append(f"**Path**: {path_str}")
# Add embeddable fields
- embeddable_fields = self._extract_embeddable_fields(entity)
+ embeddable_fields = {
+ key: value
+ for key, value in self._extract_embeddable_fields(entity).items()
+ if key not in exclude_fields
+ }
if embeddable_fields:
lines.append("")
lines.append(self._format_embeddable_fields_as_markdown(embeddable_fields))
@@ -247,7 +321,7 @@ def _format_value(self, value: Any) -> str:
def _partition_by_converter(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> Tuple[Dict[Any, List[Tuple[BaseEntity, str]]], List[BaseEntity]]:
"""Partition entities by their converter type.
@@ -313,13 +387,15 @@ def _get_converter_and_key(self, entity: BaseEntity) -> Tuple[Any, Optional[str]
async def _convert_partitions(
self,
converter_groups: Dict[Any, List[Tuple[BaseEntity, str]]],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
+ content_starts: dict[str, int],
) -> List[BaseEntity]:
"""Execute batch conversion for each converter group.
Args:
converter_groups: Dict mapping converter to (entity, key) tuples
sync_context: Sync context for logging
+ content_starts: Per-entity converter-content offsets to populate
Returns:
List of entities that failed conversion
@@ -331,7 +407,9 @@ async def _convert_partitions(
for i in range(0, len(entity_key_pairs), batch_size):
sub_batch = entity_key_pairs[i : i + batch_size]
- failures = await self._convert_sub_batch(converter, sub_batch, sync_context)
+ failures = await self._convert_sub_batch(
+ converter, sub_batch, sync_context, content_starts
+ )
failed_entities.extend(failures)
return failed_entities
@@ -340,7 +418,8 @@ async def _convert_sub_batch(
self,
converter: Any,
sub_batch: List[Tuple[BaseEntity, str]],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
+ content_starts: dict[str, int],
) -> List[BaseEntity]:
"""Convert a sub-batch of entities using the given converter.
@@ -348,6 +427,7 @@ async def _convert_sub_batch(
converter: Converter module to use
sub_batch: List of (entity, key) tuples
sync_context: Sync context for logging
+ content_starts: Per-entity converter-content offsets to populate
Returns:
List of entities that failed conversion
@@ -372,7 +452,9 @@ async def _convert_sub_batch(
failed_entities.append(entity)
continue
- entity.textual_representation += f"\n\n# Content\n\n{text_content}"
+ prefix = f"{entity.textual_representation}\n\n# Content\n\n"
+ content_starts[entity.entity_id] = len(prefix)
+ entity.textual_representation = prefix + text_content
except SyncFailureError:
# Infrastructure failure - propagate to fail entire sync
@@ -403,8 +485,8 @@ async def _handle_conversion_failures(
self,
entities: List[BaseEntity],
failed_entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
) -> None:
"""Remove failed entities and update progress.
diff --git a/backend/airweave/domains/sync_pipeline/pipeline/text_models.py b/backend/airweave/domains/sync_pipeline/pipeline/text_models.py
new file mode 100644
index 0000000000..3f2e0d3c5e
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/pipeline/text_models.py
@@ -0,0 +1,33 @@
+"""One conversion result, shared by indexing and optional retained representation reads."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from airweave.platform.entities._base import BaseEntity
+
+
+class NativeTextBody(BaseModel):
+ """Source-selected body; an empty string is present content, not missing content."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+ text: str
+ metadata_fields: tuple[str, ...] = ()
+
+
+class BuiltText(BaseModel):
+ """Content starts at a known construction boundary, never a parsed delimiter."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ entity_id: str
+ text: str
+ content_start: int | None = Field(default=None, ge=0)
+ kind: Literal["native_text", "extracted_text", "generated_text"] = "generated_text"
+
+
+class BuiltTextBatch(BaseModel):
+ """Entities retain the same strings that were passed to the chunker."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ entities: list[BaseEntity]
+ representations: tuple[BuiltText, ...]
diff --git a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py
index 237d929e24..af0f96ac56 100644
--- a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py
+++ b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py
@@ -19,12 +19,15 @@
from airweave.domains.embedders.protocols import DenseEmbedderProtocol, SparseEmbedderProtocol
from airweave.domains.sync_pipeline.exceptions import EntityProcessingError, SyncFailureError
from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder
+from airweave.domains.sync_pipeline.pipeline.text_models import BuiltTextBatch, NativeTextBody
from airweave.domains.sync_pipeline.processors.utils import filter_empty_representations
from airweave.platform.entities._base import BaseEntity, CodeFileEntity
if TYPE_CHECKING:
- from airweave.domains.sync_pipeline.contexts import SyncContext
- from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
+ from airweave.domains.sync_pipeline.processors.context import (
+ ProcessingContext,
+ ProcessingRuntime,
+ )
class ChunkEmbedProcessor:
@@ -37,27 +40,71 @@ def __init__(
sparse_embedder: SparseEmbedderProtocol,
) -> None:
"""Initialize with converter registry and embedding providers."""
+ self._converter_registry = converter_registry
self._text_builder = TextualRepresentationBuilder(converter_registry)
self._dense_embedder = dense_embedder
self._sparse_embedder = sparse_embedder
+ def supports_file_extension(self, extension: str) -> bool:
+ """Whether this processor has a configured converter for these bytes."""
+ return self._converter_registry.for_extension(extension) is not None
+
async def process(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
+ *,
+ strict: bool = False,
) -> List[BaseEntity]:
"""Process entities through full chunk+embed pipeline."""
if not entities:
return []
+ expected_ids = {entity.entity_id for entity in entities}
+ if strict and len(expected_ids) != len(entities):
+ raise EntityProcessingError("Projection mapper returned duplicate part identities")
+
# Step 1: Build textual representations
processed = await self._text_builder.build_for_batch(entities, sync_context, runtime)
+ return await self.process_built_text(
+ processed, sync_context, runtime, strict=strict, expected_ids=expected_ids
+ )
+
+ async def build_text(
+ self,
+ entities: List[BaseEntity],
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
+ *,
+ native_bodies: dict[str, NativeTextBody] | None = None,
+ ) -> BuiltTextBatch:
+ """Convert once, exposing the complete pre-chunk representation for retention."""
+ return await self._text_builder.build_with_text(
+ entities, sync_context, runtime, native_bodies=native_bodies
+ )
+
+ async def process_built_text(
+ self,
+ processed: List[BaseEntity],
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
+ *,
+ strict: bool = False,
+ expected_ids: set[str] | None = None,
+ ) -> List[BaseEntity]:
+ """Chunk already-built text without rerunning converters."""
+ if expected_ids is None:
+ expected_ids = {entity.entity_id for entity in processed}
+ if strict and len(expected_ids) != len(processed):
+ raise EntityProcessingError("Projection conversion dropped required content")
# Step 2: Filter empty representations
processed = await filter_empty_representations(
processed, sync_context, runtime, "ChunkEmbed"
)
+ if strict and {entity.entity_id for entity in processed} != expected_ids:
+ raise EntityProcessingError("Projection conversion dropped required content")
if not processed:
sync_context.logger.debug("[ChunkEmbedProcessor] No entities after text building")
return []
@@ -65,15 +112,25 @@ async def process(
# Step 3: Chunk entities
chunk_entities = await self._chunk_entities(processed, sync_context, runtime)
+ if (
+ strict
+ and {entity.airweave_system_metadata.original_entity_id for entity in chunk_entities}
+ != expected_ids
+ ):
+ raise EntityProcessingError("Projection chunking dropped required content")
+ expected_chunks = {entity.entity_id for entity in chunk_entities}
+
# Step 4: Release parent text (memory optimization)
for entity in processed:
entity.textual_representation = None
# Step 5: Embed chunks (may remove entities that fail embedding)
chunk_entities = await self._embed_entities(chunk_entities, sync_context)
+ if strict and {entity.entity_id for entity in chunk_entities} != expected_chunks:
+ raise EntityProcessingError("Projection embedding dropped required chunks")
sync_context.logger.debug(
- f"[ChunkEmbedProcessor] {len(entities)} entities -> {len(chunk_entities)} chunks"
+ f"[ChunkEmbedProcessor] {len(processed)} entities -> {len(chunk_entities)} chunks"
)
return chunk_entities
@@ -85,8 +142,8 @@ async def process(
async def _chunk_entities(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
) -> List[BaseEntity]:
"""Route entities to appropriate chunker."""
code_entities = [e for e in entities if isinstance(e, CodeFileEntity)]
@@ -107,8 +164,8 @@ async def _chunk_entities(
async def _chunk_code_entities(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
) -> List[BaseEntity]:
"""Chunk code with AST-aware CodeChunker."""
from airweave.platform.chunkers.code import CodeChunker
@@ -134,7 +191,7 @@ async def _chunk_code_entities(
async def _chunk_textual_entities(
self,
entities: List[BaseEntity],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> List[BaseEntity]:
"""Chunk text with SemanticChunker."""
from airweave.platform.chunkers.semantic import SemanticChunker
@@ -180,7 +237,7 @@ def _multiply_entities(
self,
entities: List[BaseEntity],
chunk_lists: List[List[Dict[str, Any]]],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> List[BaseEntity]:
"""Create chunk entities from chunker output."""
chunk_entities: List[BaseEntity] = []
@@ -213,7 +270,7 @@ def _multiply_entities(
async def _embed_entities(
self,
chunk_entities: List[BaseEntity],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> List[BaseEntity]:
"""Compute dense and sparse embeddings for all destinations.
@@ -279,7 +336,7 @@ def _validate_dense_dimensions(self, dense_results: List[Any]) -> None:
async def _dense_embed_with_fallback(
self,
chunk_entities: List[BaseEntity],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> Tuple[List[Any], List[BaseEntity]]:
"""Run dense embedding with fallback to individual embedding on failure.
@@ -325,7 +382,7 @@ async def _dense_embed_with_fallback(
async def _embed_individually(
self,
chunk_entities: List[BaseEntity],
- sync_context: "SyncContext",
+ sync_context: "ProcessingContext",
) -> Tuple[List[Any], List[BaseEntity]]:
"""Embed entities one-by-one, skipping any that fail.
diff --git a/backend/airweave/domains/sync_pipeline/processors/context.py b/backend/airweave/domains/sync_pipeline/processors/context.py
new file mode 100644
index 0000000000..94ba7bc3a1
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/processors/context.py
@@ -0,0 +1,26 @@
+"""Minimal processing dependencies shared by capture-era and replay processing."""
+
+from typing import Protocol
+
+from airweave.core.logging import ContextualLogger
+
+
+class ProcessingContext(Protocol):
+ """Identity and diagnostics needed by text conversion/chunking."""
+
+ logger: ContextualLogger
+ source_short_name: str
+
+
+class ProcessingTracker(Protocol):
+ """Existing processing skip accounting interface."""
+
+ async def record_skipped(self, count: int) -> None:
+ """Record that input could not be converted."""
+ ...
+
+
+class ProcessingRuntime(Protocol):
+ """Processing never requires a live source or its credentials."""
+
+ entity_tracker: ProcessingTracker
diff --git a/backend/airweave/domains/sync_pipeline/processors/entity_fields.py b/backend/airweave/domains/sync_pipeline/processors/entity_fields.py
new file mode 100644
index 0000000000..eb6d75322a
--- /dev/null
+++ b/backend/airweave/domains/sync_pipeline/processors/entity_fields.py
@@ -0,0 +1,26 @@
+"""Normalize declared entity display fields before either indexing pipeline."""
+
+from airweave.core.shared_models import AirweaveFieldFlag
+from airweave.platform.entities._base import BaseEntity
+
+
+def populate_base_fields(entity: BaseEntity) -> None:
+ """Copy explicitly flagged source fields without serializing native capture JSON."""
+ flags = {
+ "entity_id": AirweaveFieldFlag.IS_ENTITY_ID,
+ "name": AirweaveFieldFlag.IS_NAME,
+ "created_at": AirweaveFieldFlag.IS_CREATED_AT,
+ "updated_at": AirweaveFieldFlag.IS_UPDATED_AT,
+ }
+ for target, flag in flags.items():
+ if getattr(entity, target):
+ continue
+ for name, field in type(entity).model_fields.items():
+ metadata = field.json_schema_extra
+ if isinstance(metadata, dict) and metadata.get(flag.value):
+ value = getattr(entity, name)
+ if value is not None:
+ setattr(
+ entity, target, str(value) if target in ("entity_id", "name") else value
+ )
+ break
diff --git a/backend/airweave/domains/sync_pipeline/processors/utils.py b/backend/airweave/domains/sync_pipeline/processors/utils.py
index ddef0c3643..4dde762ddd 100644
--- a/backend/airweave/domains/sync_pipeline/processors/utils.py
+++ b/backend/airweave/domains/sync_pipeline/processors/utils.py
@@ -5,14 +5,16 @@
from airweave.platform.entities._base import BaseEntity
if TYPE_CHECKING:
- from airweave.domains.sync_pipeline.contexts import SyncContext
- from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
+ from airweave.domains.sync_pipeline.processors.context import (
+ ProcessingContext,
+ ProcessingRuntime,
+ )
async def filter_empty_representations(
entities: List[BaseEntity],
- sync_context: "SyncContext",
- runtime: "SyncRuntime",
+ sync_context: "ProcessingContext",
+ runtime: "ProcessingRuntime",
processor_name: str = "Processor",
) -> List[BaseEntity]:
"""Filter entities with empty textual_representation.
diff --git a/backend/airweave/domains/sync_pipeline/protocols.py b/backend/airweave/domains/sync_pipeline/protocols.py
index 6969152114..17b65947c5 100644
--- a/backend/airweave/domains/sync_pipeline/protocols.py
+++ b/backend/airweave/domains/sync_pipeline/protocols.py
@@ -12,6 +12,7 @@
if TYPE_CHECKING:
from airweave.core.context import BaseContext
+ from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
from airweave.domains.sync_pipeline.config import SyncConfig
from airweave.domains.sync_pipeline.contexts import SyncContext
from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime
@@ -105,6 +106,7 @@ async def create_orchestrator(
force_full_sync: bool = False,
execution_config: Optional[SyncConfig] = None,
access_token: Optional[str] = None,
+ capture_attempt: CaptureAttempt | None = None,
) -> SyncOrchestrator:
"""Create and return a fully-wired SyncOrchestrator."""
...
diff --git a/backend/airweave/domains/sync_pipeline/stream.py b/backend/airweave/domains/sync_pipeline/stream.py
index 5c9f7e9bd1..6d32874456 100644
--- a/backend/airweave/domains/sync_pipeline/stream.py
+++ b/backend/airweave/domains/sync_pipeline/stream.py
@@ -5,10 +5,9 @@
from typing import AsyncGenerator, Generic, Optional, TypeVar
from airweave.core.logging import ContextualLogger
-from airweave.platform.entities._base import BaseEntity
from airweave.platform.utils.error_utils import get_error_message
-T = TypeVar("T", bound=BaseEntity)
+T = TypeVar("T")
class StreamState(Enum):
diff --git a/backend/airweave/domains/sync_pipeline/tests/test_factory.py b/backend/airweave/domains/sync_pipeline/tests/test_factory.py
index 21d0ec11e2..8be28460c4 100644
--- a/backend/airweave/domains/sync_pipeline/tests/test_factory.py
+++ b/backend/airweave/domains/sync_pipeline/tests/test_factory.py
@@ -6,6 +6,8 @@
import pytest
+from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline
+from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt
from airweave.domains.sync_pipeline.factory import SourceBuildResult, SyncFactory
@@ -115,7 +117,14 @@ async def test_create_orchestrator_raises_when_source_connection_missing():
ctx.organization.id = uuid4()
db = AsyncMock()
- with pytest.raises(NotFoundException, match="Source connection record not found"):
+ with (
+ patch(
+ "airweave.domains.sync_pipeline.factory.CanonicalRecordStore.admit_job",
+ new_callable=AsyncMock,
+ return_value=0,
+ ),
+ pytest.raises(NotFoundException, match="Source connection record not found"),
+ ):
await factory.create_orchestrator(
db=db,
sync=sync,
@@ -132,7 +141,8 @@ async def test_create_orchestrator_raises_when_source_connection_missing():
@pytest.mark.asyncio
-async def test_create_orchestrator_passes_entity_repo_to_pipeline():
+@pytest.mark.parametrize("canonical", [False, True, "page", "composed"])
+async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical):
"""entity_repo is forwarded to EntityActionResolver and EntityPipeline."""
entity_repo = MagicMock()
sc_repo = MagicMock()
@@ -158,6 +168,11 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline():
db = AsyncMock()
with (
+ patch(
+ "airweave.domains.sync_pipeline.factory.CanonicalRecordStore.admit_job",
+ new_callable=AsyncMock,
+ return_value=0,
+ ) as admit,
patch("airweave.domains.sync_pipeline.factory.SyncContextBuilder") as mock_sc_builder,
patch(
"airweave.domains.sync_pipeline.factory.EntityDispatcherBuilder"
@@ -176,7 +191,35 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline():
) as mock_build_tracker,
):
mock_source = MagicMock()
+ mock_source.capture_page_source = None
mock_source.generate_entities = MagicMock(return_value=AsyncMock())
+ if canonical in {"page", "composed"}:
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.platform.configs.config import SlackConfig
+ from airweave.platform.sources.slack import SlackSource
+
+ mock_source = SlackSource(
+ auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=MagicMock()
+ )
+ mock_source.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U1")
+ mock_source._get = AsyncMock(
+ return_value={"ok": True, "team_id": "T1", "user_id": "U1"}
+ )
+ await mock_source.validate()
+ page_source = mock_source
+ if canonical == "composed":
+ from airweave.platform.sources.outlook_mail import OutlookMailSource
+
+ mock_source = OutlookMailSource(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=MagicMock(),
+ )
+ mock_source._capture_page_source = page_source
+ mock_source.generate_entities = MagicMock()
+ elif canonical:
+ mock_source.canonical_record_types = ("event",)
+ mock_source.generate_observations = MagicMock(return_value=AsyncMock())
mock_build_source.return_value = SourceBuildResult(
source=mock_source, cursor=MagicMock(), files=MagicMock(), node_selections=[]
)
@@ -194,11 +237,28 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline():
collection=collection,
connection=connection,
ctx=ctx,
+ capture_attempt=CaptureAttempt(id=uuid4(), number=1),
)
assert orchestrator is not None
- assert orchestrator.entity_pipeline._entity_repo is entity_repo
- assert orchestrator.entity_pipeline._resolver._entity_repo is entity_repo
+ assert admit.await_count == 2
+ db.refresh.assert_awaited_once_with(sc)
+ if canonical:
+ assert isinstance(orchestrator.entity_pipeline, CanonicalCapturePipeline)
+ mock_build_destinations.assert_not_called()
+ mock_disp_builder.assert_not_called()
+ if canonical in {"page", "composed"}:
+ assert orchestrator.stream is None
+ assert orchestrator.entity_pipeline.page_source is page_source
+ assert orchestrator.runtime.source is mock_source
+ if canonical == "composed":
+ mock_source.generate_entities.assert_not_called()
+ else:
+ mock_source.generate_entities.assert_not_called()
+ mock_source.generate_observations.assert_called_once()
+ else:
+ assert orchestrator.entity_pipeline._entity_repo is entity_repo
+ assert orchestrator.entity_pipeline._resolver._entity_repo is entity_repo
# ---------------------------------------------------------------------------
@@ -358,8 +418,11 @@ async def test_raises_not_found_when_no_arf_data(self):
@pytest.mark.asyncio
async def test_success_returns_source_build_result(self):
- mock_source = AsyncMock()
+ from airweave.domains.arf.replay_source import ArfReplaySource
+
+ mock_source = ArfReplaySource(sync_id=uuid4(), storage=MagicMock())
mock_source.validate = AsyncMock(return_value=True)
+ assert mock_source.capture_page_source is None
factory = _build_factory()
db = AsyncMock()
@@ -397,10 +460,12 @@ def test_passes_for_snapshot_short_name(self):
def test_passes_for_normal_source_with_invalid_snapshot_config(self):
"""Normal source with config that fails SnapshotConfig parsing → no exception."""
sc = SimpleNamespace(short_name="github", config_fields={"not": "snapshot"}, name="gh")
- SyncFactory._validate_not_completed_snapshot(sc) # should not raise (ValidationError caught)
+ SyncFactory._validate_not_completed_snapshot(
+ sc
+ ) # should not raise (ValidationError caught)
def test_raises_for_restored_snapshot_source(self):
- """Source with non-snapshot short_name but valid SnapshotConfig fields → SyncFailureError."""
+ """Restored snapshot config on a normal source raises SyncFailureError."""
from airweave.domains.sync_pipeline.exceptions import SyncFailureError
sc = SimpleNamespace(short_name="github", config_fields={}, name="restored-snap")
@@ -527,9 +592,7 @@ async def test_returns_tracker_with_initial_counts(self):
entity_count_repo.get_counts_per_sync_and_type = AsyncMock(return_value=[count_row])
factory = _build_factory(entity_count_repo=entity_count_repo)
- tracker = await factory._build_entity_tracker(
- db=db, sync=sync, sync_job=sync_job, ctx=ctx
- )
+ tracker = await factory._build_entity_tracker(db=db, sync=sync, sync_job=sync_job, ctx=ctx)
assert tracker is not None
assert tracker.job_id == sync_job.id
@@ -545,9 +608,7 @@ async def test_returns_tracker_with_empty_counts(self):
entity_count_repo.get_counts_per_sync_and_type = AsyncMock(return_value={})
factory = _build_factory(entity_count_repo=entity_count_repo)
- tracker = await factory._build_entity_tracker(
- db=db, sync=sync, sync_job=sync_job, ctx=ctx
- )
+ tracker = await factory._build_entity_tracker(db=db, sync=sync, sync_job=sync_job, ctx=ctx)
assert tracker is not None
@@ -584,3 +645,49 @@ async def test_converts_rows_to_node_selection_data(self):
assert len(result) == 1
assert result[0].source_node_id == "node-1"
assert result[0].node_type == "file"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("during_build", [False, True])
+async def test_generation_admission_stops_stale_factory(during_build):
+ """Stale workflow schema zero never bypasses authoritative admission."""
+ from airweave.domains.entities.canonical.store import StaleWriter
+
+ sc = MagicMock()
+ sc_repo = MagicMock(get_by_sync_id=AsyncMock(return_value=sc))
+ factory = _build_factory(sc_repo=sc_repo)
+ sync = SimpleNamespace(id=uuid4(), sync_config=None, provisioning_generation=0)
+ job = SimpleNamespace(id=uuid4(), sync_config=None, provisioning_generation=0)
+ db = AsyncMock()
+ order = []
+ db.refresh.side_effect = lambda _: order.append("refresh")
+
+ async def build(**kwargs):
+ order.append("build")
+ return MagicMock()
+
+ with (
+ patch(
+ "airweave.domains.sync_pipeline.factory.CanonicalRecordStore.admit_job",
+ new_callable=AsyncMock,
+ ) as admit,
+ patch.object(factory, "_build_source", side_effect=build) as source,
+ patch.object(factory, "_build_destinations", new_callable=AsyncMock) as destinations,
+ patch.object(factory, "_build_entity_tracker", new_callable=AsyncMock) as tracker,
+ ):
+ # Even if both DB generations are individually eligible, switching between
+ # checks must reject the source constructed under the first generation.
+ admit.side_effect = [4, 5] if during_build else StaleWriter("Stale job")
+ with pytest.raises(StaleWriter):
+ await factory.create_orchestrator(
+ db=db,
+ sync=sync,
+ sync_job=job,
+ collection=SimpleNamespace(sync_config=None),
+ connection=MagicMock(),
+ ctx=_make_ctx(),
+ )
+ assert source.await_count == int(during_build)
+ assert order == (["refresh", "build"] if during_build else [])
+ destinations.assert_not_awaited()
+ tracker.assert_not_awaited()
diff --git a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py
index cec2075275..e0a06bf20f 100644
--- a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py
+++ b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py
@@ -42,7 +42,7 @@ def _make_orchestrator(**overrides):
worker_pool=worker_pool,
stream=overrides.pop("stream", MagicMock()),
sync_context=sync_context,
- runtime=overrides.pop("runtime", MagicMock()),
+ runtime=overrides.pop("runtime", MagicMock(canonical_capture=None)),
access_control_pipeline=overrides.pop("access_control_pipeline", MagicMock()),
event_bus=overrides.pop("event_bus", MagicMock()),
usage_checker=usage_checker,
@@ -468,3 +468,32 @@ async def _fake_handler(_e: Exception):
):
with pytest.raises(RuntimeError, match="redis down"):
await orc.run()
+
+
+async def test_checkpoint_failure_prevents_completed_transition():
+ """A successfully processed stream is incomplete if checkpoint persistence fails."""
+ service = MagicMock()
+ service.create_or_update_cursor = AsyncMock(side_effect=ConnectionError("database unavailable"))
+ state_machine = MagicMock()
+ state_machine.transition = AsyncMock()
+ ctx = _make_sync_context()
+ ctx.execution_config = None
+ runtime = MagicMock()
+ runtime.canonical_capture = None
+ runtime.cursor = SimpleNamespace(cursor_data={"history_id": "42"}, cursor_field=None)
+ orch = _make_orchestrator(
+ sync_context=ctx,
+ runtime=runtime,
+ sync_cursor_service=service,
+ state_machine=state_machine,
+ )
+ db_context = MagicMock()
+ db_context.__aenter__ = AsyncMock(return_value=MagicMock())
+ db_context.__aexit__ = AsyncMock(return_value=False)
+ with patch(
+ "airweave.domains.sync_pipeline.orchestrator.get_db_context", return_value=db_context
+ ):
+ with pytest.raises(ConnectionError, match="database unavailable"):
+ await orch._complete_sync()
+ service.create_or_update_cursor.assert_awaited_once()
+ state_machine.transition.assert_not_awaited()
diff --git a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator_coverage.py b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator_coverage.py
index 2bf4a0a241..0d63dcb826 100644
--- a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator_coverage.py
+++ b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator_coverage.py
@@ -114,6 +114,7 @@ async def empty_stream():
stream.get_entities = empty_stream
runtime = MagicMock()
+ runtime.canonical_capture = None
runtime.source.source_name = "test_source"
runtime.source.short_name = "test_source"
runtime.source.supports_access_control = False
diff --git a/backend/airweave/domains/sync_pipeline/tests/test_text_builder.py b/backend/airweave/domains/sync_pipeline/tests/test_text_builder.py
index de1193451d..0b942e65cb 100644
--- a/backend/airweave/domains/sync_pipeline/tests/test_text_builder.py
+++ b/backend/airweave/domains/sync_pipeline/tests/test_text_builder.py
@@ -4,7 +4,6 @@
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock
-from uuid import uuid4
import pytest
@@ -57,7 +56,7 @@ async def test_sync_failure_error_propagates(self):
entity = _make_entity()
with pytest.raises(SyncFailureError, match="infra down"):
- await builder._convert_sub_batch(converter, [(entity, "/path")], ctx)
+ await builder._convert_sub_batch(converter, [(entity, "/path")], ctx, {})
@pytest.mark.asyncio
async def test_entity_processing_error_skips_sub_batch(self):
@@ -67,7 +66,7 @@ async def test_entity_processing_error_skips_sub_batch(self):
ctx = FakeSyncContext()
entity = _make_entity()
- failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx)
+ failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx, {})
assert failures == [entity]
ctx.logger.warning.assert_called()
@@ -80,7 +79,7 @@ async def test_generic_exception_marks_failed_but_continues(self):
ctx = FakeSyncContext()
entity = _make_entity()
- failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx)
+ failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx, {})
assert failures == [entity]
ctx.logger.error.assert_called()
@@ -93,7 +92,7 @@ async def test_successful_conversion_appends_content(self):
ctx = FakeSyncContext()
entity = _make_entity()
- failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx)
+ failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx, {})
assert failures == []
assert "Hello world content" in entity.textual_representation
@@ -106,7 +105,7 @@ async def test_no_content_returns_entity_as_failure(self):
ctx = FakeSyncContext()
entity = _make_entity()
- failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx)
+ failures = await builder._convert_sub_batch(converter, [(entity, "/path")], ctx, {})
assert failures == [entity]
@@ -154,8 +153,8 @@ def test_file_entity_unsupported_extension_raises(self):
builder._get_converter_and_key(entity)
def test_file_entity_no_converter_for_extension_raises(self):
- from airweave.platform.entities._base import FileEntity
from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS
+ from airweave.platform.entities._base import FileEntity
supported_ext = next(iter(SUPPORTED_FILE_EXTENSIONS))
registry = MagicMock()
diff --git a/backend/airweave/domains/syncs/cursors/service.py b/backend/airweave/domains/syncs/cursors/service.py
index c362553abb..d86317ded7 100644
--- a/backend/airweave/domains/syncs/cursors/service.py
+++ b/backend/airweave/domains/syncs/cursors/service.py
@@ -32,7 +32,7 @@ async def get_cursor_data(self, db: AsyncSession, sync_id: UUID, ctx: BaseContex
return {}
except Exception as e:
logger.warning(f"Failed to load cursor data for sync {sync_id}: {e}")
- return {}
+ raise
async def get_cursor_field(
self, db: AsyncSession, sync_id: UUID, ctx: ApiContext
@@ -54,7 +54,7 @@ async def get_cursor_field(
return None
except Exception as e:
logger.warning(f"Failed to load cursor field for sync {sync_id}: {e}")
- return None
+ raise
async def create_or_update_cursor(
self,
@@ -63,7 +63,7 @@ async def create_or_update_cursor(
cursor_data: dict,
ctx: ApiContext,
cursor_field: Optional[str] = None,
- ) -> Optional[schemas.SyncCursor]:
+ ) -> schemas.SyncCursor:
"""Create or update cursor data for a sync.
Args:
@@ -74,7 +74,7 @@ async def create_or_update_cursor(
cursor_field: Optional cursor field name
Returns:
- Created or updated sync cursor, None if operation failed
+ Created or updated sync cursor. Persistence failures propagate to the caller.
"""
try:
cursor_create = schemas.SyncCursorCreate(
@@ -93,7 +93,7 @@ async def create_or_update_cursor(
except Exception as e:
logger.error(f"Failed to create/update cursor for sync {sync_id}: {e}")
- return None
+ raise
async def update_cursor_data(
self,
@@ -111,7 +111,7 @@ async def update_cursor_data(
ctx: API context
Returns:
- Updated sync cursor, None if operation failed
+ Updated sync cursor, or None if no cursor exists. Persistence failures propagate.
"""
try:
cursor = await crud.sync_cursor.update_cursor_data(
@@ -127,7 +127,7 @@ async def update_cursor_data(
except Exception as e:
logger.error(f"Failed to update cursor data for sync {sync_id}: {e}")
- return None
+ raise
async def delete_cursor(
self,
@@ -157,4 +157,4 @@ async def delete_cursor(
except Exception as e:
logger.error(f"Failed to delete cursor for sync {sync_id}: {e}")
- return False
+ raise
diff --git a/backend/airweave/domains/syncs/cursors/tests/test_failure_semantics.py b/backend/airweave/domains/syncs/cursors/tests/test_failure_semantics.py
new file mode 100644
index 0000000000..9f0b03ca0c
--- /dev/null
+++ b/backend/airweave/domains/syncs/cursors/tests/test_failure_semantics.py
@@ -0,0 +1,39 @@
+"""Persistence failure must never masquerade as an empty or saved checkpoint."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+
+from airweave.domains.syncs.cursors.service import SyncCursorService
+
+
+@pytest.mark.parametrize(
+ ("operation", "crud_operation", "extra"),
+ [
+ ("get_cursor_data", "get_by_sync_id", {}),
+ ("get_cursor_field", "get_by_sync_id", {}),
+ ("create_or_update_cursor", "create_or_update", {"cursor_data": {"position": "42"}}),
+ ("update_cursor_data", "update_cursor_data", {"cursor_data": {"position": "42"}}),
+ ("delete_cursor", "delete_by_sync_id", {}),
+ ],
+)
+async def test_database_failure_propagates(operation, crud_operation, extra):
+ failure = ConnectionError("database unavailable")
+ with patch(
+ f"airweave.domains.syncs.cursors.service.crud.sync_cursor.{crud_operation}",
+ new=AsyncMock(side_effect=failure),
+ ):
+ with pytest.raises(ConnectionError) as raised:
+ await getattr(SyncCursorService(), operation)(
+ db=MagicMock(), sync_id=uuid4(), ctx=MagicMock(), **extra
+ )
+ assert raised.value is failure
+
+
+async def test_missing_checkpoint_is_distinct_from_database_failure():
+ with patch(
+ "airweave.domains.syncs.cursors.service.crud.sync_cursor.get_by_sync_id",
+ new=AsyncMock(return_value=None),
+ ):
+ assert await SyncCursorService().get_cursor_data(MagicMock(), uuid4(), MagicMock()) == {}
diff --git a/backend/airweave/domains/syncs/fakes/repository.py b/backend/airweave/domains/syncs/fakes/repository.py
index 66e3f6ad74..6b080be045 100644
--- a/backend/airweave/domains/syncs/fakes/repository.py
+++ b/backend/airweave/domains/syncs/fakes/repository.py
@@ -74,9 +74,10 @@ async def create(
obj_in: SyncCreate,
ctx: ApiContext,
uow: Optional[UnitOfWork] = None,
+ initial_pipeline_version: int = 1,
) -> schemas.Sync:
"""Record call and return canned result."""
- self._calls.append(("create", db, obj_in, ctx, uow))
+ self._calls.append(("create", db, obj_in, ctx, uow, initial_pipeline_version))
if self._create_result is None:
raise RuntimeError("FakeSyncRepository.create_result not configured")
return self._create_result
diff --git a/backend/airweave/domains/syncs/fakes/service.py b/backend/airweave/domains/syncs/fakes/service.py
index eade30d2e8..9ebfc3b6f9 100644
--- a/backend/airweave/domains/syncs/fakes/service.py
+++ b/backend/airweave/domains/syncs/fakes/service.py
@@ -3,6 +3,7 @@
from typing import Dict, List, Optional, Tuple
from uuid import UUID
+from pydantic import JsonValue
from sqlalchemy.ext.asyncio import AsyncSession
from airweave import schemas
@@ -73,6 +74,8 @@ async def create(
run_immediately: bool,
ctx: ApiContext,
uow: UnitOfWork,
+ defer_execution: bool = False,
+ source_config: dict[str, JsonValue] | None = None,
) -> SyncProvisionResult:
self._calls.append(("create", name, source_connection_id, collection_id))
if self._should_raise:
diff --git a/backend/airweave/domains/syncs/protocols.py b/backend/airweave/domains/syncs/protocols.py
index 7d17b55fe9..5f93139e86 100644
--- a/backend/airweave/domains/syncs/protocols.py
+++ b/backend/airweave/domains/syncs/protocols.py
@@ -5,6 +5,7 @@
from typing import List, Optional, Protocol, Tuple
from uuid import UUID
+from pydantic import JsonValue
from sqlalchemy.ext.asyncio import AsyncSession
from airweave import schemas
@@ -53,6 +54,7 @@ async def create(
obj_in: SyncCreate,
ctx: ApiContext,
uow: Optional[UnitOfWork] = None,
+ initial_pipeline_version: int = 1,
) -> schemas.Sync:
"""Create a new sync with its connection associations."""
...
@@ -122,6 +124,8 @@ async def create(
run_immediately: bool,
ctx: ApiContext,
uow: UnitOfWork,
+ defer_execution: bool = False,
+ source_config: dict[str, JsonValue] | None = None,
) -> SyncProvisionResult:
"""Create sync + optional job + Temporal schedule atomically."""
...
diff --git a/backend/airweave/domains/syncs/repository.py b/backend/airweave/domains/syncs/repository.py
index e48edc2f2f..0b09870acd 100644
--- a/backend/airweave/domains/syncs/repository.py
+++ b/backend/airweave/domains/syncs/repository.py
@@ -52,9 +52,16 @@ async def create(
obj_in: SyncCreate,
ctx: ApiContext,
uow: Optional[UnitOfWork] = None,
+ initial_pipeline_version: int = 1,
) -> schemas.Sync:
"""Create a new sync with its connection associations."""
- return await crud.sync.create(db=db, obj_in=obj_in, ctx=ctx, uow=uow)
+ return await crud.sync.create(
+ db=db,
+ obj_in=obj_in,
+ ctx=ctx,
+ uow=uow,
+ initial_pipeline_version=initial_pipeline_version,
+ )
async def update(
self,
diff --git a/backend/airweave/domains/syncs/service.py b/backend/airweave/domains/syncs/service.py
index f8315b2a3b..8ed9a5bab9 100644
--- a/backend/airweave/domains/syncs/service.py
+++ b/backend/airweave/domains/syncs/service.py
@@ -16,6 +16,7 @@
from uuid import UUID
from fastapi import HTTPException
+from pydantic import JsonValue
from sqlalchemy.ext.asyncio import AsyncSession
from airweave import schemas
@@ -29,6 +30,10 @@
)
from airweave.db.session import get_db_context
from airweave.db.unit_of_work import UnitOfWork
+from airweave.domains.entities.canonical.page_source import CanonicalPageSource
+from airweave.domains.entities.canonical.search_metadata import SEARCH_METADATA_PIPELINE_VERSION
+from airweave.domains.entities.canonical.source import CanonicalSource
+from airweave.domains.owned_provisioning.guard import require_provider_sync
from airweave.domains.sources.exceptions.classifier import classify_error
from airweave.domains.sources.types import SourceRegistryEntry
from airweave.domains.sync_pipeline.config import SyncConfig
@@ -54,6 +59,7 @@
TemporalScheduleServiceProtocol,
TemporalWorkflowServiceProtocol,
)
+from airweave.platform.configs.config import OutlookCalendarConfig, OutlookMailConfig, StripeConfig
from airweave.schemas.source_connection import ScheduleConfig
from airweave.schemas.sync import SyncCreate
from airweave.schemas.sync_job import SyncJobCreate
@@ -106,6 +112,8 @@ async def create(
run_immediately: bool,
ctx: ApiContext,
uow: UnitOfWork,
+ defer_execution: bool = False,
+ source_config: dict[str, JsonValue] | None = None,
) -> SyncProvisionResult:
"""Create sync + optional job + Temporal schedule atomically.
@@ -124,6 +132,22 @@ async def create(
if cron:
self._validate_cron_for_source(cron, source_entry)
+ canonical = isinstance(
+ source_entry.source_class_ref, (CanonicalSource, CanonicalPageSource)
+ )
+ if source_entry.short_name in {"outlook_mail", "outlook_calendar"}:
+ # Outlook composes its canonical adapter only for explicit original capture.
+ config_type = (
+ OutlookMailConfig
+ if source_entry.short_name == "outlook_mail"
+ else OutlookCalendarConfig
+ )
+ canonical = config_type.model_validate(source_config or {}).capture_originals
+ elif source_entry.short_name == "stripe":
+ canonical = (
+ StripeConfig.model_validate(source_config or {}).original_capture is not None
+ )
+
sync_schema, sync_job_schema = await self._create_sync_records(
uow.session,
name=f"Sync for {name}",
@@ -131,11 +155,12 @@ async def create(
destination_connection_ids=destination_connection_ids,
cron_schedule=cron,
run_immediately=run_immediately,
+ initial_pipeline_version=(SEARCH_METADATA_PIPELINE_VERSION if canonical else 1),
ctx=ctx,
uow=uow,
)
- if cron:
+ if cron and not defer_execution:
await self._temporal_schedule_service.create_or_update_schedule(
sync_id=sync_schema.id,
cron_schedule=cron,
@@ -306,6 +331,8 @@ async def cancel_job(
detail=f"Cannot cancel job in {sync_job.status} state",
)
+ await require_provider_sync(db, sync_job.sync_id, ctx.organization.id)
+
if sync_job.status == SyncJobStatus.PENDING:
# PENDING → CANCELLED directly (PENDING → CANCELLING is invalid)
await self._job_state_machine.transition(
@@ -470,6 +497,7 @@ async def _create_sync_records(
destination_connection_ids: List[UUID],
cron_schedule: Optional[str],
run_immediately: bool,
+ initial_pipeline_version: int,
ctx: ApiContext,
uow: UnitOfWork,
) -> Tuple[schemas.Sync, Optional[schemas.SyncJob]]:
@@ -486,7 +514,13 @@ async def _create_sync_records(
run_immediately=run_immediately,
)
- sync_schema = await self._sync_repo.create(uow.session, obj_in=sync_in, ctx=ctx, uow=uow)
+ sync_schema = await self._sync_repo.create(
+ uow.session,
+ obj_in=sync_in,
+ ctx=ctx,
+ uow=uow,
+ initial_pipeline_version=initial_pipeline_version,
+ )
await uow.session.flush()
sync_job_schema: Optional[schemas.SyncJob] = None
diff --git a/backend/airweave/domains/syncs/tests/test_service.py b/backend/airweave/domains/syncs/tests/test_service.py
index ae0e4d4efa..7b7bb849db 100644
--- a/backend/airweave/domains/syncs/tests/test_service.py
+++ b/backend/airweave/domains/syncs/tests/test_service.py
@@ -825,7 +825,9 @@ async def test_cancel_job_not_found():
svc = _build_svc(sync_job_repo=job_repo)
with pytest.raises(HTTPException) as exc_info:
- await svc.cancel_job(AsyncMock(), job_id=uuid4(), ctx=_mock_ctx())
+ await svc.cancel_job(
+ AsyncMock(scalar=AsyncMock(return_value=None)), job_id=uuid4(), ctx=_mock_ctx()
+ )
assert exc_info.value.status_code == 404
@@ -838,7 +840,9 @@ async def test_cancel_job_wrong_status():
svc = _build_svc(sync_job_repo=job_repo)
with pytest.raises(HTTPException) as exc_info:
- await svc.cancel_job(AsyncMock(), job_id=uuid4(), ctx=_mock_ctx())
+ await svc.cancel_job(
+ AsyncMock(scalar=AsyncMock(return_value=None)), job_id=uuid4(), ctx=_mock_ctx()
+ )
assert exc_info.value.status_code == 400
@@ -856,7 +860,7 @@ async def test_cancel_job_success():
}
job_sm = AsyncMock()
- db = AsyncMock()
+ db = AsyncMock(scalar=AsyncMock(return_value=None))
svc = _build_svc(
sync_job_repo=job_repo,
@@ -880,7 +884,7 @@ async def test_cancel_pending_job_transitions_directly_to_cancelled():
temporal = AsyncMock()
job_sm = AsyncMock()
- db = AsyncMock()
+ db = AsyncMock(scalar=AsyncMock(return_value=None))
svc = _build_svc(
sync_job_repo=job_repo,
@@ -909,7 +913,7 @@ async def test_cancel_job_workflow_not_found_marks_cancelled():
}
job_sm = AsyncMock()
- db = AsyncMock()
+ db = AsyncMock(scalar=AsyncMock(return_value=None))
svc = _build_svc(
sync_job_repo=job_repo,
@@ -942,7 +946,9 @@ async def test_cancel_job_temporal_failure():
svc = _build_svc(sync_job_repo=job_repo, temporal_workflow_service=temporal)
with pytest.raises(HTTPException) as exc_info:
- await svc.cancel_job(AsyncMock(), job_id=job_id, ctx=_mock_ctx())
+ await svc.cancel_job(
+ AsyncMock(scalar=AsyncMock(return_value=None)), job_id=job_id, ctx=_mock_ctx()
+ )
assert exc_info.value.status_code == 502
@@ -953,6 +959,7 @@ async def test_cancel_job_temporal_failure():
def _mock_source_entry(*, short_name="github", continuous=False, federated=False):
entry = MagicMock()
+ entry.source_class_ref = object
entry.short_name = short_name
entry.supports_continuous = continuous
entry.federated_search = federated
@@ -1065,7 +1072,31 @@ async def test_create_no_cron_no_run_immediately():
@pytest.mark.asyncio
-async def test_create_with_cron_calls_temporal_schedule():
+@pytest.mark.parametrize(
+ "source_kind, source_config, expected_version",
+ [
+ ("legacy", None, 1),
+ ("gmail", None, 2),
+ ("slack", None, 2),
+ ("outlook_mail", None, 1),
+ ("stripe", None, 1),
+ (
+ "stripe",
+ {
+ "original_capture": {
+ "expected_account_id": "acct_selected",
+ "livemode": False,
+ "api_version": "2025-06-30.basil",
+ }
+ },
+ 2,
+ ),
+ ("outlook_mail", {"capture_originals": True, "expected_principal_id": "native-owner"}, 2),
+ ],
+)
+async def test_create_with_cron_calls_temporal_schedule(
+ source_kind, source_config, expected_version
+):
from airweave.schemas.source_connection import ScheduleConfig
sync_repo = AsyncMock()
@@ -1086,6 +1117,19 @@ async def test_create_with_cron_calls_temporal_schedule():
temporal_schedule_service=temporal_sched,
)
+ from airweave.platform.sources.gmail import GmailSource
+ from airweave.platform.sources.outlook_mail import OutlookMailSource
+ from airweave.platform.sources.slack import SlackSource
+ from airweave.platform.sources.stripe import StripeSource
+
+ source_entry = _mock_source_entry(short_name=source_kind)
+ source_entry.source_class_ref = {
+ "legacy": object,
+ "gmail": GmailSource,
+ "slack": SlackSource,
+ "outlook_mail": OutlookMailSource,
+ "stripe": StripeSource,
+ }[source_kind]
result = await svc.create(
AsyncMock(),
name="test",
@@ -1093,7 +1137,8 @@ async def test_create_with_cron_calls_temporal_schedule():
destination_connection_ids=[uuid4()],
collection_id=uuid4(),
collection_readable_id="col-x",
- source_entry=_mock_source_entry(),
+ source_entry=source_entry,
+ source_config=source_config,
schedule_config=ScheduleConfig(cron="0 6 * * *"),
run_immediately=False,
ctx=_mock_ctx(),
@@ -1102,6 +1147,8 @@ async def test_create_with_cron_calls_temporal_schedule():
assert result is not None
assert result.sync_id == mock_sync.id
temporal_sched.create_or_update_schedule.assert_awaited_once()
+ assert sync_repo.create.call_args.kwargs["initial_pipeline_version"] == expected_version
+ assert "index_pipeline_version" not in sync_repo.create.call_args.kwargs["obj_in"].model_dump()
# ---------------------------------------------------------------------------
@@ -1202,3 +1249,19 @@ async def test_schedule_cleanup_handles_error():
ctx = _mock_ctx()
await svc._schedule_cleanup(uuid4(), uuid4(), uuid4(), ctx)
ctx.logger.error.assert_called_once()
+
+
+async def test_provider_cancellation_rejects_native_job_before_external_actions():
+ job_repo, temporal, state = AsyncMock(), AsyncMock(), AsyncMock()
+ job = _orm_sync_job(job_id=uuid4(), status=SyncJobStatus.RUNNING)
+ job_repo.get.return_value = job
+ svc = _build_svc(
+ sync_job_repo=job_repo, temporal_workflow_service=temporal, job_state_machine=state
+ )
+ with pytest.raises(HTTPException) as error:
+ await svc.cancel_job(
+ AsyncMock(scalar=AsyncMock(return_value=uuid4())), job_id=job.id, ctx=_mock_ctx()
+ )
+ assert error.value.status_code == 409
+ state.transition.assert_not_awaited()
+ temporal.cancel_sync_job_workflow.assert_not_awaited()
diff --git a/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py b/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py
new file mode 100644
index 0000000000..e810c8b789
--- /dev/null
+++ b/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py
@@ -0,0 +1,52 @@
+"""Bounded reclamation on the existing system maintenance workflow."""
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+from temporalio import activity
+
+from airweave.db.session import get_db_context
+from airweave.domains.entities.canonical.projection_gc import ProjectionGCStore
+from airweave.domains.storage.protocols import StorageBackend
+from airweave.platform.destinations.vespa.client import VespaClient
+
+
+@dataclass
+class CleanupProjectionGenerationsActivity:
+ """A tick deletes at most 100 exact document IDs; failures remain durable."""
+
+ storage: StorageBackend
+
+ @activity.defn(name="cleanup_projection_generations_activity")
+ async def run(self) -> None:
+ """Use existing SQL/Temporal infrastructure, without creating schedules here."""
+ store = ProjectionGCStore()
+ now = datetime.now(timezone.utc)
+ async with get_db_context() as db:
+ due = await store.due(db, now=now, limit=20)
+ if not due:
+ return
+ client = await VespaClient.connect()
+ remaining = 100
+ try:
+ for generation in due:
+ if remaining <= 0:
+ break
+ async with get_db_context() as db:
+ page = await store.claim(db, generation, now=now, limit=remaining)
+ if page is None:
+ continue
+ error = None
+ try:
+ await client.delete_documents(
+ [(doc.schema_name, doc.document_id) for doc in page.documents]
+ )
+ for key in page.artifact_keys:
+ await self.storage.delete_file(key)
+ except Exception as failure:
+ error = type(failure).__name__
+ async with get_db_context() as db:
+ await store.acknowledge(db, page, now=datetime.now(timezone.utc), error=error)
+ remaining -= len(page.documents) + len(page.artifact_keys)
+ finally:
+ await client.close()
diff --git a/backend/airweave/domains/temporal/activities/discover_native_projection.py b/backend/airweave/domains/temporal/activities/discover_native_projection.py
new file mode 100644
index 0000000000..33ed3c7007
--- /dev/null
+++ b/backend/airweave/domains/temporal/activities/discover_native_projection.py
@@ -0,0 +1,36 @@
+"""Discover committed canonical work without a second dispatch journal."""
+
+from dataclasses import dataclass
+from uuid import UUID
+
+from temporalio import activity
+
+from airweave.db.session import get_db_context
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.source import indexed_record_types
+from airweave.domains.sources.protocols import SourceRegistryProtocol
+
+
+@dataclass
+class DiscoverNativeProjectionActivity:
+ """Read the existing pending state; no credentials or content cross Temporal."""
+
+ source_registry: SourceRegistryProtocol
+
+ @activity.defn(name="discover_native_projection_activity")
+ async def run(self, after_id: str | None = None) -> dict:
+ """Discover native imports and provider records even while capture is active."""
+ source_names = ("almanac",) + tuple(
+ entry.short_name
+ for entry in self.source_registry.list_all()
+ if entry.short_name != "almanac"
+ and indexed_record_types(entry.short_name, self.source_registry)
+ )
+ async with get_db_context() as db:
+ page = await CanonicalProjectionStore().pending_sources(
+ db,
+ source_names=source_names,
+ after_id=UUID(after_id) if after_id else None,
+ limit=20,
+ )
+ return page.model_dump(mode="json")
diff --git a/backend/airweave/domains/temporal/activities/project_canonical_records.py b/backend/airweave/domains/temporal/activities/project_canonical_records.py
new file mode 100644
index 0000000000..5dab931d0d
--- /dev/null
+++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py
@@ -0,0 +1,73 @@
+"""Project owned records without reconnecting or consulting the provider."""
+
+from dataclasses import dataclass
+from uuid import UUID
+
+from temporalio import activity
+
+from airweave.core.logging import logger
+from airweave.db.session import get_db_context
+from airweave.domains.entities.canonical.projection_models import ProjectionBatchResult
+from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+from airweave.domains.entities.canonical.projector import CanonicalProjector
+from airweave.domains.entities.canonical.source import indexed_record_types
+from airweave.domains.sources.protocols import SourceRegistryProtocol
+from airweave.platform.destinations.vespa.destination import VespaDestination
+
+
+@dataclass
+class ProjectCanonicalRecordsActivity:
+ """One bounded page; Temporal retries infrastructure failures independently."""
+
+ projector: CanonicalProjector
+ source_registry: SourceRegistryProtocol
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def run(
+ self,
+ organization_id: str,
+ sync_id: str,
+ after_id: str | None = None,
+ skip_failed: bool = False,
+ ) -> dict:
+ """Resolve tenant collection scope fresh; never receive OAuth credentials."""
+ organization = UUID(organization_id)
+ sync = UUID(sync_id)
+ async with get_db_context() as db:
+ binding = await CanonicalProjectionStore().binding(db, organization, sync)
+ if binding is None:
+ return ProjectionBatchResult().model_dump(mode="json")
+ source_name, collection_id = binding.source_name, binding.collection_id
+ if not indexed_record_types(source_name, self.source_registry):
+ return ProjectionBatchResult().model_dump(mode="json")
+ # Do not construct an index client for sources with no pending records.
+ pending = await CanonicalProjectionStore().pending(
+ db,
+ organization,
+ sync,
+ after_id=UUID(after_id) if after_id else None,
+ limit=1,
+ skip_failed=skip_failed,
+ )
+ if not pending:
+ return ProjectionBatchResult().model_dump(mode="json")
+ destination = await VespaDestination.create(
+ collection_id=collection_id,
+ organization_id=organization,
+ logger=logger,
+ soft_fail=False,
+ )
+ try:
+ options = {"skip_failed": True} if skip_failed else {}
+ result = await self.projector.batch(
+ organization,
+ sync,
+ source_name,
+ destination,
+ logger,
+ after_id=UUID(after_id) if after_id else None,
+ **options,
+ )
+ return result.model_dump(mode="json")
+ finally:
+ await destination.close_connection()
diff --git a/backend/airweave/domains/temporal/activities/tests/test_discover_native_projection.py b/backend/airweave/domains/temporal/activities/tests/test_discover_native_projection.py
new file mode 100644
index 0000000000..d26ed38fe2
--- /dev/null
+++ b/backend/airweave/domains/temporal/activities/tests/test_discover_native_projection.py
@@ -0,0 +1,43 @@
+"""Discovery uses registered indexing capabilities, never provider authentication calls."""
+
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+from uuid import uuid4
+
+from airweave.domains.entities.canonical.projection_models import ProjectionSourcePage
+from airweave.domains.temporal.activities import discover_native_projection as module
+
+
+async def test_discovers_native_and_canonical_capabilities_with_bounded_cursor(monkeypatch):
+ entries = {
+ name: SimpleNamespace(
+ short_name=name,
+ source_class_ref=SimpleNamespace(canonical_record_types=kinds),
+ )
+ for name, kinds in (
+ ("slack", ("channel", "message", "file")),
+ ("gmail", ("message",)),
+ ("legacy", ()),
+ )
+ }
+ registry = Mock()
+ registry.list_all.return_value = list(entries.values())
+ registry.get.side_effect = entries.__getitem__
+ db = object()
+
+ @asynccontextmanager
+ async def sessions():
+ yield db
+
+ pending = AsyncMock(return_value=ProjectionSourcePage(sources=(), next_cursor=None))
+ monkeypatch.setattr(module, "get_db_context", sessions)
+ monkeypatch.setattr(
+ module, "CanonicalProjectionStore", lambda: SimpleNamespace(pending_sources=pending)
+ )
+ after = uuid4()
+ result = await module.DiscoverNativeProjectionActivity(registry).run(str(after))
+ assert result == {"sources": [], "next_cursor": None}
+ pending.assert_awaited_once_with(
+ db, source_names=("almanac", "slack", "gmail"), after_id=after, limit=20
+ )
diff --git a/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py b/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py
new file mode 100644
index 0000000000..83d507354d
--- /dev/null
+++ b/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py
@@ -0,0 +1,147 @@
+"""Real source/collection scope resolution; only the remote projector is replaced."""
+
+# ruff: noqa: F811
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+from uuid import UUID, uuid4
+
+import pytest
+
+from airweave.domains.entities.canonical.projection_models import ProjectionBatchResult
+from airweave.domains.entities.canonical.requests import RecordIdentity
+from airweave.domains.entities.canonical.tests.conftest import database, source # noqa: F401
+from airweave.domains.entities.canonical.tests.helpers import capture, observation
+from airweave.domains.temporal.activities import project_canonical_records as module
+from airweave.models import Collection, SourceConnection
+from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata
+
+
+@pytest.fixture
+async def activity_setup(database, source, monkeypatch):
+ service, fence = source
+ metadata = VectorDbDeploymentMetadata(
+ dense_embedder="test", sparse_embedder="test", embedding_dimensions=3
+ )
+ collection_id, connection_id = uuid4(), uuid4()
+ async with database() as db:
+ db.add(metadata)
+ await db.flush()
+ db.add(
+ Collection(
+ id=collection_id,
+ name="Native",
+ readable_id="native",
+ organization_id=fence.organization_id,
+ vector_db_deployment_metadata_id=metadata.id,
+ )
+ )
+ db.add(
+ SourceConnection(
+ id=connection_id,
+ organization_id=fence.organization_id,
+ sync_id=fence.sync_id,
+ name="Native",
+ is_authenticated=True,
+ short_name="almanac",
+ readable_collection_id="native",
+ )
+ )
+ await db.commit()
+ await capture(
+ database,
+ service,
+ fence,
+ observation(
+ identity=RecordIdentity(record_type="knowledge", native_id="one"),
+ payload={"body": "नमस्ते"},
+ ),
+ )
+ monkeypatch.setattr(module, "get_db_context", database)
+ destination = SimpleNamespace(close_connection=AsyncMock())
+ create = AsyncMock(return_value=destination)
+ monkeypatch.setattr(module.VespaDestination, "create", create)
+ projector = SimpleNamespace(batch=AsyncMock(return_value=ProjectionBatchResult(published=1)))
+ registry = Mock()
+ registry.get.side_effect = AssertionError("Native sources must not enter provider registry")
+ activity = module.ProjectCanonicalRecordsActivity(projector, registry)
+ return activity, fence, connection_id, collection_id, destination, create
+
+
+@pytest.mark.parametrize("provider", ["almanac", "gmail", "legacy"])
+async def test_capability_routes_native_and_provider(activity_setup, database, provider):
+ activity, fence, connection_id, collection_id, destination, create = activity_setup
+ if provider != "almanac":
+ async with database() as db:
+ connection = await db.get(SourceConnection, connection_id)
+ connection.short_name = provider
+ await db.commit()
+ activity.source_registry.get.side_effect = None
+ activity.source_registry.get.return_value = SimpleNamespace(
+ source_class_ref=SimpleNamespace(
+ canonical_record_types=("message",) if provider == "gmail" else ()
+ )
+ )
+ result = await activity.run(str(fence.organization_id), str(fence.sync_id), str(UUID(int=0)))
+ if provider == "legacy":
+ assert result == ProjectionBatchResult().model_dump(mode="json")
+ create.assert_not_awaited()
+ activity.projector.batch.assert_not_awaited()
+ return
+ assert result["published"] == 1
+ create.assert_awaited_once_with(
+ collection_id=collection_id,
+ organization_id=fence.organization_id,
+ logger=module.logger,
+ soft_fail=False,
+ )
+ activity.projector.batch.assert_awaited_once_with(
+ fence.organization_id,
+ fence.sync_id,
+ provider,
+ destination,
+ module.logger,
+ after_id=UUID(int=0),
+ )
+ destination.close_connection.assert_awaited_once()
+ if provider == "almanac":
+ activity.source_registry.get.assert_not_called()
+ else:
+ activity.source_registry.get.assert_called_once_with(provider)
+
+
+async def test_foreign_organization_and_empty_page_do_not_connect(activity_setup):
+ activity, fence, _, _, destination, create = activity_setup
+ assert await activity.run(
+ str(uuid4()), str(fence.sync_id)
+ ) == ProjectionBatchResult().model_dump(mode="json")
+ # This valid source has no pending records beyond the maximum UUID.
+ assert await activity.run(
+ str(fence.organization_id), str(fence.sync_id), str(UUID(int=2**128 - 1))
+ ) == ProjectionBatchResult().model_dump(mode="json")
+ activity.source_registry.get.assert_not_called()
+ activity.projector.batch.assert_not_awaited()
+ create.assert_not_awaited()
+ destination.close_connection.assert_not_awaited()
+
+
+async def test_projector_failure_propagates_and_closes_destination(activity_setup):
+ activity, fence, _, _, destination, _ = activity_setup
+ activity.projector.batch.side_effect = RuntimeError("synthetic projection failure")
+ with pytest.raises(RuntimeError, match="synthetic projection failure"):
+ await activity.run(str(fence.organization_id), str(fence.sync_id))
+ destination.close_connection.assert_awaited_once()
+
+
+async def test_deauthenticated_binding_never_constructs_destination(activity_setup, database):
+ activity, fence, connection_id, _, destination, create = activity_setup
+ async with database() as db:
+ connection = await db.get(SourceConnection, connection_id)
+ connection.is_authenticated = False
+ await db.commit()
+ assert await activity.run(str(fence.organization_id), str(fence.sync_id)) == (
+ ProjectionBatchResult().model_dump(mode="json")
+ )
+ activity.source_registry.get.assert_not_called()
+ activity.projector.batch.assert_not_awaited()
+ create.assert_not_awaited()
+ destination.close_connection.assert_not_awaited()
diff --git a/backend/airweave/domains/temporal/protocols.py b/backend/airweave/domains/temporal/protocols.py
index 519423331d..b936567b8b 100644
--- a/backend/airweave/domains/temporal/protocols.py
+++ b/backend/airweave/domains/temporal/protocols.py
@@ -82,6 +82,9 @@ async def delete_all_schedules_for_sync(
sync_id: UUID,
db: AsyncSession,
ctx: ApiContext,
+ *,
+ uow: Optional[UnitOfWork] = None,
+ strict: bool = False,
) -> None:
"""Delete all schedules associated with a sync."""
...
diff --git a/backend/airweave/domains/temporal/schedule_service.py b/backend/airweave/domains/temporal/schedule_service.py
index 5d5fedf7e6..2615b80f84 100644
--- a/backend/airweave/domains/temporal/schedule_service.py
+++ b/backend/airweave/domains/temporal/schedule_service.py
@@ -200,7 +200,6 @@ async def _create_schedule(
{
"temporal_schedule_id": schedule_id,
"sync_type": sync_type,
- "status": "ACTIVE",
"cron_schedule": cron_expression,
},
ctx,
@@ -265,11 +264,16 @@ async def _delete_schedule_by_id(
sync_id: UUID,
db: AsyncSession,
ctx: ApiContext,
+ uow: Optional[UnitOfWork] = None,
) -> None:
"""Delete a single Temporal schedule and clear sync DB fields."""
client = await self._get_client()
handle = client.get_schedule_handle(schedule_id)
- await handle.delete()
+ try:
+ await handle.delete()
+ except RPCError as exc:
+ if exc.status != RPCStatusCode.NOT_FOUND:
+ raise
sync_obj = await self._sync_repo.get_without_connections(db, sync_id, ctx)
await self._sync_repo.update(
@@ -281,6 +285,7 @@ async def _delete_schedule_by_id(
"sync_type": "full",
},
ctx,
+ uow=uow,
)
logger.info(f"Deleted schedule {schedule_id}")
@@ -385,7 +390,7 @@ async def create_or_update_schedule(
if sync.temporal_schedule_id:
status = await self._check_schedule_exists(sync.temporal_schedule_id)
if status["exists"]:
- await self.delete_all_schedules_for_sync(sync_id, db, ctx)
+ await self.delete_all_schedules_for_sync(sync_id, db, ctx, uow=uow, strict=True)
else:
logger.warning(
f"Schedule {sync.temporal_schedule_id} not found in Temporal "
@@ -430,12 +435,17 @@ async def delete_all_schedules_for_sync(
sync_id: UUID,
db: AsyncSession,
ctx: ApiContext,
+ *,
+ uow: Optional[UnitOfWork] = None,
+ strict: bool = False,
) -> None:
"""Delete all schedules (regular + minute + daily cleanup) for a sync."""
for sid in sched_ids.all_schedule_ids(sync_id):
try:
- await self._delete_schedule_by_id(sid, sync_id, db, ctx)
+ await self._delete_schedule_by_id(sid, sync_id, db, ctx, uow=uow)
except Exception as e:
+ if strict:
+ raise
logger.info(f"Schedule {sid} not deleted (may not exist): {e}")
async def delete_schedule_handle(self, schedule_id: str) -> None:
diff --git a/backend/airweave/domains/temporal/service.py b/backend/airweave/domains/temporal/service.py
index 48a38d69c9..a1ac484ae1 100644
--- a/backend/airweave/domains/temporal/service.py
+++ b/backend/airweave/domains/temporal/service.py
@@ -12,6 +12,7 @@
from typing import List, Optional
from temporalio.client import WorkflowHandle
+from temporalio.common import WorkflowIDReusePolicy
from temporalio.service import RPCError, RPCStatusCode
from airweave import schemas
@@ -79,6 +80,7 @@ async def run_source_connection_workflow(
force_full_sync,
],
id=workflow_id,
+ id_reuse_policy=WorkflowIDReusePolicy.REJECT_DUPLICATE,
task_queue=settings.TEMPORAL_TASK_QUEUE,
)
diff --git a/backend/airweave/domains/temporal/tests/test_temporal_schedule_service.py b/backend/airweave/domains/temporal/tests/test_temporal_schedule_service.py
index 4bd79df8fa..cb010be570 100644
--- a/backend/airweave/domains/temporal/tests/test_temporal_schedule_service.py
+++ b/backend/airweave/domains/temporal/tests/test_temporal_schedule_service.py
@@ -12,7 +12,6 @@
from uuid import UUID, uuid4
import pytest
-from fastapi import HTTPException
from temporalio.service import RPCError, RPCStatusCode
from airweave.domains.temporal.exceptions import InvalidCronExpressionError
@@ -628,7 +627,7 @@ async def test_delete_all_schedules_for_sync(case: DeleteAllCase):
svc = _build_svc()
call_log = []
- async def mock_delete(schedule_id, sync_id, db, ctx):
+ async def mock_delete(schedule_id, sync_id, db, ctx, *, uow=None):
call_log.append(schedule_id)
for prefix in case.delete_raises:
if schedule_id.startswith(prefix):
@@ -856,10 +855,7 @@ async def test_pause_schedules_pauses_all_prefixes():
f"minute-sync-{sync_id}",
f"daily-cleanup-{sync_id}",
}
- actual_ids = {
- call.args[0]
- for call in mock_client.get_schedule_handle.call_args_list
- }
+ actual_ids = {call.args[0] for call in mock_client.get_schedule_handle.call_args_list}
assert actual_ids == expected_ids
@@ -870,9 +866,7 @@ async def test_pause_schedules_swallows_not_found():
svc = _build_svc()
handle = AsyncMock()
- handle.pause.side_effect = _rpc_error(
- "not found", RPCStatusCode.NOT_FOUND
- )
+ handle.pause.side_effect = _rpc_error("not found", RPCStatusCode.NOT_FOUND)
mock_client = MagicMock()
mock_client.get_schedule_handle.return_value = handle
svc._get_client = AsyncMock(return_value=mock_client)
@@ -903,9 +897,7 @@ async def test_unpause_schedules_swallows_not_found():
svc = _build_svc()
handle = AsyncMock()
- handle.unpause.side_effect = _rpc_error(
- "not found", RPCStatusCode.NOT_FOUND
- )
+ handle.unpause.side_effect = _rpc_error("not found", RPCStatusCode.NOT_FOUND)
mock_client = MagicMock()
mock_client.get_schedule_handle.return_value = handle
svc._get_client = AsyncMock(return_value=mock_client)
@@ -953,14 +945,18 @@ async def test_get_schedules_for_sync_returns_metadata():
entries = [
_mock_list_entry(
- f"minute-sync-{sync_id}", paused=True,
+ f"minute-sync-{sync_id}",
+ paused=True,
note="Credential error: api_key_invalid",
- next_times=[next_time], recent_count=3,
+ next_times=[next_time],
+ recent_count=3,
),
_mock_list_entry(
- f"daily-cleanup-{sync_id}", paused=True,
+ f"daily-cleanup-{sync_id}",
+ paused=True,
note="Credential error: api_key_invalid",
- next_times=[], recent_count=1,
+ next_times=[],
+ recent_count=1,
),
]
@@ -997,3 +993,74 @@ async def test_get_schedules_for_sync_empty():
result = await svc.get_schedules_for_sync(uuid4())
assert result == []
+
+
+async def test_owned_cleanup_propagates_failure_and_keeps_caller_transaction():
+ """A failed Temporal delete cannot acknowledge disconnect or release its DB lock."""
+ svc = _build_svc()
+ failure = RuntimeError("Temporal unavailable")
+ svc._delete_schedule_by_id = AsyncMock(side_effect=failure)
+ db, uow, ctx = AsyncMock(), AsyncMock(), _mock_ctx()
+ with pytest.raises(RuntimeError, match="Temporal unavailable"):
+ await svc.delete_all_schedules_for_sync(SYNC_ID, db, ctx, uow=uow, strict=True)
+ assert svc._delete_schedule_by_id.call_args.kwargs["uow"] is uow
+ db.commit.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_owned_reconcile_replaces_orphan_schedule_and_relinks_current_generation():
+ """A lost create reply leaves remote state but no committed DB schedule link."""
+ from types import SimpleNamespace
+
+ from airweave.domains.temporal import schedule_ids
+
+ sid = schedule_ids.sync_schedule_id(SYNC_ID)
+ remote = {sid: SimpleNamespace(action=SimpleNamespace(args=[{"provisioning_generation": 1}]))}
+ sync = _mock_sync_model(temporal_schedule_id=None)
+ repo = AsyncMock()
+ repo.get_without_connections.return_value = sync
+
+ async def save(db, obj, values, ctx, **kwargs):
+ for key, value in values.items():
+ setattr(obj, key, value)
+
+ repo.update.side_effect = save
+ client = MagicMock()
+
+ def handle(key):
+ async def delete():
+ if remote.pop(key, None) is None:
+ raise _rpc_error("missing", RPCStatusCode.NOT_FOUND)
+
+ async def describe():
+ if key not in remote:
+ raise _rpc_error("missing", RPCStatusCode.NOT_FOUND)
+ return SimpleNamespace(schedule=remote[key])
+
+ return SimpleNamespace(delete=delete, describe=describe)
+
+ async def create(key, schedule, **kwargs):
+ assert key not in remote
+ remote[key] = schedule
+
+ client.get_schedule_handle.side_effect = handle
+ client.create_schedule = AsyncMock(side_effect=create)
+ svc = _build_svc(sync_repo=repo)
+ db, ctx, uow = AsyncMock(), _mock_ctx(), MagicMock()
+ with (
+ patch.object(svc, "_get_client", new=AsyncMock(return_value=client)),
+ patch.object(
+ svc,
+ "_gather_schedule_data",
+ new=AsyncMock(return_value=({"provisioning_generation": 2}, {}, {})),
+ ),
+ ):
+ # These are the two real schedule-service calls made by owned reconciliation.
+ await svc.delete_all_schedules_for_sync(SYNC_ID, db, ctx, uow=uow, strict=True)
+ await svc.create_or_update_schedule(SYNC_ID, "0 0 * * *", db, ctx, uow)
+ assert list(remote) == [sid]
+ assert remote[sid].action.args[0]["provisioning_generation"] == 2
+ assert sync.temporal_schedule_id == sid
+ assert sync.cron_schedule == "0 0 * * *"
+ assert all(call.kwargs["uow"] is uow for call in repo.update.call_args_list)
+ db.commit.assert_not_awaited()
diff --git a/backend/airweave/domains/temporal/worker/__init__.py b/backend/airweave/domains/temporal/worker/__init__.py
index a622a74416..861b8f0d46 100644
--- a/backend/airweave/domains/temporal/worker/__init__.py
+++ b/backend/airweave/domains/temporal/worker/__init__.py
@@ -10,7 +10,6 @@
import asyncio
import signal
from datetime import timedelta
-from typing import Any
from temporalio.worker import Worker
@@ -63,12 +62,14 @@ def __init__(self, config: WorkerConfig) -> None:
self._runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(
- bind_address=f"0.0.0.0:{config.sdk_metrics_port}",
+ bind_address=f"{config.bind_host}:{config.sdk_metrics_port}",
),
),
)
self._worker: Worker | None = None
self._state = WorkerState()
+ self._stop_task: asyncio.Task[None] | None = None
+ self._startup_lock = asyncio.Lock()
registry = CollectorRegistry()
self._control_server = WorkerControlServer(
@@ -81,65 +82,79 @@ def __init__(self, config: WorkerConfig) -> None:
async def start(self) -> None:
"""Start the worker and control server."""
- # Start control server first (non-blocking)
- try:
- await self._control_server.start()
- except Exception as e:
- logger.warning(f"Failed to start control server (metrics unavailable): {e}")
-
- from airweave.domains.temporal.client import get_client as get_temporal_client
-
- client = await get_temporal_client(runtime=self._runtime)
- logger.info(f"Starting Temporal worker on task queue: {self._config.task_queue}")
-
- # Create worker
- self._worker = Worker(
- client,
- task_queue=self._config.task_queue,
- workflows=get_workflows(),
- activities=create_activities(),
- workflow_runner=self._get_sandbox_runner(),
- max_concurrent_workflow_task_polls=self._config.max_concurrent_workflow_polls,
- max_concurrent_activity_task_polls=self._config.max_concurrent_activity_polls,
- sticky_queue_schedule_to_start_timeout=self._config.sticky_queue_schedule_to_start_timeout,
- nonsticky_to_sticky_poll_ratio=self._config.nonsticky_to_sticky_poll_ratio,
- default_heartbeat_throttle_interval=self._config.default_heartbeat_throttle_interval,
- max_heartbeat_throttle_interval=self._config.max_heartbeat_throttle_interval,
- graceful_shutdown_timeout=timedelta(
- seconds=self._config.graceful_shutdown_timeout_seconds
- ),
- )
+ async with self._startup_lock:
+ if self._stop_task is not None:
+ return
+ # Start control server first (non-blocking)
+ try:
+ await self._control_server.start()
+ except Exception as e:
+ logger.warning(f"Failed to start control server (metrics unavailable): {e}")
+
+ if self._stop_task is not None:
+ return
+
+ from airweave.domains.temporal.client import get_client as get_temporal_client
+
+ client = await get_temporal_client(runtime=self._runtime)
+ if self._stop_task is not None:
+ return
+ logger.info(f"Starting Temporal worker on task queue: {self._config.task_queue}")
+
+ # Create worker
+ self._worker = Worker(
+ client,
+ task_queue=self._config.task_queue,
+ workflows=get_workflows(),
+ activities=create_activities(),
+ workflow_runner=self._get_sandbox_runner(),
+ max_concurrent_workflow_task_polls=self._config.max_concurrent_workflow_polls,
+ max_concurrent_activity_task_polls=self._config.max_concurrent_activity_polls,
+ sticky_queue_schedule_to_start_timeout=self._config.sticky_queue_schedule_to_start_timeout,
+ nonsticky_to_sticky_poll_ratio=self._config.nonsticky_to_sticky_poll_ratio,
+ default_heartbeat_throttle_interval=self._config.default_heartbeat_throttle_interval,
+ max_heartbeat_throttle_interval=self._config.max_heartbeat_throttle_interval,
+ graceful_shutdown_timeout=timedelta(
+ seconds=self._config.graceful_shutdown_timeout_seconds
+ ),
+ )
- # Wire up drain callback
- self._state.on_drain_requested = self._on_drain
+ # Wire up drain callback
+ self._state.on_drain_requested = self._on_drain
- # Start worker (blocks until shutdown)
- self._state.running = True
- logger.info(
- f"Worker started with graceful shutdown timeout: "
- f"{self._config.graceful_shutdown_timeout_seconds}s"
- )
+ # Start worker (blocks until shutdown)
+ self._state.running = True
+ logger.info(
+ f"Worker started with graceful shutdown timeout: "
+ f"{self._config.graceful_shutdown_timeout_seconds}s"
+ )
await self._worker.run()
async def stop(self) -> None:
- """Stop the worker and control server."""
- if self._worker and self._state.running:
- logger.info("Stopping worker gracefully")
- self._state.running = False
- await self._worker.shutdown()
+ """Wait for the single shutdown shared by signals, drain and main cleanup."""
+ if self._stop_task is None:
+ self._stop_task = asyncio.create_task(self._stop())
+ await asyncio.shield(self._stop_task)
+
+ async def _stop(self) -> None:
+ """Shut down polling before releasing the control server and client."""
+ async with self._startup_lock:
+ if self._worker and self._state.running:
+ logger.info("Stopping worker gracefully")
+ self._state.running = False
+ await self._worker.shutdown()
- await self._control_server.stop()
+ await self._control_server.stop()
- from airweave.domains.temporal.client import close as close_temporal_client
+ from airweave.domains.temporal.client import close as close_temporal_client
- await close_temporal_client()
+ await close_temporal_client()
async def _on_drain(self) -> None:
"""Called when drain is requested via control server."""
try:
logger.info("Calling worker.shutdown() - stops polling for new work")
- if self._worker:
- await self._worker.shutdown()
+ await self.stop()
logger.info("Worker shutdown complete - process will exit")
except Exception as e:
logger.error(f"Error during worker shutdown: {e}")
@@ -173,25 +188,26 @@ async def main() -> None:
initialize_container(settings)
logger.info("Container initialized successfully")
- # 2. Require OCR backend for the sync worker
+ # OCR is an optional document capability, not a worker startup dependency.
if container_mod.container.ocr_provider is None:
- logger.error(
- "Temporal worker requires an OCR backend. "
- "Set MISTRAL_API_KEY or DOCLING_BASE_URL and restart."
+ logger.warning(
+ "OCR is not configured. Text extraction and source capture remain available; "
+ "documents requiring OCR stay unindexed. Set MISTRAL_API_KEY or "
+ "DOCLING_BASE_URL to enable OCR and retry their projection."
)
- raise SystemExit(1)
# 3. Create worker with config
config = WorkerConfig.from_settings()
worker = TemporalWorker(config)
# 5. Set up signal handlers
- def signal_handler(signum: int, frame: Any) -> None:
+ def signal_handler(signum: int) -> None:
logger.info(f"Received signal {signum}, shutting down...")
asyncio.create_task(worker.stop())
- signal.signal(signal.SIGINT, signal_handler)
- signal.signal(signal.SIGTERM, signal_handler)
+ loop = asyncio.get_running_loop()
+ for signum in (signal.SIGINT, signal.SIGTERM):
+ loop.add_signal_handler(signum, signal_handler, signum)
# 6. Run worker
try:
@@ -199,7 +215,11 @@ def signal_handler(signum: int, frame: Any) -> None:
except KeyboardInterrupt:
logger.info("Received keyboard interrupt, shutting down...")
finally:
- await worker.stop()
+ try:
+ await worker.stop()
+ finally:
+ for signum in (signal.SIGINT, signal.SIGTERM):
+ loop.remove_signal_handler(signum)
if __name__ == "__main__":
diff --git a/backend/airweave/domains/temporal/worker/config.py b/backend/airweave/domains/temporal/worker/config.py
index 4d0c30cd59..5304a2a022 100644
--- a/backend/airweave/domains/temporal/worker/config.py
+++ b/backend/airweave/domains/temporal/worker/config.py
@@ -12,6 +12,7 @@ class WorkerConfig:
Attributes:
task_queue: Temporal task queue name
+ bind_host: Bind address for control server and SDK metrics
metrics_port: Port for control server (health, metrics, drain)
graceful_shutdown_timeout_seconds: How long to wait for activities to complete
@@ -31,6 +32,8 @@ class WorkerConfig:
metrics_port: int
graceful_shutdown_timeout_seconds: int
+ bind_host: str = "0.0.0.0"
+
# Polling concurrency
max_concurrent_workflow_polls: int = 8
max_concurrent_activity_polls: int = 16
@@ -55,6 +58,7 @@ def from_settings(cls) -> "WorkerConfig":
return cls(
task_queue=settings.TEMPORAL_TASK_QUEUE,
metrics_port=settings.WORKER_METRICS_PORT,
+ bind_host=settings.WORKER_BIND_HOST,
graceful_shutdown_timeout_seconds=settings.TEMPORAL_GRACEFUL_SHUTDOWN_TIMEOUT,
disable_sandbox=settings.TEMPORAL_DISABLE_SANDBOX,
sdk_metrics_port=settings.TEMPORAL_SDK_METRICS_PORT,
diff --git a/backend/airweave/domains/temporal/worker/control_server.py b/backend/airweave/domains/temporal/worker/control_server.py
index 7b7349f406..a875457984 100644
--- a/backend/airweave/domains/temporal/worker/control_server.py
+++ b/backend/airweave/domains/temporal/worker/control_server.py
@@ -78,11 +78,11 @@ async def start(self) -> None:
self._runner = web.AppRunner(app)
await self._runner.setup()
- site = web.TCPSite(self._runner, "0.0.0.0", self._config.metrics_port)
+ site = web.TCPSite(self._runner, self._config.bind_host, self._config.metrics_port)
await site.start()
logger.info(
- f"Control server started on 0.0.0.0:{self._config.metrics_port} "
+ f"Control server started on {self._config.bind_host}:{self._config.metrics_port} "
f"(endpoints: /health, /metrics, /status, /drain)"
)
diff --git a/backend/airweave/domains/temporal/worker/tests/test_config.py b/backend/airweave/domains/temporal/worker/tests/test_config.py
index ae13090cc3..335077f392 100644
--- a/backend/airweave/domains/temporal/worker/tests/test_config.py
+++ b/backend/airweave/domains/temporal/worker/tests/test_config.py
@@ -14,6 +14,7 @@ def test_sdk_metrics_port_default():
graceful_shutdown_timeout_seconds=30,
)
assert config.sdk_metrics_port == 9090
+ assert config.bind_host == "0.0.0.0"
def test_sdk_metrics_port_override():
@@ -32,6 +33,7 @@ def test_from_settings_wires_sdk_metrics_port():
with patch("airweave.domains.temporal.worker.config.settings") as mock:
mock.TEMPORAL_TASK_QUEUE = "q"
mock.WORKER_METRICS_PORT = 8080
+ mock.WORKER_BIND_HOST = "0.0.0.0"
mock.TEMPORAL_GRACEFUL_SHUTDOWN_TIMEOUT = 30
mock.TEMPORAL_DISABLE_SANDBOX = False
mock.TEMPORAL_SDK_METRICS_PORT = 7777
@@ -46,6 +48,7 @@ def test_from_settings_wires_all_fields():
with patch("airweave.domains.temporal.worker.config.settings") as mock:
mock.TEMPORAL_TASK_QUEUE = "my-queue"
mock.WORKER_METRICS_PORT = 9091
+ mock.WORKER_BIND_HOST = "127.0.0.1"
mock.TEMPORAL_GRACEFUL_SHUTDOWN_TIMEOUT = 60
mock.TEMPORAL_DISABLE_SANDBOX = True
mock.TEMPORAL_SDK_METRICS_PORT = 9999
@@ -54,6 +57,7 @@ def test_from_settings_wires_all_fields():
assert config.task_queue == "my-queue"
assert config.metrics_port == 9091
+ assert config.bind_host == "127.0.0.1"
assert config.graceful_shutdown_timeout_seconds == 60
assert config.disable_sandbox is True
assert config.sdk_metrics_port == 9999
diff --git a/backend/airweave/domains/temporal/worker/tests/test_wiring.py b/backend/airweave/domains/temporal/worker/tests/test_wiring.py
index 13a59fc1d4..5948f2e0a9 100644
--- a/backend/airweave/domains/temporal/worker/tests/test_wiring.py
+++ b/backend/airweave/domains/temporal/worker/tests/test_wiring.py
@@ -32,7 +32,17 @@ def test_create_activities_returns_list():
result = create_activities()
assert isinstance(result, list)
- assert len(result) == 7
+ assert len(result) == 10
+ activity_classes = {type(method.__self__).__name__ for method in result}
+ assert "ProjectCanonicalRecordsActivity" in activity_classes
+ assert "CleanupProjectionGenerationsActivity" in activity_classes
+ assert "DiscoverNativeProjectionActivity" in activity_classes
+ discovery = next(
+ method.__self__
+ for method in result
+ if type(method.__self__).__name__ == "DiscoverNativeProjectionActivity"
+ )
+ assert discovery.source_registry is mock_container.source_registry
@pytest.mark.unit
@@ -51,8 +61,10 @@ def test_get_workflows_returns_classes():
result = get_workflows()
assert isinstance(result, list)
- assert len(result) == 4
+ assert len(result) == 6
class_names = [cls.__name__ for cls in result]
+ assert "ProjectCanonicalRecordsWorkflow" in class_names
+ assert "RecoverNativeProjectionWorkflow" in class_names
assert "RunSourceConnectionWorkflow" in class_names
assert "CleanupStuckSyncJobsWorkflow" in class_names
assert "CleanupSyncDataWorkflow" in class_names
diff --git a/backend/airweave/domains/temporal/worker/tests/test_worker.py b/backend/airweave/domains/temporal/worker/tests/test_worker.py
index 339edf7b87..c690a3ec9e 100644
--- a/backend/airweave/domains/temporal/worker/tests/test_worker.py
+++ b/backend/airweave/domains/temporal/worker/tests/test_worker.py
@@ -5,9 +5,11 @@
WorkerMetricsRegistry) is constructed for real so tests exercise actual wiring.
"""
+import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
+from aiohttp import ClientSession
from temporalio.runtime import PrometheusConfig, TelemetryConfig
from airweave.domains.temporal.worker.config import WorkerConfig
@@ -16,11 +18,11 @@
def _make_config(**overrides):
"""Build a minimal WorkerConfig for testing."""
- defaults = dict(
- task_queue="test-queue",
- metrics_port=8080,
- graceful_shutdown_timeout_seconds=30,
- )
+ defaults = {
+ "task_queue": "test-queue",
+ "metrics_port": 8080,
+ "graceful_shutdown_timeout_seconds": 30,
+ }
defaults.update(overrides)
return WorkerConfig(**defaults)
@@ -33,7 +35,7 @@ def test_init_creates_runtime_with_correct_config(mock_runtime_cls):
"""__init__ builds a real TelemetryConfig(PrometheusConfig) and passes it to Runtime."""
from airweave.domains.temporal.worker import TemporalWorker
- config = _make_config(sdk_metrics_port=9999)
+ config = _make_config(sdk_metrics_port=9999, bind_host="127.0.0.1")
worker = TemporalWorker(config)
# Runtime was called exactly once
@@ -46,7 +48,7 @@ def test_init_creates_runtime_with_correct_config(mock_runtime_cls):
# Its metrics member is a real PrometheusConfig with the right address
assert isinstance(telemetry.metrics, PrometheusConfig)
- assert telemetry.metrics.bind_address == "0.0.0.0:9999"
+ assert telemetry.metrics.bind_address == "127.0.0.1:9999"
# The worker holds the Runtime *instance*
assert worker._runtime is mock_runtime_cls.return_value
@@ -183,3 +185,160 @@ async def test_stop_skips_shutdown_when_not_running(mock_runtime_cls, mock_clien
worker._control_server.stop.assert_awaited_once()
mock_client_close.assert_awaited_once()
+
+
+@patch("temporalio.runtime.Runtime")
+async def test_control_server_binds_only_loopback(mock_runtime_cls):
+ """Real aiohttp listener uses the same configured host as SDK metrics."""
+ from airweave.domains.temporal.worker import TemporalWorker
+
+ worker = TemporalWorker(_make_config(bind_host="127.0.0.1", metrics_port=0))
+ server = worker._control_server
+ try:
+ await server.start()
+ assert server._runner is not None
+ addresses = server._runner.addresses
+ assert len(addresses) == 1
+ host, port = addresses[0]
+ assert host == "127.0.0.1"
+ async with ClientSession() as client:
+ async with client.get(f"http://127.0.0.1:{port}/health") as response:
+ assert response.status == 503
+ assert await response.text() == "NOT_RUNNING"
+ finally:
+ await server.stop()
+
+
+@patch("airweave.domains.temporal.client.close", new_callable=AsyncMock)
+@patch("temporalio.runtime.Runtime")
+async def test_concurrent_stop_waits_for_one_complete_shutdown(mock_runtime_cls, mock_client_close):
+ """Signal and main-finally callers share SDK shutdown and subsequent cleanup."""
+ from airweave.domains.temporal.worker import TemporalWorker
+
+ worker = TemporalWorker(_make_config())
+ entered = asyncio.Event()
+ release = asyncio.Event()
+
+ async def shutdown():
+ entered.set()
+ await release.wait()
+
+ worker._worker = AsyncMock()
+ worker._worker.shutdown.side_effect = shutdown
+ worker._state.running = True
+ worker._control_server.stop = AsyncMock()
+ first = asyncio.create_task(worker.stop())
+ await entered.wait()
+ second = asyncio.create_task(worker.stop())
+ try:
+ await asyncio.sleep(0)
+ assert not second.done()
+ worker._control_server.stop.assert_not_awaited()
+ mock_client_close.assert_not_awaited()
+ finally:
+ release.set()
+ await asyncio.gather(first, second)
+ await worker.stop()
+ worker._worker.shutdown.assert_awaited_once()
+ worker._control_server.stop.assert_awaited_once()
+ mock_client_close.assert_awaited_once()
+
+
+@patch("airweave.domains.temporal.client.close", new_callable=AsyncMock)
+@patch("temporalio.runtime.Runtime")
+async def test_cancelled_stop_waiter_does_not_cancel_shutdown(mock_runtime_cls, mock_client_close):
+ """Cancellation of one caller leaves shutdown owned by the shared task."""
+ from airweave.domains.temporal.worker import TemporalWorker
+
+ worker = TemporalWorker(_make_config())
+ entered = asyncio.Event()
+ release = asyncio.Event()
+
+ async def shutdown():
+ entered.set()
+ await release.wait()
+
+ worker._worker = AsyncMock()
+ worker._worker.shutdown.side_effect = shutdown
+ worker._state.running = True
+ worker._control_server.stop = AsyncMock()
+ caller = asyncio.create_task(worker.stop())
+ await entered.wait()
+ caller.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await caller
+ release.set()
+ await worker.stop()
+ worker._worker.shutdown.assert_awaited_once()
+ worker._control_server.stop.assert_awaited_once()
+ mock_client_close.assert_awaited_once()
+
+
+@patch("temporalio.runtime.Runtime")
+async def test_stop_failure_remains_observable_to_later_callers(mock_runtime_cls):
+ """A failed shutdown cannot become an apparent successful repeated stop."""
+ from airweave.domains.temporal.worker import TemporalWorker
+
+ worker = TemporalWorker(_make_config())
+ worker._worker = AsyncMock()
+ worker._worker.shutdown.side_effect = RuntimeError("shutdown failed")
+ worker._state.running = True
+ worker._control_server.stop = AsyncMock()
+ for _ in range(2):
+ with pytest.raises(RuntimeError, match="shutdown failed"):
+ await worker.stop()
+ worker._worker.shutdown.assert_awaited_once()
+ worker._control_server.stop.assert_not_awaited()
+
+
+@patch("airweave.domains.temporal.worker.create_activities", return_value=[])
+@patch("airweave.domains.temporal.worker.get_workflows", return_value=[])
+@pytest.mark.parametrize("blocked_phase", ["control", "client"])
+@patch("airweave.domains.temporal.client.close", new_callable=AsyncMock)
+@patch("airweave.domains.temporal.worker.Worker")
+@patch("temporalio.runtime.Runtime")
+async def test_stop_during_startup_waits_then_prevents_polling(
+ mock_runtime_cls,
+ mock_worker_cls,
+ mock_client_close,
+ mock_workflows,
+ mock_activities,
+ blocked_phase,
+):
+ """Cleanup waits for resource creation and no SDK worker starts after stop."""
+ from airweave.domains.temporal.worker import TemporalWorker
+
+ worker = TemporalWorker(_make_config())
+ entered = asyncio.Event()
+ release = asyncio.Event()
+
+ async def blocked(*args, **kwargs):
+ entered.set()
+ await release.wait()
+ return MagicMock()
+
+ worker._control_server.start = AsyncMock(
+ side_effect=blocked if blocked_phase == "control" else None
+ )
+ worker._control_server.stop = AsyncMock()
+ mock_worker_cls.return_value.run = AsyncMock()
+ with patch("airweave.domains.temporal.client.get_client", new_callable=AsyncMock) as client:
+ if blocked_phase == "client":
+ client.side_effect = blocked
+ starting = asyncio.create_task(worker.start())
+ await entered.wait()
+ stopping = asyncio.create_task(worker.stop())
+ try:
+ await asyncio.sleep(0)
+ await asyncio.sleep(0)
+ assert not stopping.done()
+ worker._control_server.stop.assert_not_awaited()
+ mock_client_close.assert_not_awaited()
+ finally:
+ release.set()
+ await asyncio.gather(starting, stopping)
+ mock_worker_cls.assert_not_called()
+ if blocked_phase == "control":
+ client.assert_not_awaited()
+ worker._control_server.stop.assert_awaited_once()
+ mock_client_close.assert_awaited_once()
diff --git a/backend/airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py b/backend/airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py
index d9caf2f067..3bd7406eb7 100644
--- a/backend/airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py
+++ b/backend/airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py
@@ -1,58 +1,79 @@
-"""Tests for the OCR guard in the Temporal worker main()."""
+"""OCR capability does not prevent the real worker composition from starting."""
-from unittest.mock import patch
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+import httpx
import pytest
import airweave.core.container as container_mod
+from airweave.core.config import settings
+from airweave.core.container import factory
+from airweave.domains.embedders import config as embedding_config
+from airweave.domains.temporal import worker
@pytest.mark.asyncio
-async def test_raises_system_exit_when_ocr_is_none(test_container):
- """Worker must abort when the container has no OCR provider."""
- no_ocr = test_container.replace(ocr_provider=None)
- saved = container_mod.container
+@pytest.mark.parametrize("health_status", [200, 503])
+async def test_worker_composes_without_ocr_and_starts_supported_activities(
+ monkeypatch, tmp_path, health_status
+):
+ """Real DI composition; substitute Temporal lifecycle and inference health HTTP."""
+ configured = settings.model_copy(
+ update={
+ "MISTRAL_API_KEY": "",
+ "DOCLING_BASE_URL": "",
+ "STORAGE_BACKEND": "filesystem",
+ "STORAGE_PATH": str(tmp_path),
+ "STRIPE_ENABLED": False,
+ "ANALYTICS_ENABLED": False,
+ }
+ )
+ monkeypatch.setattr(worker, "settings", configured)
+ monkeypatch.setattr(factory, "DENSE_EMBEDDER", "local_minilm")
+ monkeypatch.setattr(factory, "EMBEDDING_DIMENSIONS", 384)
+ monkeypatch.setattr(embedding_config, "DENSE_EMBEDDER", "local_minilm")
+ monkeypatch.setattr(embedding_config, "EMBEDDING_DIMENSIONS", 384)
+ # Stub only the external inference health HTTP call, not embedding validation.
+ original_client = httpx.Client
- try:
- container_mod.container = no_ocr
- with patch.object(
- container_mod,
- "initialize_container",
- ):
- from airweave.domains.temporal.worker import main
-
- with pytest.raises(SystemExit) as exc_info:
- await main()
-
- assert exc_info.value.code == 1
- finally:
- container_mod.container = saved
-
-
-@pytest.mark.asyncio
-async def test_passes_guard_when_ocr_present(test_container):
- """Worker proceeds past OCR guard when provider is present."""
- saved = container_mod.container
-
- try:
- container_mod.container = test_container
- with (
- patch.object(
- container_mod,
- "initialize_container",
+ def health_client(*args, **kwargs):
+ return original_client(
+ transport=httpx.MockTransport(
+ lambda request: httpx.Response(health_status)
+ if request.url.path == "/.well-known/ready"
+ else httpx.Response(503)
),
- patch(
- "airweave.domains.temporal.worker.WorkerConfig.from_settings",
- side_effect=SystemExit(99),
- ),
- ):
- from airweave.domains.temporal.worker import main
-
- with pytest.raises(SystemExit) as exc_info:
- await main()
+ **kwargs,
+ )
- # Exit code 99 means we passed the OCR guard and hit
- # the WorkerConfig stub — the guard did not fire.
- assert exc_info.value.code == 99
- finally:
- container_mod.container = saved
+ monkeypatch.setattr(
+ embedding_config,
+ "httpx",
+ SimpleNamespace(
+ Client=health_client,
+ Timeout=httpx.Timeout,
+ ConnectError=httpx.ConnectError,
+ TimeoutException=httpx.TimeoutException,
+ HTTPStatusError=httpx.HTTPStatusError,
+ ),
+ )
+ monkeypatch.setattr(container_mod, "container", None)
+ lifecycle = MagicMock(start=AsyncMock(), stop=AsyncMock())
+ monkeypatch.setattr(worker, "TemporalWorker", lambda _: lifecycle)
+ loop = asyncio.get_running_loop()
+ monkeypatch.setattr(loop, "add_signal_handler", lambda *_: None)
+ monkeypatch.setattr(loop, "remove_signal_handler", lambda *_: True)
+ if health_status != 200:
+ with pytest.raises(embedding_config.EmbeddingConfigError, match="not reachable"):
+ await worker.main()
+ lifecycle.start.assert_not_awaited()
+ return
+ await worker.main()
+ assert container_mod.container.ocr_provider is None
+ assert container_mod.container.converter_registry.for_extension(".txt") is not None
+ assert container_mod.container.converter_registry.for_extension(".pdf") is not None
+ assert container_mod.container.converter_registry.for_extension(".png") is None
+ lifecycle.start.assert_awaited_once()
+ lifecycle.stop.assert_awaited_once()
diff --git a/backend/airweave/domains/temporal/worker/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py
index 58defe7211..f6da596670 100644
--- a/backend/airweave/domains/temporal/worker/wiring.py
+++ b/backend/airweave/domains/temporal/worker/wiring.py
@@ -33,6 +33,20 @@ def create_activities() -> list:
if container is None:
raise RuntimeError("Container not initialized — cannot wire activities")
+ from airweave.db.session import get_db_context
+ from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore
+ from airweave.domains.entities.canonical.projector import CanonicalProjector
+ from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor
+ from airweave.domains.temporal.activities.cleanup_projection_generations import (
+ CleanupProjectionGenerationsActivity,
+ )
+ from airweave.domains.temporal.activities.discover_native_projection import (
+ DiscoverNativeProjectionActivity,
+ )
+ from airweave.domains.temporal.activities.project_canonical_records import (
+ ProjectCanonicalRecordsActivity,
+ )
+
email_service = container.email_service
event_bus = container.event_bus
sync_service = container.sync_service
@@ -50,6 +64,21 @@ def create_activities() -> list:
logger.debug("Wiring activities with container dependencies")
return [
+ DiscoverNativeProjectionActivity(container.source_registry).run,
+ CleanupProjectionGenerationsActivity(container.storage_backend).run,
+ ProjectCanonicalRecordsActivity(
+ projector=CanonicalProjector(
+ CanonicalProjectionStore(),
+ get_db_context,
+ ChunkEmbedProcessor(
+ container.converter_registry,
+ container.dense_embedder,
+ container.sparse_embedder,
+ ),
+ container.storage_backend,
+ ),
+ source_registry=container.source_registry,
+ ).run,
RunSyncActivity(
sync_service=sync_service,
sync_repo=sync_repo,
@@ -101,8 +130,16 @@ def get_workflows() -> list:
CleanupSyncDataWorkflow,
RunSourceConnectionWorkflow,
)
+ from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+ )
+ from airweave.domains.temporal.workflows.recover_native_projection import (
+ RecoverNativeProjectionWorkflow,
+ )
return [
+ RecoverNativeProjectionWorkflow,
+ ProjectCanonicalRecordsWorkflow,
RunSourceConnectionWorkflow,
CleanupStuckSyncJobsWorkflow,
CleanupSyncDataWorkflow,
diff --git a/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py b/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py
index 3896f18511..14f9cebd65 100644
--- a/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py
+++ b/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py
@@ -5,10 +5,17 @@
from datetime import timedelta
from temporalio import workflow
-from temporalio.common import RetryPolicy
+from temporalio.common import RetryPolicy, WorkflowIDReusePolicy
+from temporalio.exceptions import WorkflowAlreadyStartedError
with workflow.unsafe.imports_passed_through():
from airweave.domains.temporal.activities import cleanup_stuck_sync_jobs_activity
+ from airweave.domains.temporal.activities.cleanup_projection_generations import (
+ CleanupProjectionGenerationsActivity,
+ )
+ from airweave.domains.temporal.workflows.recover_native_projection import (
+ RecoverNativeProjectionWorkflow,
+ )
_CLEANUP_TIMEOUT = timedelta(minutes=5)
_CLEANUP_RETRY = RetryPolicy(
@@ -29,9 +36,27 @@ class CleanupStuckSyncJobsWorkflow:
@workflow.run
async def run(self) -> None:
- """Run the cleanup workflow."""
- await workflow.execute_activity(
- cleanup_stuck_sync_jobs_activity, # type: ignore[arg-type]
- start_to_close_timeout=_CLEANUP_TIMEOUT,
- retry_policy=_CLEANUP_RETRY,
- )
+ """Preserve cleanup failures while independently dispatching native recovery."""
+ try:
+ await workflow.execute_activity(
+ cleanup_stuck_sync_jobs_activity, # type: ignore[arg-type]
+ start_to_close_timeout=_CLEANUP_TIMEOUT,
+ retry_policy=_CLEANUP_RETRY,
+ )
+ if workflow.patched("canonical-generation-gc-v1"):
+ await workflow.execute_activity(
+ CleanupProjectionGenerationsActivity.run,
+ start_to_close_timeout=_CLEANUP_TIMEOUT,
+ retry_policy=_CLEANUP_RETRY,
+ )
+ finally:
+ if workflow.patched("native-projection-recovery-v1"):
+ try:
+ await workflow.start_child_workflow(
+ RecoverNativeProjectionWorkflow.run,
+ id="native-projection-recovery",
+ id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+ except WorkflowAlreadyStartedError:
+ pass
diff --git a/backend/airweave/domains/temporal/workflows/project_canonical_records.py b/backend/airweave/domains/temporal/workflows/project_canonical_records.py
new file mode 100644
index 0000000000..84371d9d09
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/project_canonical_records.py
@@ -0,0 +1,57 @@
+"""Independent retryable projection on the existing Temporal worker/task queue."""
+
+from datetime import timedelta
+
+from temporalio import workflow
+from temporalio.common import RetryPolicy
+from temporalio.exceptions import ApplicationError
+
+with workflow.unsafe.imports_passed_through():
+ from airweave.domains.temporal.activities.project_canonical_records import (
+ ProjectCanonicalRecordsActivity,
+ )
+
+
+@workflow.defn
+class ProjectCanonicalRecordsWorkflow:
+ """Scan pending rows; publication failures never change capture-job status."""
+
+ @workflow.run
+ async def run(
+ self,
+ organization_id: str,
+ sync_id: str,
+ after_id: str | None = None,
+ failures: int = 0,
+ sweep: int = 0,
+ skip_failed: bool = False,
+ ) -> None:
+ """Drain bounded pages; retry failed sweeps without recapturing providers."""
+ for _ in range(100):
+ result = await workflow.execute_activity(
+ ProjectCanonicalRecordsActivity.run,
+ args=(
+ [organization_id, sync_id, after_id, True]
+ if skip_failed
+ else [organization_id, sync_id, after_id]
+ ),
+ start_to_close_timeout=timedelta(minutes=30),
+ retry_policy=RetryPolicy(maximum_attempts=3),
+ )
+ failures += result["failed"]
+ after_id = result["after_id"]
+ if not result["has_more"]:
+ if not failures:
+ return
+ if skip_failed or sweep >= 2:
+ raise ApplicationError(
+ "Canonical projection still has failed records; retained as pending",
+ non_retryable=True,
+ )
+ await workflow.sleep(timedelta(seconds=30 * (2**sweep)))
+ workflow.continue_as_new(args=[organization_id, sync_id, None, 0, sweep + 1])
+ # Bound history even for very large sources.
+ args = [organization_id, sync_id, after_id, failures, sweep]
+ if skip_failed:
+ args.append(True)
+ workflow.continue_as_new(args=args)
diff --git a/backend/airweave/domains/temporal/workflows/recover_native_projection.py b/backend/airweave/domains/temporal/workflows/recover_native_projection.py
new file mode 100644
index 0000000000..f637ce431a
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/recover_native_projection.py
@@ -0,0 +1,58 @@
+"""Recover canonical publications independently of capture and cleanup outcomes."""
+
+from datetime import timedelta
+
+from temporalio import workflow
+from temporalio.common import RetryPolicy, WorkflowIDReusePolicy
+from temporalio.exceptions import WorkflowAlreadyStartedError
+
+with workflow.unsafe.imports_passed_through():
+ from airweave.domains.entities.canonical.projection_models import (
+ ProjectionSourcePage,
+ canonical_projection_workflow_id,
+ )
+ from airweave.domains.temporal.activities.discover_native_projection import (
+ DiscoverNativeProjectionActivity,
+ )
+ from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+ )
+
+
+@workflow.defn
+class RecoverNativeProjectionWorkflow:
+ """Page through durable pending records; never recapture or retry known failures."""
+
+ @workflow.run
+ async def run(self, after_sync_id: str | None = None) -> None:
+ """Advance every source page independently of its child projection outcome."""
+ # Keyset progress lives in workflow history, not another checkpoint table.
+ # A bad/active source cannot keep every later source off the first page.
+ for _ in range(100):
+ page = ProjectionSourcePage.model_validate(
+ await workflow.execute_activity(
+ DiscoverNativeProjectionActivity.run,
+ after_sync_id,
+ start_to_close_timeout=timedelta(minutes=5),
+ retry_policy=RetryPolicy(maximum_attempts=3),
+ )
+ )
+ for source in page.sources:
+ try:
+ await workflow.start_child_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[str(source.organization_id), str(source.sync_id), None, 0, 0, True],
+ id=(
+ canonical_projection_workflow_id(source.organization_id, source.sync_id)
+ if workflow.patched("canonical-projection-identity-v1")
+ else f"native-projection:{source.organization_id}:{source.sync_id}"
+ ),
+ id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+ except WorkflowAlreadyStartedError:
+ pass # An existing run owns this source; keep walking other pages.
+ if page.next_cursor is None:
+ return
+ after_sync_id = str(page.next_cursor)
+ workflow.continue_as_new(args=[after_sync_id])
diff --git a/backend/airweave/domains/temporal/workflows/run_source_connection.py b/backend/airweave/domains/temporal/workflows/run_source_connection.py
index f2966a83bc..8e822fb736 100644
--- a/backend/airweave/domains/temporal/workflows/run_source_connection.py
+++ b/backend/airweave/domains/temporal/workflows/run_source_connection.py
@@ -12,14 +12,18 @@
from typing import Any, Dict, Optional
from temporalio import workflow
-from temporalio.common import RetryPolicy
+from temporalio.common import RetryPolicy, WorkflowIDReusePolicy
from temporalio.exceptions import (
ActivityError,
ApplicationError,
+ WorkflowAlreadyStartedError,
is_cancelled_exception,
)
with workflow.unsafe.imports_passed_through():
+ from airweave.domains.entities.canonical.projection_models import (
+ canonical_projection_workflow_id,
+ )
from airweave.domains.temporal.activities import (
create_sync_job_activity,
run_sync_activity,
@@ -31,6 +35,9 @@
CLASSIFIED_USER_ERROR_TYPE,
ORPHANED_SYNC_ERROR_TYPE,
)
+ from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+ )
# ---------------------------------------------------------------------------
# Execution policies
@@ -189,22 +196,50 @@ async def _execute_sync(
local_development = ctx_dict.get("local_development", False)
heartbeat_timeout = _HEARTBEAT_TIMEOUT_LOCAL if local_development else _HEARTBEAT_TIMEOUT
- await workflow.execute_activity(
- run_sync_activity,
- args=[
- sync_dict,
- sync_job_dict,
- collection_dict,
- connection_dict,
- ctx_dict,
- access_token,
- force_full_sync,
- ],
- start_to_close_timeout=_SYNC_TIMEOUT,
- heartbeat_timeout=heartbeat_timeout,
- cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
- retry_policy=_NO_RETRY,
- )
+ try:
+ await workflow.execute_activity(
+ run_sync_activity,
+ args=[
+ sync_dict,
+ sync_job_dict,
+ collection_dict,
+ connection_dict,
+ ctx_dict,
+ access_token,
+ force_full_sync,
+ ],
+ start_to_close_timeout=_SYNC_TIMEOUT,
+ heartbeat_timeout=heartbeat_timeout,
+ cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+ retry_policy=_NO_RETRY,
+ )
+ finally:
+ # Starting the child is durable workflow history. Capture may be partial;
+ # projection consumes only committed rows and owns its own retries/status.
+ organization_id = str(ctx_dict["organization"]["id"])
+ sync_id = str(sync_dict["id"])
+ if workflow.patched("canonical-projection-identity-v1"):
+ try:
+ await asyncio.shield(
+ workflow.start_child_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[organization_id, sync_id, None, 0, 0, True],
+ id=canonical_projection_workflow_id(organization_id, sync_id),
+ id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+ )
+ except WorkflowAlreadyStartedError:
+ pass # Maintenance or an explicit operator run already owns projection.
+ else:
+ await asyncio.shield(
+ workflow.start_child_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[organization_id, sync_id],
+ id=f"projection:{sync_dict['id']}:{workflow.info().run_id}",
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+ )
# ------------------------------------------------------------------
# Phase 3: State transitions
diff --git a/backend/airweave/domains/temporal/workflows/tests/test_canonical_projection.py b/backend/airweave/domains/temporal/workflows/tests/test_canonical_projection.py
new file mode 100644
index 0000000000..f805dcc275
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/tests/test_canonical_projection.py
@@ -0,0 +1,42 @@
+"""Real Temporal test-server proof of independent bounded projection retries."""
+
+from uuid import uuid4
+
+from temporalio import activity
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+)
+
+
+async def test_projection_retries_failed_sweep_without_provider_capture():
+ calls = []
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(organization: str, sync: str, after: str | None = None) -> dict:
+ calls.append((organization, sync, after))
+ return {
+ "after_id": None,
+ "has_more": False,
+ "failed": int(len(calls) == 1),
+ "published": int(len(calls) > 1),
+ "superseded": 0,
+ }
+
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ async with Worker(
+ env.client,
+ task_queue="canonical-projection-test",
+ workflows=[ProjectCanonicalRecordsWorkflow],
+ activities=[project],
+ ):
+ await env.client.execute_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[str(uuid4()), str(uuid4())],
+ id=str(uuid4()),
+ task_queue="canonical-projection-test",
+ )
+ assert len(calls) == 2
+ assert calls[0] == calls[1]
diff --git a/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py b/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py
index 31ade8b619..66f3447945 100644
--- a/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py
+++ b/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py
@@ -3,6 +3,7 @@
import uuid
import pytest
+from temporalio import activity
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
@@ -12,6 +13,9 @@
from airweave.domains.temporal.workflows.cleanup_sync_data import (
CleanupSyncDataWorkflow,
)
+from airweave.domains.temporal.workflows.recover_native_projection import (
+ RecoverNativeProjectionWorkflow,
+)
from .conftest import (
COLLECTION_ID,
@@ -27,20 +31,32 @@
@pytest.mark.unit
async def test_cleanup_stuck_sync_jobs_calls_activity():
+ @activity.defn(name="cleanup_projection_generations_activity")
+ async def cleanup_generations():
+ recorder.record("cleanup_generations", ())
+
+ @activity.defn(name="discover_native_projection_activity")
+ async def discover_native(after: str | None = None):
+ return {"sources": [], "next_cursor": None}
+
recorder = ActivityRecorder()
async with await WorkflowEnvironment.start_time_skipping() as env:
async with Worker(
env.client,
task_queue=TASK_QUEUE,
- workflows=[CleanupStuckSyncJobsWorkflow],
- activities=[mock_cleanup_stuck_sync_jobs(recorder)],
+ workflows=[CleanupStuckSyncJobsWorkflow, RecoverNativeProjectionWorkflow],
+ activities=[
+ mock_cleanup_stuck_sync_jobs(recorder), cleanup_generations, discover_native
+ ],
):
await env.client.execute_workflow(
CleanupStuckSyncJobsWorkflow.run,
id=f"test-cleanup-stuck-{uuid.uuid4()}",
task_queue=TASK_QUEUE,
)
+ await env.client.get_workflow_handle("native-projection-recovery").result()
+ assert recorder.called("cleanup_generations")
assert recorder.called("cleanup_stuck")
assert recorder.call_count("cleanup_stuck") == 1
diff --git a/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py b/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py
new file mode 100644
index 0000000000..f621d831b8
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py
@@ -0,0 +1,308 @@
+"""Real Temporal child deduplication, page progress and bounded failure recovery."""
+
+import asyncio
+from datetime import timedelta
+from uuid import UUID, uuid4
+
+import pytest
+from temporalio import activity, workflow
+from temporalio.client import WorkflowFailureError
+from temporalio.common import RetryPolicy
+from temporalio.exceptions import ApplicationError
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Replayer, Worker
+
+from airweave.domains.temporal.workflows.cleanup_stuck_sync_jobs import CleanupStuckSyncJobsWorkflow
+from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+)
+from airweave.domains.temporal.workflows.recover_native_projection import (
+ RecoverNativeProjectionWorkflow,
+)
+
+
+@workflow.defn(name="CleanupStuckSyncJobsWorkflow")
+class PreviousCleanupWorkflow:
+ """Previous command sequence, used to generate real pre-recovery history."""
+
+ @workflow.run
+ async def run(self) -> None:
+ retry = RetryPolicy(
+ maximum_attempts=3,
+ initial_interval=timedelta(seconds=10),
+ maximum_interval=timedelta(seconds=60),
+ )
+ await workflow.execute_activity(
+ "cleanup_stuck_sync_jobs_activity",
+ start_to_close_timeout=timedelta(minutes=5),
+ retry_policy=retry,
+ )
+ if workflow.patched("canonical-generation-gc-v1"):
+ await workflow.execute_activity(
+ "cleanup_projection_generations_activity",
+ start_to_close_timeout=timedelta(minutes=5),
+ retry_policy=retry,
+ )
+
+
+async def test_previous_maintenance_history_replays_without_recovery_commands():
+ @activity.defn(name="cleanup_stuck_sync_jobs_activity")
+ async def cleanup():
+ pass
+
+ @activity.defn(name="cleanup_projection_generations_activity")
+ async def generations():
+ pass
+
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "native-replay-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[PreviousCleanupWorkflow],
+ activities=[cleanup, generations],
+ ):
+ handle = await env.client.start_workflow(
+ PreviousCleanupWorkflow.run, id=uuid4().hex, task_queue=queue
+ )
+ await handle.result()
+ history = await handle.fetch_history()
+ await Replayer(workflows=[CleanupStuckSyncJobsWorkflow]).replay_workflow(history)
+
+
+async def test_recovery_walks_later_page_despite_active_source_and_cleanup_failure():
+ organization = str(uuid4())
+ sources = [str(UUID(int=i)) for i in range(1, 22)]
+ calls, cursors = [], []
+ started, release = asyncio.Event(), asyncio.Event()
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(org: str, sync: str, after: str | None = None, skip_failed: bool = False):
+ assert org == organization and skip_failed
+ calls.append(sync)
+ if sync == sources[0]:
+ started.set()
+ await release.wait()
+ return {"after_id": None, "has_more": False, "failed": 0, "published": 1}
+
+ @activity.defn(name="discover_native_projection_activity")
+ async def discover(after: str | None = None):
+ cursors.append(after)
+ page = sources[:20] if after is None else sources[20:]
+ return {
+ "sources": [{"organization_id": organization, "sync_id": sync} for sync in page],
+ "next_cursor": sources[19] if after is None else None,
+ }
+
+ @activity.defn(name="cleanup_stuck_sync_jobs_activity")
+ async def cleanup():
+ raise ApplicationError("Unrelated provider cleanup failed", non_retryable=True)
+
+ @activity.defn(name="cleanup_projection_generations_activity")
+ async def generations():
+ pass
+
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "native-recovery-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[
+ CleanupStuckSyncJobsWorkflow,
+ RecoverNativeProjectionWorkflow,
+ ProjectCanonicalRecordsWorkflow,
+ ],
+ activities=[project, discover, cleanup, generations],
+ ):
+ active = await env.client.start_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[organization, sources[0], None, 0, 0, True],
+ id=f"canonical-projection:{organization}:{sources[0]}",
+ task_queue=queue,
+ )
+ await asyncio.wait_for(started.wait(), timeout=20)
+ try:
+ with pytest.raises(WorkflowFailureError):
+ await env.client.execute_workflow(
+ CleanupStuckSyncJobsWorkflow.run, id=uuid4().hex, task_queue=queue
+ )
+ await env.client.get_workflow_handle("native-projection-recovery").result()
+ release.set()
+ await active.result()
+ for sync in sources[1:]:
+ await env.client.get_workflow_handle(
+ f"canonical-projection:{organization}:{sync}"
+ ).result()
+ finally:
+ release.set()
+ assert cursors == [None, sources[19]]
+ assert sorted(calls) == sorted(sources)
+
+
+async def test_automatic_recovery_reports_failure_without_restarting_failed_sweeps():
+ calls = []
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(org: str, sync: str, after: str | None = None, skip_failed: bool = False):
+ calls.append(skip_failed)
+ return {"after_id": None, "has_more": False, "failed": 1, "published": 0}
+
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "native-failure-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[ProjectCanonicalRecordsWorkflow],
+ activities=[project],
+ ):
+ with pytest.raises(WorkflowFailureError):
+ await env.client.execute_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[str(uuid4()), str(uuid4()), None, 0, 0, True],
+ id=uuid4().hex,
+ task_queue=queue,
+ )
+ assert calls == [True]
+
+
+async def test_maintenance_projects_during_capture_and_finalizer_reuses_active_projection():
+ """Capture need not finish for search publication, and both triggers share ownership."""
+ from airweave.domains.temporal.workflows.run_source_connection import (
+ RunSourceConnectionWorkflow,
+ )
+
+ from .conftest import (
+ ORG_ID,
+ SYNC_ID,
+ ActivityRecorder,
+ make_collection_dict,
+ make_connection_dict,
+ make_ctx_dict,
+ make_sync_dict,
+ make_sync_job_dict,
+ mock_transition_sync_job,
+ )
+
+ capturing, projected = asyncio.Event(), asyncio.Event()
+ finish_capture, finish_projection = asyncio.Event(), asyncio.Event()
+ calls = []
+
+ @activity.defn(name="run_sync_activity")
+ async def capture(
+ sync: dict,
+ job: dict,
+ collection: dict,
+ connection: dict,
+ context: dict,
+ token: str | None,
+ force: bool,
+ ):
+ capturing.set()
+ await finish_capture.wait()
+
+ @activity.defn(name="discover_native_projection_activity")
+ async def discover(after: str | None = None):
+ assert capturing.is_set() and not finish_capture.is_set()
+ return {"sources": [{"organization_id": ORG_ID, "sync_id": SYNC_ID}], "next_cursor": None}
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(org: str, sync: str, after: str | None = None, skip_failed: bool = False):
+ assert skip_failed
+ calls.append(sync)
+ projected.set()
+ await finish_projection.wait()
+ return {"after_id": None, "has_more": False, "failed": 0, "published": 1}
+
+ recorder = ActivityRecorder()
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "concurrent-projection-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[
+ RunSourceConnectionWorkflow,
+ RecoverNativeProjectionWorkflow,
+ ProjectCanonicalRecordsWorkflow,
+ ],
+ activities=[capture, discover, project, mock_transition_sync_job(recorder)],
+ ):
+ handle = await env.client.start_workflow(
+ RunSourceConnectionWorkflow.run,
+ args=[
+ make_sync_dict(),
+ make_sync_job_dict(),
+ make_collection_dict(),
+ make_connection_dict(),
+ make_ctx_dict(),
+ ],
+ id=uuid4().hex,
+ task_queue=queue,
+ )
+ try:
+ await asyncio.wait_for(capturing.wait(), 20)
+ await env.client.execute_workflow(
+ RecoverNativeProjectionWorkflow.run, id=uuid4().hex, task_queue=queue
+ )
+ await asyncio.wait_for(projected.wait(), 20)
+ assert not finish_capture.is_set()
+ finish_capture.set()
+ await handle.result()
+ assert recorder.called("transition_completed")
+ assert calls == [SYNC_ID]
+ finally:
+ finish_capture.set()
+ finish_projection.set()
+ await env.client.get_workflow_handle(
+ f"canonical-projection:{ORG_ID}:{SYNC_ID}"
+ ).result()
+
+
+@workflow.defn(name="RecoverNativeProjectionWorkflow")
+class PreviousRecoveryWorkflow:
+ """The historical native-only child ID must remain replayable."""
+
+ @workflow.run
+ async def run(self, after_sync_id: str | None = None) -> None:
+ page = await workflow.execute_activity(
+ "discover_native_projection_activity",
+ after_sync_id,
+ start_to_close_timeout=timedelta(minutes=5),
+ retry_policy=RetryPolicy(maximum_attempts=3),
+ )
+ for source in page["sources"]:
+ await workflow.start_child_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[source["organization_id"], source["sync_id"], None, 0, 0, True],
+ id=f"native-projection:{source['organization_id']}:{source['sync_id']}",
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+
+
+async def test_previous_source_projection_identity_replays():
+ organization, sync = str(uuid4()), str(uuid4())
+
+ @activity.defn(name="discover_native_projection_activity")
+ async def discover(after: str | None = None):
+ return {
+ "sources": [{"organization_id": organization, "sync_id": sync}],
+ "next_cursor": None,
+ }
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(org: str, sync: str, after: str | None = None, skip_failed: bool = False):
+ return {"after_id": None, "has_more": False, "failed": 0, "published": 1}
+
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "previous-projection-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[PreviousRecoveryWorkflow, ProjectCanonicalRecordsWorkflow],
+ activities=[discover, project],
+ ):
+ handle = await env.client.start_workflow(
+ PreviousRecoveryWorkflow.run, id=uuid4().hex, task_queue=queue
+ )
+ await handle.result()
+ history = await handle.fetch_history()
+ await Replayer(workflows=[RecoverNativeProjectionWorkflow]).replay_workflow(history)
diff --git a/backend/airweave/domains/temporal/workflows/tests/test_reproject_command.py b/backend/airweave/domains/temporal/workflows/tests/test_reproject_command.py
new file mode 100644
index 0000000000..08950a1ed9
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/tests/test_reproject_command.py
@@ -0,0 +1,47 @@
+"""Operator retry intent must not disappear behind automatic projection."""
+
+from argparse import Namespace
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import pytest
+from temporalio.exceptions import WorkflowAlreadyStartedError
+
+from airweave.domains.entities.canonical.reprojection import ReprojectionPlan
+from scripts import reproject_canonical
+
+
+async def test_busy_projection_preserves_committed_version_but_reports_retry(monkeypatch, capsys):
+ organization, sync = uuid4(), uuid4()
+ args = Namespace(organization=organization, sync=sync, expect=2, target=3, apply=True)
+ plan = ReprojectionPlan(
+ organization_id=organization,
+ sync_id=sync,
+ source_connection_id=uuid4(),
+ provider="slack",
+ current_version=2,
+ target_version=3,
+ captured_records=1,
+ target_pending_records=1,
+ workflow_id=f"canonical-projection:{organization}:{sync}",
+ applied=True,
+ )
+ session = AsyncMock()
+ context = MagicMock()
+ context.__aenter__ = AsyncMock(return_value=session)
+ context.__aexit__ = AsyncMock(return_value=False)
+ monkeypatch.setattr(reproject_canonical, "AsyncSessionLocal", lambda: context)
+ monkeypatch.setattr(reproject_canonical, "plan_reprojection", AsyncMock(return_value=plan))
+ client = AsyncMock()
+ client.start_workflow.side_effect = WorkflowAlreadyStartedError(plan.workflow_id, "projection")
+ monkeypatch.setattr(reproject_canonical, "get_client", AsyncMock(return_value=client))
+
+ with pytest.raises(SystemExit, match="retry this exact command"):
+ await reproject_canonical.run(args)
+
+ session.commit.assert_awaited_once()
+ assert "accepted" not in capsys.readouterr().out
+ call = client.start_workflow.call_args.kwargs
+ assert call["id"] == plan.workflow_id
+ assert call["args"] == [str(organization), str(sync)]
+ assert "id_conflict_policy" not in call
diff --git a/backend/airweave/domains/temporal/workflows/tests/test_source_projection_replay.py b/backend/airweave/domains/temporal/workflows/tests/test_source_projection_replay.py
new file mode 100644
index 0000000000..4321c6af75
--- /dev/null
+++ b/backend/airweave/domains/temporal/workflows/tests/test_source_projection_replay.py
@@ -0,0 +1,156 @@
+"""Replay real pre-shared-ID source histories against the current workflow."""
+
+import asyncio
+from datetime import timedelta
+from typing import Any
+from uuid import uuid4
+
+import pytest
+from temporalio import activity, workflow
+from temporalio.client import WorkflowFailureError
+from temporalio.common import RetryPolicy
+from temporalio.exceptions import ApplicationError
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Replayer, Worker
+
+from airweave.domains.temporal.activities import run_sync_activity
+from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+)
+from airweave.domains.temporal.workflows.run_source_connection import RunSourceConnectionWorkflow
+
+from .conftest import (
+ ActivityRecorder,
+ make_collection_dict,
+ make_connection_dict,
+ make_ctx_dict,
+ make_sync_dict,
+ make_sync_job_dict,
+ mock_run_sync,
+ mock_transition_sync_job,
+)
+
+
+@workflow.defn(name="RunSourceConnectionWorkflow")
+class PreviousSourceWorkflow(RunSourceConnectionWorkflow):
+ """Keep the former finalizer's ID, arguments and execution policies verbatim."""
+
+ @workflow.run
+ async def run(
+ self,
+ sync_dict: dict[str, Any],
+ sync_job_dict: dict[str, Any] | None,
+ collection_dict: dict[str, Any],
+ connection_dict: dict[str, Any],
+ ctx_dict: dict[str, Any],
+ access_token: str | None = None,
+ force_full_sync: bool = False,
+ ) -> None:
+ await super().run(
+ sync_dict,
+ sync_job_dict,
+ collection_dict,
+ connection_dict,
+ ctx_dict,
+ access_token,
+ force_full_sync,
+ )
+
+ async def _execute_sync(
+ self,
+ sync_dict,
+ sync_job_dict,
+ collection_dict,
+ connection_dict,
+ ctx_dict,
+ access_token,
+ force_full_sync,
+ ) -> None:
+ heartbeat = (
+ timedelta(hours=1)
+ if ctx_dict.get("local_development", False)
+ else timedelta(minutes=15)
+ )
+ try:
+ await workflow.execute_activity(
+ run_sync_activity,
+ args=[
+ sync_dict,
+ sync_job_dict,
+ collection_dict,
+ connection_dict,
+ ctx_dict,
+ access_token,
+ force_full_sync,
+ ],
+ start_to_close_timeout=timedelta(days=7),
+ heartbeat_timeout=heartbeat,
+ cancellation_type=workflow.ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+ retry_policy=RetryPolicy(maximum_attempts=1),
+ )
+ finally:
+ await asyncio.shield(
+ workflow.start_child_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[str(ctx_dict["organization"]["id"]), str(sync_dict["id"])],
+ id=f"projection:{sync_dict['id']}:{workflow.info().run_id}",
+ parent_close_policy=workflow.ParentClosePolicy.ABANDON,
+ )
+ )
+
+
+@pytest.mark.parametrize("local,fails", [(False, False), (True, True)])
+async def test_old_source_projection_finalizer_history_replays(local, fails):
+ recorder = ActivityRecorder()
+
+ @activity.defn(name="project_canonical_records_activity")
+ async def project(organization: str, sync: str, after: str | None = None):
+ return {"after_id": None, "has_more": False, "failed": 0, "published": 0}
+
+ ctx = make_ctx_dict()
+ ctx["local_development"] = local
+ sync = make_sync_dict()
+ error = ApplicationError("Synthetic capture failure", non_retryable=True) if fails else None
+ async with await WorkflowEnvironment.start_time_skipping() as env:
+ queue = "source-projection-replay-" + uuid4().hex
+ async with Worker(
+ env.client,
+ task_queue=queue,
+ workflows=[PreviousSourceWorkflow, ProjectCanonicalRecordsWorkflow],
+ activities=[
+ mock_run_sync(recorder, error),
+ mock_transition_sync_job(recorder),
+ project,
+ ],
+ ):
+ handle = await env.client.start_workflow(
+ PreviousSourceWorkflow.run,
+ args=[
+ sync,
+ make_sync_job_dict(),
+ make_collection_dict(),
+ make_connection_dict(),
+ ctx,
+ "synthetic-token",
+ True,
+ ],
+ id=uuid4().hex,
+ task_queue=queue,
+ )
+ if fails:
+ with pytest.raises(WorkflowFailureError):
+ await handle.result()
+ else:
+ await handle.result()
+ history = await handle.fetch_history()
+ children = [
+ event.start_child_workflow_execution_initiated_event_attributes
+ for event in history.events
+ if event.HasField("start_child_workflow_execution_initiated_event_attributes")
+ ]
+ assert len(children) == 1
+ assert children[0].workflow_id.startswith(f"projection:{sync['id']}:")
+ assert len(children[0].input.payloads) == 2
+ await env.client.get_workflow_handle(children[0].workflow_id).result()
+ await Replayer(workflows=[RunSourceConnectionWorkflow]).replay_workflow(history)
+ assert recorder.called("transition_failed" if fails else "transition_completed")
diff --git a/backend/airweave/main.py b/backend/airweave/main.py
index 2702ea968b..2064278cb6 100644
--- a/backend/airweave/main.py
+++ b/backend/airweave/main.py
@@ -4,9 +4,6 @@
and unhandled exceptions.
"""
-import os
-import subprocess
-import sys
from contextlib import asynccontextmanager
from fastapi import FastAPI
@@ -34,6 +31,7 @@
)
from airweave.api.router import TrailingSlashRouter
from airweave.api.v1.api import api_router
+from airweave.api.v1.endpoints.records import record_error_response
from airweave.core.config import settings
from airweave.core.exceptions import (
AirweaveException,
@@ -44,16 +42,16 @@
RateLimitExceededException,
)
from airweave.core.logging import logger
-from airweave.db.init_db import init_db
from airweave.db.session import AsyncSessionLocal
from airweave.domains.embedders.config import validate_embedding_config
+from airweave.domains.entities.canonical.store import CanonicalStoreError
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Lifespan context manager for startup and shutdown events.
- Initializes the DI container, runs alembic migrations, and syncs platform components.
+ Initialize runtime dependencies against an explicitly provisioned database.
"""
# Initialize the dependency injection container (fail fast if wiring is broken)
from airweave.core import container as container_mod
@@ -64,19 +62,6 @@ async def lifespan(app: FastAPI):
logger.info("Container initialized successfully")
async with AsyncSessionLocal() as db:
- if settings.RUN_ALEMBIC_MIGRATIONS:
- logger.info("Running alembic migrations...")
- env = os.environ.copy()
- backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
- env["PYTHONPATH"] = backend_dir
- subprocess.run(
- [sys.executable, "-m", "alembic", "upgrade", "head"],
- check=True,
- cwd=backend_dir,
- env=env,
- )
- await init_db(db)
-
# Reconcile embedding config against DB deployment metadata
await validate_embedding_config(db)
@@ -86,6 +71,10 @@ async def lifespan(app: FastAPI):
await container_mod.container.temporal_schedule_service.ensure_system_schedules()
logger.info("System Temporal schedules initialized successfully")
except Exception as e:
+ if "temporal" in settings.health_critical_probes:
+ # The cached health probe cannot initialize a missing connection.
+ # Let the process supervisor retry startup for a required service.
+ raise
logger.warning(
f"Failed to initialize system schedules (Temporal may not be available): {e}"
)
@@ -141,6 +130,7 @@ async def lifespan(app: FastAPI):
# Register custom Airweave exception handlers
app.exception_handler(AirweaveException)(airweave_exception_handler)
+app.exception_handler(CanonicalStoreError)(record_error_response)
# Default CORS origins - white labels and environment variables can extend this
CORS_ORIGINS = [
diff --git a/backend/airweave/models/__init__.py b/backend/airweave/models/__init__.py
index f755072eee..c95bf93808 100644
--- a/backend/airweave/models/__init__.py
+++ b/backend/airweave/models/__init__.py
@@ -4,10 +4,12 @@
from .api_key import APIKey
from .auth_provider import AuthProvider
from .billing_period import BillingPeriod
+from .capture_scan import CaptureScan
from .collection import Collection
from .connection import Connection
from .connection_init_session import ConnectionInitSession
from .entity import Entity
+from .entity_change import EntityChange
from .entity_count import EntityCount
from .entity_definition import EntityDefinition
from .entity_relation import EntityRelation
@@ -16,6 +18,7 @@
from .node_selection import NodeSelection
from .organization import Organization
from .organization_billing import OrganizationBilling
+from .owned_provisioning import OwnedProvisioning
from .processed_webhook_event import ProcessedWebhookEvent
from .redirect_session import RedirectSession
from .search_query import SearchQuery
@@ -31,12 +34,16 @@
from .vector_db_deployment_metadata import VectorDbDeploymentMetadata
__all__ = [
+ "OwnedProvisioning",
"AccessControlMembership",
"APIKey",
"AuthProvider",
"BillingPeriod",
"Collection",
"Entity",
+ "EntityChange",
+ "CaptureScan",
+ "ProjectionGeneration",
"EntityCount",
"Connection",
"ConnectionInitSession",
@@ -61,3 +68,5 @@
"UserOrganization",
"VectorDbDeploymentMetadata",
]
+
+from .projection_generation import ProjectionGeneration
diff --git a/backend/airweave/models/capture_scan.py b/backend/airweave/models/capture_scan.py
new file mode 100644
index 0000000000..fcae496086
--- /dev/null
+++ b/backend/airweave/models/capture_scan.py
@@ -0,0 +1,61 @@
+"""One recoverable whole-scope scan, independent of worker attempts."""
+
+from datetime import datetime
+from uuid import UUID
+
+from sqlalchemy import (
+ BigInteger,
+ CheckConstraint,
+ DateTime,
+ ForeignKey,
+ Index,
+ String,
+ UniqueConstraint,
+ text,
+)
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airweave.models._base import OrganizationBase
+
+
+class CaptureScan(OrganizationBase):
+ """Only the canonical capture transaction boundary mutates this state."""
+
+ __tablename__ = "capture_scan"
+ sync_id: Mapped[UUID] = mapped_column(ForeignKey("sync.id", ondelete="CASCADE"))
+ scope_key: Mapped[str] = mapped_column(String)
+ record_type: Mapped[str] = mapped_column(String)
+ container_id: Mapped[str | None] = mapped_column(String)
+ parent_record_id: Mapped[UUID | None] = mapped_column(
+ ForeignKey("entity.id", ondelete="CASCADE")
+ )
+ parent_visibility_epoch: Mapped[int | None] = mapped_column(BigInteger)
+ parent_verified_attempt_id: Mapped[UUID | None] = mapped_column(nullable=True)
+ parent_verified_revision: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
+ membership_attempt_id: Mapped[UUID | None]
+ cycle_id: Mapped[UUID]
+ sweep_id: Mapped[UUID]
+ revision: Mapped[int] = mapped_column(BigInteger)
+ phase: Mapped[str] = mapped_column(String)
+ fingerprint: Mapped[str] = mapped_column(String(64))
+ continuation: Mapped[dict] = mapped_column(JSONB)
+ execution_state: Mapped[dict | None] = mapped_column(JSONB, nullable=True)
+ started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
+ completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
+
+ __table_args__ = (
+ Index(
+ "idx_capture_scan_parent",
+ "sync_id",
+ "parent_record_id",
+ "record_type",
+ unique=True,
+ postgresql_where=text("parent_record_id IS NOT NULL"),
+ ),
+ UniqueConstraint("sync_id", "scope_key", name="uq_capture_scan_scope"),
+ CheckConstraint(
+ "phase IN ('collecting','reconciling','complete')", name="ck_capture_scan_phase"
+ ),
+ CheckConstraint("revision >= 1", name="ck_capture_scan_revision"),
+ )
diff --git a/backend/airweave/models/entity.py b/backend/airweave/models/entity.py
index a5aa81522d..2f3682a01f 100644
--- a/backend/airweave/models/entity.py
+++ b/backend/airweave/models/entity.py
@@ -1,9 +1,20 @@
"""Entity model."""
+from datetime import datetime
from typing import TYPE_CHECKING, Optional
from uuid import UUID
-from sqlalchemy import ForeignKey, Index, String, UniqueConstraint
+from sqlalchemy import (
+ BigInteger,
+ CheckConstraint,
+ DateTime,
+ ForeignKey,
+ Index,
+ String,
+ UniqueConstraint,
+ text,
+)
+from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship
from airweave.models._base import OrganizationBase
@@ -25,8 +36,8 @@ class Entity(OrganizationBase):
index=False, # Disabled: billions of rows, not selective, better indexes on sync_id
)
- sync_job_id: Mapped[UUID] = mapped_column(
- ForeignKey("sync_job.id", ondelete="CASCADE", name="fk_entity_sync_job_id"), nullable=False
+ sync_job_id: Mapped[Optional[UUID]] = mapped_column(
+ ForeignKey("sync_job.id", ondelete="SET NULL", name="fk_entity_sync_job_id"), nullable=True
)
sync_id: Mapped[UUID] = mapped_column(
ForeignKey("sync.id", ondelete="CASCADE", name="fk_entity_sync_id"), nullable=False
@@ -39,8 +50,35 @@ class Entity(OrganizationBase):
)
hash: Mapped[str] = mapped_column(String, nullable=False)
+ # Canonical source state. Revision zero identifies legacy metadata-only rows.
+ native_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ container_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ parent_record_type: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ parent_native_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ parent_container_id: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ visibility_epoch: Mapped[int] = mapped_column(BigInteger, default=1, server_default="1")
+ parent_visibility_epoch: Mapped[Optional[int]] = mapped_column(BigInteger)
+ source_payload: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
+ payload_schema_version: Mapped[int] = mapped_column(default=1, server_default="1")
+ record_revision: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ capture_hash: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ content_hash: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+ source_created_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
+ source_updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
+ observed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
+ deleted_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
+ removal_reason: Mapped[Optional[str]] = mapped_column(String)
+ completeness: Mapped[Optional[str]] = mapped_column(String)
+ blob_references: Mapped[Optional[list]] = mapped_column(JSONB)
+ last_seen_run_id: Mapped[Optional[UUID]] = mapped_column(nullable=True)
+ indexed_generation: Mapped[Optional[UUID]] = mapped_column(nullable=True)
+ indexed_chunk_count: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)
+ indexed_revision: Mapped[Optional[int]] = mapped_column(BigInteger)
+ indexed_pipeline_version: Mapped[Optional[int]] = mapped_column(BigInteger)
+ projection_error: Mapped[Optional[str]] = mapped_column(String)
+
# Add back references
- sync_job: Mapped["SyncJob"] = relationship(
+ sync_job: Mapped[Optional["SyncJob"]] = relationship(
"SyncJob",
back_populates="entities",
lazy="noload",
@@ -53,6 +91,36 @@ class Entity(OrganizationBase):
)
__table_args__ = (
+ CheckConstraint(
+ "visibility_epoch >= 1 AND "
+ "(parent_visibility_epoch IS NULL OR parent_visibility_epoch >= 1)",
+ name="ck_entity_visibility_epoch",
+ ),
+ Index(
+ "idx_entity_canonical_identity",
+ "sync_id",
+ "entity_definition_short_name",
+ "native_id",
+ "container_id",
+ ),
+ Index(
+ "idx_entity_canonical_parent",
+ "sync_id",
+ "parent_record_type",
+ "parent_native_id",
+ "parent_container_id",
+ ),
+ Index(
+ "idx_entity_canonical_scope", "sync_id", "entity_definition_short_name", "container_id"
+ ),
+ Index(
+ "idx_entity_pending_revision",
+ "sync_id",
+ "id",
+ postgresql_where=text(
+ "record_revision > 0 AND indexed_revision IS DISTINCT FROM record_revision"
+ ),
+ ),
UniqueConstraint(
"sync_id",
"entity_id",
diff --git a/backend/airweave/models/entity_change.py b/backend/airweave/models/entity_change.py
new file mode 100644
index 0000000000..da8c05450e
--- /dev/null
+++ b/backend/airweave/models/entity_change.py
@@ -0,0 +1,24 @@
+"""Immutable snapshots of observed canonical record revisions."""
+
+from uuid import UUID
+
+from sqlalchemy import BigInteger, ForeignKey, String, UniqueConstraint
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airweave.models._base import OrganizationBase
+
+
+class EntityChange(OrganizationBase):
+ """Per-sync commit-ordered journal; snapshot survives later record updates."""
+
+ __tablename__ = "entity_change"
+
+ sync_id: Mapped[UUID] = mapped_column(ForeignKey("sync.id", ondelete="CASCADE"))
+ entity_record_id: Mapped[UUID] = mapped_column(ForeignKey("entity.id", ondelete="CASCADE"))
+ sequence: Mapped[int] = mapped_column(BigInteger)
+ record_revision: Mapped[int] = mapped_column(BigInteger)
+ kind: Mapped[str] = mapped_column(String)
+ snapshot: Mapped[dict] = mapped_column(JSONB)
+
+ __table_args__ = (UniqueConstraint("sync_id", "sequence", name="uq_entity_change_sequence"),)
diff --git a/backend/airweave/models/owned_provisioning.py b/backend/airweave/models/owned_provisioning.py
new file mode 100644
index 0000000000..3857c1f5ea
--- /dev/null
+++ b/backend/airweave/models/owned_provisioning.py
@@ -0,0 +1,44 @@
+"""One desired Almanac account/source relationship, not another record store."""
+
+from datetime import datetime
+from uuid import UUID
+
+from sqlalchemy import BigInteger, DateTime, ForeignKey, String, UniqueConstraint
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airweave.models._base import OrganizationBase
+
+
+class OwnedProvisioning(OrganizationBase):
+ """Durable idempotency and execution intent for the fixed Almanac client."""
+
+ __tablename__ = "owned_provisioning"
+ client_namespace: Mapped[str] = mapped_column(String(32), nullable=False)
+ account_id: Mapped[UUID] = mapped_column(nullable=False)
+ generation: Mapped[int] = mapped_column(BigInteger, nullable=False)
+ request_hash: Mapped[str] = mapped_column(String(64), nullable=False)
+ request_payload: Mapped[dict] = mapped_column(JSONB, nullable=False)
+ desired_state: Mapped[str] = mapped_column(String(20), nullable=False)
+ observed_generation: Mapped[int] = mapped_column(BigInteger, default=0, nullable=False)
+ source_connection_id: Mapped[UUID | None] = mapped_column(
+ ForeignKey("source_connection.id", ondelete="RESTRICT"), nullable=True
+ )
+ sync_id: Mapped[UUID | None] = mapped_column(
+ ForeignKey("sync.id", ondelete="RESTRICT"), nullable=True
+ )
+ initial_job_id: Mapped[UUID | None] = mapped_column(
+ ForeignKey("sync_job.id", ondelete="SET NULL"), nullable=True
+ )
+ cancellation_job_ids: Mapped[list[str]] = mapped_column(JSONB, default=list, nullable=False)
+ verified_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
+
+ __table_args__ = (
+ UniqueConstraint(
+ "organization_id",
+ "client_namespace",
+ "account_id",
+ name="uq_owned_provisioning_account",
+ ),
+ UniqueConstraint("sync_id", name="uq_owned_provisioning_sync"),
+ )
diff --git a/backend/airweave/models/projection_generation.py b/backend/airweave/models/projection_generation.py
new file mode 100644
index 0000000000..7bb510f27e
--- /dev/null
+++ b/backend/airweave/models/projection_generation.py
@@ -0,0 +1,37 @@
+"""Body-free deletion manifest survives removal of the source and its records."""
+
+from datetime import datetime
+from uuid import UUID
+
+from sqlalchemy import BigInteger, DateTime, Index, String
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airweave.models._base import Base
+
+
+class ProjectionGeneration(Base):
+ """Immutable feed identity; retired generations can never publish again."""
+
+ __tablename__ = "projection_generation"
+ organization_id: Mapped[UUID] = mapped_column(index=True)
+ sync_id: Mapped[UUID] = mapped_column(index=True)
+ collection_id: Mapped[UUID] = mapped_column()
+ record_id: Mapped[UUID] = mapped_column(index=True)
+ revision: Mapped[int] = mapped_column(BigInteger)
+ pipeline_version: Mapped[int] = mapped_column(BigInteger)
+ documents: Mapped[list] = mapped_column(JSONB)
+ extraction_coverage: Mapped[dict | None] = mapped_column(
+ JSONB(none_as_null=True), nullable=True
+ )
+ text_representations: Mapped[list | None] = mapped_column(
+ JSONB(none_as_null=True), nullable=True
+ )
+ retired_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+ next_gc_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True)
+ last_gc_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+ delete_cursor: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ gc_passes: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ gc_attempt: Mapped[UUID | None] = mapped_column(nullable=True)
+ gc_error: Mapped[str | None] = mapped_column(String(200), nullable=True)
+ __table_args__ = (Index("idx_projection_gc_due", "next_gc_at", "id"),)
diff --git a/backend/airweave/models/sync.py b/backend/airweave/models/sync.py
index baa0ce4c49..4948495376 100644
--- a/backend/airweave/models/sync.py
+++ b/backend/airweave/models/sync.py
@@ -5,7 +5,7 @@
from typing import TYPE_CHECKING, Any, Optional
from uuid import UUID
-from sqlalchemy import JSON, DateTime, String, event, text
+from sqlalchemy import JSON, BigInteger, DateTime, String, event, text
from sqlalchemy import Connection as SAConnection
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship
@@ -41,6 +41,18 @@ class Sync(OrganizationBase, UserMixin):
sync_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
sync_config: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
+ provisioning_generation: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ provisioning_ready_generation: Mapped[int] = mapped_column(
+ BigInteger, default=0, server_default="0"
+ )
+
+ observed_change_sequence: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ writer_epoch: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ writer_job_id: Mapped[Optional[UUID]] = mapped_column(nullable=True)
+ writer_attempt_id: Mapped[Optional[UUID]] = mapped_column(nullable=True)
+ writer_attempt_number: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
+ index_pipeline_version: Mapped[int] = mapped_column(BigInteger, default=1, server_default="1")
+
jobs: Mapped[list["SyncJob"]] = relationship(
"SyncJob",
back_populates="sync",
diff --git a/backend/airweave/models/sync_job.py b/backend/airweave/models/sync_job.py
index 2dc87524df..ffbdf312b5 100644
--- a/backend/airweave/models/sync_job.py
+++ b/backend/airweave/models/sync_job.py
@@ -4,7 +4,7 @@
from typing import TYPE_CHECKING, Optional
from uuid import UUID
-from sqlalchemy import JSON, Boolean, DateTime, ForeignKey, Index, Integer, String
+from sqlalchemy import JSON, BigInteger, Boolean, DateTime, ForeignKey, Index, Integer, String
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship
@@ -24,6 +24,7 @@ class SyncJob(OrganizationBase, UserMixin):
sync_id: Mapped[UUID] = mapped_column(
ForeignKey("sync.id", ondelete="CASCADE", name="fk_sync_job_sync_id"), nullable=False
)
+ provisioning_generation: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
status: Mapped[str] = mapped_column(String(50), default=SyncJobStatus.PENDING.value)
started_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
@@ -50,7 +51,6 @@ class SyncJob(OrganizationBase, UserMixin):
"Entity",
back_populates="sync_job",
lazy="noload",
- cascade="all, delete-orphan",
passive_deletes=True,
)
diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py
index 77d8d004c7..4d90d582ed 100644
--- a/backend/airweave/platform/configs/config.py
+++ b/backend/airweave/platform/configs/config.py
@@ -1,9 +1,18 @@
"""Configuration classes for platform components."""
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
from typing import Literal, Optional
-
-from pydantic import Field, field_validator
+from uuid import UUID
+
+from pydantic import (
+ AwareDatetime,
+ BaseModel,
+ ConfigDict,
+ Field,
+ StrictBool,
+ field_validator,
+ model_validator,
+)
from airweave.platform.configs._base import BaseConfig, RequiredTemplateConfig
from airweave.platform.utils.ssrf import validate_host, validate_url
@@ -34,9 +43,9 @@ class ApolloConfig(SourceConfig):
class AttioConfig(SourceConfig):
- """Attio configuration schema."""
+ """Attested workspace boundary for native CRM capture."""
- pass
+ workspace_id: UUID = Field(description="Native Attio workspace UUID; verified using /v2/self")
class BitbucketConfig(SourceConfig):
@@ -175,52 +184,51 @@ class ElasticsearchConfig(SourceConfig):
pass
-class GitHubConfig(SourceConfig):
- """Github configuration schema."""
+class GitHubRepositorySelection(BaseModel):
+ """Stable authorized repository/owner identities plus a refreshable native route."""
- repo_name: str = Field(
- title="Repository Name",
- description="Repository to sync in owner/repo format (e.g., 'airweave-ai/airweave')",
- min_length=3,
- pattern=r"^[a-zA-Z0-9_-]+/[a-zA-Z0-9_.-]+$",
- )
- branch: str = Field(
- default="",
- title="Branch name",
- description=(
- "Specific branch to sync (e.g., 'main', 'development'). "
- "If empty, uses the default branch."
- ),
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ repository_id: int = Field(gt=0, strict=True)
+ owner_id: int = Field(gt=0, strict=True)
+ full_name: str = Field(pattern=r"^[a-zA-Z0-9_-]+/[a-zA-Z0-9_.-]+$", max_length=300)
+ ref: str | None = Field(
+ default=None,
+ min_length=1,
+ max_length=1024,
+ description="Selected branch name; defaults to the repository default branch",
)
- sync_pull_requests: bool = Field(
- default=False,
- title="Sync Pull Requests",
- description=(
- "Sync merged pull requests and their review comments. "
- "Enables searching over PR descriptions, discussions, and code review feedback."
- ),
+
+
+class GitHubConfig(SourceConfig):
+ """Selected repositories captured through one owned account and durable page engine."""
+
+ model_config = ConfigDict(extra="forbid")
+ expected_user_id: int = Field(
+ gt=0,
+ strict=True,
+ description="Expected native GitHub user ID; verified using authenticated /user",
)
+ repositories: tuple[GitHubRepositorySelection, ...] = Field(min_length=1, max_length=500)
+ include_code: bool = True
+ include_conversations: bool = True
- @field_validator("repo_name")
+ @model_validator(mode="before")
@classmethod
- def validate_repo_name(cls, v: str) -> str:
- """Validate repository name is in owner/repo format."""
- if not v or not v.strip():
- raise ValueError("Repository name is required")
- v = v.strip()
- if "/" not in v:
- raise ValueError(
- "Repository must be in 'owner/repo' format (e.g., 'airweave-ai/airweave')"
- )
- parts = v.split("/")
- if len(parts) != 2:
- raise ValueError(
- "Repository must be in 'owner/repo' format (e.g., 'airweave-ai/airweave')"
- )
- owner, repo = parts
- if not owner or not repo:
- raise ValueError("Both owner and repository name must be non-empty")
- return v
+ def require_stable_selection(cls, value):
+ """Old name-only configurations cannot silently authorize a replacement repository."""
+ if isinstance(value, dict) and "repo_name" in value:
+ raise ValueError("Refresh GitHub repository selections with repository and owner IDs")
+ return value
+
+ @model_validator(mode="after")
+ def unique_selection(self) -> "GitHubConfig":
+ """One selected ref per repository; no ambiguous duplicate scope owners."""
+ ids = [selection.repository_id for selection in self.repositories]
+ if len(ids) != len(set(ids)):
+ raise ValueError("GitHub repository selections must be unique")
+ if not self.include_code and not self.include_conversations:
+ raise ValueError("Select GitHub code or conversations to capture")
+ return self
class GitLabConfig(SourceConfig):
@@ -245,6 +253,15 @@ class GitLabConfig(SourceConfig):
class GmailConfig(SourceConfig):
"""Gmail configuration schema."""
+ expected_mailbox: str | None = Field(
+ default=None,
+ pattern=r"^[^\s@]+@[^\s@]+$",
+ description=(
+ "Native mailbox identity attested by the trusted account binding. "
+ "Required for owned capture."
+ ),
+ )
+
after_date: Optional[str] = Field(
None,
title="After Date",
@@ -307,10 +324,56 @@ def validate_date_format(cls, value):
return value.replace("-", "/")
+class CalendarOccurrenceWindow(BaseConfig):
+ """Explicit operator-requested capture window; dates must carry timezones."""
+
+ start: AwareDatetime
+ end: AwareDatetime
+
+ @model_validator(mode="after")
+ def bounded(self):
+ """Keep one explicit capture request bounded and nonempty."""
+ if self.start.microsecond or self.end.microsecond:
+ raise ValueError("Calendar capture bounds require whole-second precision")
+ if not timedelta(0) < self.end - self.start <= timedelta(days=366):
+ raise ValueError("Calendar capture window must be positive and at most366days")
+ return self
+
+
class GoogleCalendarConfig(SourceConfig):
- """Google Calendar configuration schema."""
+ """Rolling expanded capture, or an explicit window through source configuration."""
- pass
+ expected_primary_calendar_id: str | None = Field(
+ default=None,
+ min_length=1,
+ pattern=r"^\S+$",
+ description="Trusted primary calendar ID; required for owned capture.",
+ )
+
+ calendar_ids: tuple[str, ...] | None = Field(default=None, max_length=250)
+
+ @field_validator("calendar_ids")
+ @classmethod
+ def unique_calendars(cls, value: tuple[str, ...] | None) -> tuple[str, ...] | None:
+ """None selects all, empty selects none; only exact native IDs are accepted."""
+ if value is not None and (
+ len(set(value)) != len(value)
+ or any(not item.strip() or item != item.strip() or item == "primary" for item in value)
+ ):
+ raise ValueError("Use unique exact calendar IDs, not blank values or primary alias")
+ return value
+
+ occurrence_past_days: int = Field(default=30, ge=0, le=180)
+ occurrence_future_days: int = Field(default=90, ge=1, le=186)
+ occurrence_window: CalendarOccurrenceWindow | None = None
+
+ def resolved_window(self) -> CalendarOccurrenceWindow:
+ """Freeze relative bounds once for the whole paginated source run."""
+ now = datetime.now(timezone.utc).replace(microsecond=0)
+ return self.occurrence_window or CalendarOccurrenceWindow(
+ start=now - timedelta(days=self.occurrence_past_days),
+ end=now + timedelta(days=self.occurrence_future_days),
+ )
class GoogleDocsConfig(SourceConfig):
@@ -332,6 +395,18 @@ class GoogleDocsConfig(SourceConfig):
class GoogleDriveConfig(SourceConfig):
"""Google Drive configuration schema."""
+ expected_permission_id: str | None = Field(
+ default=None,
+ min_length=1,
+ pattern=r"^\S+$",
+ description="Trusted native Drive permission ID; required for owned canonical capture.",
+ )
+
+ capture_native_sheets: bool = Field(
+ default=False,
+ description="Enable bounded native Sheets capture after verifying account API access.",
+ )
+
include_patterns: list[str] = Field(
default=[],
title="Include Patterns",
@@ -428,9 +503,22 @@ class JiraConfig(SourceConfig):
class LinearConfig(SourceConfig):
- """Linear configuration schema."""
+ """Explicit workspace and team scope for owned Linear originals."""
- pass
+ workspace_id: UUID = Field(description="Expected native Linear organization UUID")
+ team_ids: tuple[UUID, ...] = Field(
+ min_length=1,
+ max_length=100,
+ description="Exact selected native team UUIDs; no implicit workspace-wide default",
+ )
+
+ @field_validator("team_ids")
+ @classmethod
+ def unique_teams(cls, value: tuple[UUID, ...]) -> tuple[UUID, ...]:
+ """Reject ambiguous duplicate configuration rather than silently rewriting it."""
+ if len(value) != len(set(value)):
+ raise ValueError("Linear team IDs must be unique")
+ return value
class MondayConfig(SourceConfig):
@@ -446,9 +534,11 @@ class MySQLConfig(SourceConfig):
class NotionConfig(SourceConfig):
- """Notion configuration schema."""
+ """Verified workspace and connection bot identity for owned Notion capture."""
- pass
+ model_config = ConfigDict(extra="forbid")
+ expected_workspace_id: UUID
+ expected_bot_id: UUID
class OneDriveConfig(SourceConfig):
@@ -490,14 +580,65 @@ class OracleConfig(SourceConfig):
class OutlookCalendarConfig(SourceConfig):
- """Outlook Calendar configuration schema."""
+ """Owned calendar originals with bounded occurrence expansion on a fresh source."""
- pass
+ capture_originals: bool = Field(
+ default=False,
+ strict=True,
+ description="Use owned immutable-event capture; legacy source IDs are not migrated.",
+ )
+ occurrence_past_days: int = Field(default=30, ge=0, le=180)
+ occurrence_future_days: int = Field(default=90, ge=1, le=186)
+ occurrence_window: CalendarOccurrenceWindow | None = None
+
+ @model_validator(mode="after")
+ def require_capture_principal(self):
+ """Owned capture cannot inherit an unattested legacy mailbox identity."""
+ if self.capture_originals and self.expected_principal_id is None:
+ raise ValueError("Original Outlook Calendar capture requires expected_principal_id")
+ return self
+
+ def resolved_window(self) -> CalendarOccurrenceWindow:
+ """Freeze relative expansion bounds once before the durable cycle starts."""
+ now = datetime.now(timezone.utc).replace(microsecond=0)
+ return self.occurrence_window or CalendarOccurrenceWindow(
+ start=now - timedelta(days=self.occurrence_past_days),
+ end=now + timedelta(days=self.occurrence_future_days),
+ )
+
+ expected_principal_id: str | None = Field(
+ default=None,
+ min_length=1,
+ pattern=r"^\S+$",
+ description="Trusted Microsoft Graph /me ID; required for managed connections.",
+ )
class OutlookMailConfig(SourceConfig):
"""Outlook Mail configuration schema."""
+ capture_originals: bool = Field(
+ default=False,
+ strict=True,
+ description=(
+ "Use owned immutable-message capture on a fresh source; legacy IDs are not migrated."
+ ),
+ )
+
+ @model_validator(mode="after")
+ def require_capture_principal(self):
+ """Original capture must not rely on an unattested legacy mailbox."""
+ if self.capture_originals and self.expected_principal_id is None:
+ raise ValueError("Original Outlook capture requires expected_principal_id")
+ return self
+
+ expected_principal_id: str | None = Field(
+ default=None,
+ min_length=1,
+ pattern=r"^\S+$",
+ description="Trusted Microsoft Graph /me ID; required for managed connections.",
+ )
+
after_date: Optional[str] = Field(
None,
title="After Date",
@@ -726,9 +867,23 @@ def validate_host_ssrf(cls, v: str) -> str:
class SlackConfig(SourceConfig):
- """Slack configuration schema."""
+ """Trusted native workspace/user pair for owned Slack capture."""
- pass
+ capture_files: bool = Field(
+ default=False,
+ description="Opt in to message-owned file children, each committed independently; "
+ "leave disabled until reader and migration qualification complete.",
+ )
+
+ expected_team_id: str | None = Field(default=None, min_length=1, pattern=r"^\S+$")
+ expected_user_id: str | None = Field(default=None, min_length=1, pattern=r"^\S+$")
+
+ @model_validator(mode="after")
+ def complete_principal(self):
+ """A workspace alone cannot attest a user's private-channel visibility."""
+ if (self.expected_team_id is None) != (self.expected_user_id is None):
+ raise ValueError("Slack identity requires both team and user IDs")
+ return self
class SQLServerConfig(SourceConfig):
@@ -743,10 +898,29 @@ class SQliteConfig(SourceConfig):
pass
+class StripeCaptureConfig(BaseModel):
+ """Trusted immutable source identity; different accounts/modes need fresh sources."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ expected_account_id: str = Field(pattern=r"^acct_[A-Za-z0-9]+$")
+ livemode: StrictBool
+ api_version: str = Field(pattern=r"^\d{4}-\d{2}-\d{2}(?:\.[a-z]+)?$")
+ connected_account_id: str | None = Field(default=None, pattern=r"^acct_[A-Za-z0-9]+$")
+
+ @model_validator(mode="after")
+ def consistent_connect_context(self):
+ """The selected Connect account is the effective native account, not its platform."""
+ if self.connected_account_id is not None and (
+ self.connected_account_id != self.expected_account_id
+ ):
+ raise ValueError("Stripe Connect context must match expected_account_id")
+ return self
+
+
class StripeConfig(SourceConfig):
- """Stripe configuration schema."""
+ """Legacy entity ingestion, or explicitly bound original capture on a fresh source."""
- pass
+ original_capture: StripeCaptureConfig | None = None
class PipedriveConfig(SourceConfig):
@@ -1183,6 +1357,12 @@ class ComposioConfig(AuthProviderConfig):
description="Account ID for the Composio connection",
)
+ user_id: str | None = Field(
+ default=None,
+ title="User ID",
+ description="Explicit Composio user binding required for managed tool sessions",
+ )
+
class CustomConfig(AuthProviderConfig):
"""Custom Auth Provider configuration schema."""
diff --git a/backend/airweave/platform/cursors/google_calendar.py b/backend/airweave/platform/cursors/google_calendar.py
new file mode 100644
index 0000000000..89f78db7ae
--- /dev/null
+++ b/backend/airweave/platform/cursors/google_calendar.py
@@ -0,0 +1,25 @@
+"""Calendar delta tokens and successfully observed expanded-window coverage."""
+
+from uuid import UUID
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
+
+from airweave.platform.cursors._base import BaseCursor
+
+
+class CalendarWindowCoverage(BaseModel):
+ """A completed observed scan, never a frozen provider snapshot."""
+
+ model_config = ConfigDict(extra="forbid")
+ start: AwareDatetime
+ end: AwareDatetime
+ timezone: str
+ completed_at: AwareDatetime
+ scan_id: UUID
+
+
+class GoogleCalendarCursor(BaseCursor):
+ """Master deltas and expanded rows have independent source semantics."""
+
+ calendar_tokens: dict[str, str] = Field(default_factory=dict)
+ occurrence_coverage: dict[str, CalendarWindowCoverage] = Field(default_factory=dict)
diff --git a/backend/airweave/platform/cursors/google_drive.py b/backend/airweave/platform/cursors/google_drive.py
index bb5ab55f91..bacfe2fae6 100644
--- a/backend/airweave/platform/cursors/google_drive.py
+++ b/backend/airweave/platform/cursors/google_drive.py
@@ -16,6 +16,8 @@ class GoogleDriveCursor(BaseCursor):
Reference: https://developers.google.com/drive/api/guides/manage-changes
"""
+ canonical_page_token: str = ""
+
start_page_token: str = Field(
default="",
description="Drive Changes API page token for tracking incremental changes",
diff --git a/backend/airweave/platform/destinations/vespa/client.py b/backend/airweave/platform/destinations/vespa/client.py
index b279c317a2..7ca700b43e 100644
--- a/backend/airweave/platform/destinations/vespa/client.py
+++ b/backend/airweave/platform/destinations/vespa/client.py
@@ -420,10 +420,18 @@ def _parse_vespa_document_id(full_id: str) -> Tuple[Optional[str], Optional[str]
schema = prefix_parts[2]
return schema, doc_id
+ async def delete_documents(self, doc_ids: List[Tuple[str, str]]) -> None:
+ """Strict idempotent exact-ID deletion, with bounded HTTP concurrency."""
+ if any(schema not in ALL_VESPA_SCHEMAS for schema, _ in doc_ids):
+ raise ValueError("Unknown Vespa schema in deletion manifest")
+ async with httpx.AsyncClient(timeout=settings.VESPA_TIMEOUT) as client:
+ await self._delete_by_doc_ids(doc_ids, http_client=client, strict=True)
+
async def _delete_by_doc_ids(
self,
doc_ids: List[Tuple[str, str]],
http_client: httpx.AsyncClient,
+ *, strict: bool = False,
) -> int:
"""Delete documents by their Vespa document IDs using parallel direct DELETEs.
@@ -432,6 +440,7 @@ async def _delete_by_doc_ids(
Args:
doc_ids: List of (schema, doc_id) tuples
http_client: httpx client for issuing DELETE requests
+ strict: Fail partial deletion and accept already-absent documents as complete.
Returns:
Number of successfully deleted documents
@@ -448,7 +457,7 @@ async def _delete_one(schema: str, doc_id: str) -> bool:
url = f"{base_url}/document/v1/airweave/{schema}/docid/{quote(doc_id, safe='')}"
async with semaphore:
resp = await http_client.delete(url)
- return resp.status_code == 200
+ return resp.status_code in (200, 404) if strict else resp.status_code == 200
start = time.perf_counter()
tasks = [_delete_one(schema, doc_id) for schema, doc_id in doc_ids]
@@ -461,6 +470,8 @@ async def _delete_one(schema: str, doc_id: str) -> bool:
else:
failed += 1
+ if failed > 0 and strict:
+ raise RuntimeError(f"Vespa deletion incomplete: {failed} documents failed")
if failed > 0:
self._logger.warning(
f"[VespaClient] Direct delete: {deleted} ok, {failed} failed "
diff --git a/backend/airweave/platform/destinations/vespa/destination.py b/backend/airweave/platform/destinations/vespa/destination.py
index b92cd14d79..1b88ebc8dd 100644
--- a/backend/airweave/platform/destinations/vespa/destination.py
+++ b/backend/airweave/platform/destinations/vespa/destination.py
@@ -26,6 +26,7 @@
from airweave.platform.destinations.vespa.client import VespaClient
from airweave.platform.destinations.vespa.query_builder import QueryBuilder
from airweave.platform.destinations.vespa.transformer import EntityTransformer
+from airweave.platform.destinations.vespa.types import VespaDocument
from airweave.platform.entities._base import BaseEntity
from airweave.schemas.search import AirweaveTemporalConfig
from airweave.schemas.search_result import AirweaveSearchResult
@@ -140,11 +141,30 @@ async def setup_collection(self, vector_size: Optional[int] = None) -> None:
"""
self.logger.debug("Vespa schema is managed via vespa-deploy, skipping setup_collection")
- async def bulk_insert(self, entities: List[BaseEntity]) -> None:
+ def prepare_documents(self, entities: List[BaseEntity]) -> Dict[str, List[VespaDocument]]:
+ """Strictly transform once so the durable manifest matches the actual feed."""
+ documents = self._transformer.transform_batch(entities)
+ if sum(len(group) for group in documents.values()) != len(entities):
+ raise RuntimeError("Vespa transformation dropped required projection chunks")
+ return documents
+
+ async def feed_prepared(self, documents: Dict[str, List[VespaDocument]]) -> None:
+ """Feed precisely the already-manifested documents, failing partial results."""
+ if self._client is None:
+ raise RuntimeError("Vespa client not initialized")
+ expected = sum(len(group) for group in documents.values())
+ result = await self._client.feed_documents(documents)
+ if result.failed_docs:
+ self._handle_feed_failures(result.failed_docs, expected)
+ if result.success_count != expected:
+ raise RuntimeError("Vespa feed did not acknowledge every prepared document")
+
+ async def bulk_insert(self, entities: List[BaseEntity], *, strict: bool = False) -> None:
"""Transform entities and batch feed to Vespa.
Args:
entities: List of entities to insert
+ strict: Reject any dropped transformation before feeding.
"""
if not entities:
return
@@ -162,6 +182,8 @@ async def bulk_insert(self, entities: List[BaseEntity]) -> None:
# Convert to dict format for client
docs_dict = dict(docs_by_schema.items())
total_docs = sum(len(docs) for docs in docs_dict.values())
+ if strict and total_docs != len(entities):
+ raise RuntimeError("Vespa transformation dropped required projection chunks")
transform_ms = (time.perf_counter() - transform_start) * 1000
self.logger.info(
diff --git a/backend/airweave/platform/destinations/vespa/filter_translator.py b/backend/airweave/platform/destinations/vespa/filter_translator.py
index 80eeac8492..0ca5d78744 100644
--- a/backend/airweave/platform/destinations/vespa/filter_translator.py
+++ b/backend/airweave/platform/destinations/vespa/filter_translator.py
@@ -3,7 +3,7 @@
Pure transformation logic with no I/O dependencies.
"""
-from datetime import datetime
+from datetime import datetime, timezone
from typing import Any, Dict, Optional
from airweave.core.logging import ContextualLogger
@@ -33,8 +33,8 @@
"access.viewers": "access_viewers",
}
-# Fields stored as epoch milliseconds in Vespa
-EPOCH_MS_FIELDS = {"created_at", "updated_at"}
+# Fields stored as epoch seconds in Vespa
+EPOCH_SECONDS_FIELDS = {"created_at", "updated_at"}
class FilterTranslator:
@@ -201,9 +201,9 @@ def _translate_range_condition(self, condition: Dict[str, Any]) -> str:
for op, symbol in [("gt", ">"), ("gte", ">="), ("lt", "<"), ("lte", "<=")]:
if op in range_cond:
value = range_cond[op]
- # Convert ISO datetime strings to epoch milliseconds for date fields
- if key in EPOCH_MS_FIELDS and isinstance(value, str):
- value = self._parse_datetime_to_epoch_ms(value)
+ # Convert ISO datetime strings to epoch seconds for date fields
+ if key in EPOCH_SECONDS_FIELDS and isinstance(value, str):
+ value = self._parse_datetime_to_epoch_seconds(value)
parts.append(f"{key} {symbol} {value}")
return " AND ".join(parts) if parts else ""
@@ -234,13 +234,18 @@ def _escape_yql_value(self, value: str) -> str:
"""Escape special characters for YQL string literals."""
return value.replace("\\", "\\\\").replace('"', '\\"')
- def _parse_datetime_to_epoch_ms(self, value: str) -> int:
- """Parse ISO datetime string to epoch milliseconds."""
+ def _parse_datetime_to_epoch_seconds(self, value: str) -> int | float | str:
+ """Parse ISO datetime string to epoch seconds."""
try:
if value.endswith("Z"):
value = value[:-1] + "+00:00"
dt = datetime.fromisoformat(value)
- return int(dt.timestamp() * 1000)
+ # Existing filters accept zone-less ISO values; interpret them as UTC,
+ # never as the deployment host timezone. New product inputs are aware.
+ if dt.tzinfo is None:
+ dt = dt.replace(tzinfo=timezone.utc)
+ seconds = dt.timestamp()
+ return int(seconds) if seconds.is_integer() else seconds
except (ValueError, AttributeError) as e:
self._logger.warning(f"[FilterTranslator] Failed to parse datetime '{value}': {e}")
return value
diff --git a/backend/airweave/platform/destinations/vespa/transformer.py b/backend/airweave/platform/destinations/vespa/transformer.py
index d118f6d621..835e5f48d9 100644
--- a/backend/airweave/platform/destinations/vespa/transformer.py
+++ b/backend/airweave/platform/destinations/vespa/transformer.py
@@ -323,6 +323,16 @@ def _build_system_metadata(self, entity: BaseEntity) -> Dict[str, Any]:
if value is not None:
meta_fields[field_name] = transform(value) if transform else value
+ # Explicit canonical fields; do not derive them from BaseEntity dates.
+ meta_fields.update(
+ {
+ "canonical_record_type": meta.canonical_record_type,
+ "source_created_us": meta.source_created_us,
+ "source_updated_us": meta.source_updated_us,
+ "source_created_known": meta.source_created_known,
+ "source_updated_known": meta.source_updated_known,
+ }
+ )
return meta_fields
def _add_type_specific_fields(self, fields: Dict[str, Any], entity: BaseEntity) -> None:
diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py
index 7496964bb6..9d177da0db 100644
--- a/backend/airweave/platform/entities/__init__.py
+++ b/backend/airweave/platform/entities/__init__.py
@@ -29,12 +29,7 @@
AsanaTaskEntity,
AsanaWorkspaceEntity,
)
-from .attio import (
- AttioListEntity,
- AttioNoteEntity,
- AttioObjectEntity,
- AttioRecordEntity,
-)
+from .attio import AttioOriginalEntity
from .bitbucket import (
BitbucketCodeFileEntity,
BitbucketDirectoryEntity,
@@ -121,6 +116,7 @@
GitHubPullRequestEntity,
GithubRepoEntity,
GitHubRepositoryEntity,
+ GitHubTextEntity,
)
from .gitlab import (
GitLabCodeFileEntity,
@@ -147,6 +143,7 @@
GoogleDriveDriveEntity,
GoogleDriveFileDeletionEntity,
GoogleDriveFileEntity,
+ GoogleDriveFolderEntity,
)
from .google_slides import (
GoogleSlidesPresentationEntity,
@@ -180,6 +177,7 @@
LinearAttachmentEntity,
LinearCommentEntity,
LinearIssueEntity,
+ LinearLinkedAttachmentEntity,
LinearProjectEntity,
LinearTeamEntity,
LinearUserEntity,
@@ -195,6 +193,7 @@
from .notion import (
NotionDatabaseEntity,
NotionFileEntity,
+ NotionOriginalEntity,
NotionPageEntity,
NotionPropertyEntity,
)
@@ -287,7 +286,7 @@
SlabPostEntity,
SlabTopicEntity,
)
-from .slack import SlackMessageEntity
+from .slack import SlackAttachmentEntity, SlackChannelEntity, SlackMessageEntity
from .slite import SliteNoteEntity
from .stripe import (
StripeBalanceEntity,
@@ -335,6 +334,7 @@
TrelloMemberEntity,
)
from .web import WebFileEntity
+from .wispr import WisprMeetingEntity, WisprNoteEntity
from .word import WordDocumentEntity
from .zendesk import (
ZendeskAttachmentEntity,
@@ -383,12 +383,7 @@
AsanaTaskEntity,
AsanaWorkspaceEntity,
],
- "attio": [
- AttioListEntity,
- AttioNoteEntity,
- AttioObjectEntity,
- AttioRecordEntity,
- ],
+ "attio": [AttioOriginalEntity],
"bitbucket": [
BitbucketCodeFileEntity,
BitbucketDirectoryEntity,
@@ -466,6 +461,7 @@
],
"github": [
GitHubCodeFileEntity,
+ GitHubTextEntity,
GithubContentEntity,
GitHubDirectoryEntity,
GitHubFileDeletionEntity,
@@ -518,6 +514,7 @@
GoogleDocsDocumentEntity,
],
"google_drive": [
+ GoogleDriveFolderEntity,
GoogleDriveDriveEntity,
GoogleDriveFileDeletionEntity,
GoogleDriveFileEntity,
@@ -548,6 +545,7 @@
LinearAttachmentEntity,
LinearCommentEntity,
LinearIssueEntity,
+ LinearLinkedAttachmentEntity,
LinearProjectEntity,
LinearTeamEntity,
LinearUserEntity,
@@ -563,6 +561,7 @@
"notion": [
NotionDatabaseEntity,
NotionFileEntity,
+ NotionOriginalEntity,
NotionPageEntity,
NotionPropertyEntity,
],
@@ -651,6 +650,8 @@
SlabTopicEntity,
],
"slack": [
+ SlackAttachmentEntity,
+ SlackChannelEntity,
SlackMessageEntity,
],
"stripe": [
@@ -704,6 +705,7 @@
"web": [
WebFileEntity,
],
+ "wispr": [WisprMeetingEntity, WisprNoteEntity],
"word": [
WordDocumentEntity,
],
diff --git a/backend/airweave/platform/entities/_base.py b/backend/airweave/platform/entities/_base.py
index 8b3fba968f..95b974db47 100644
--- a/backend/airweave/platform/entities/_base.py
+++ b/backend/airweave/platform/entities/_base.py
@@ -62,6 +62,13 @@ class AirweaveSystemMetadata(BaseModel):
None, description="ID of the sync job this entity belongs to."
)
+ # Canonical projector-owned metadata. Provider entity fallbacks do not set these.
+ canonical_record_type: Optional[str] = None
+ source_created_us: Optional[int] = None
+ source_updated_us: Optional[int] = None
+ source_created_known: Optional[int] = None
+ source_updated_known: Optional[int] = None
+
# Set during hash computation
hash: Optional[str] = Field(None, description="Hash of the content used for change detection.")
diff --git a/backend/airweave/platform/entities/attio.py b/backend/airweave/platform/entities/attio.py
index 2924d96330..21ec84e087 100644
--- a/backend/airweave/platform/entities/attio.py
+++ b/backend/airweave/platform/entities/attio.py
@@ -1,11 +1,7 @@
-"""Attio entity schemas.
-
-Attio is a CRM platform that organizes data into Objects (Companies, People, Deals)
-and Lists (custom collections). Each object/list contains Records with custom attributes.
-"""
+"""Offline search schema for retained Attio originals."""
from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Literal
from pydantic import computed_field
@@ -13,164 +9,25 @@
from airweave.platform.entities._base import BaseEntity
-class AttioObjectEntity(BaseEntity):
- """Schema for Attio Object (e.g., Companies, People, Deals)."""
-
- object_id: str = AirweaveField(..., description="Attio object ID", is_entity_id=True)
- name: str = AirweaveField(
- ..., description="Display name of the object", is_name=True, embeddable=True
- )
- created_at: Optional[datetime] = AirweaveField(
- None, description="When the object was created", is_created_at=True
- )
-
- singular_noun: str = AirweaveField(
- ..., description="Singular name of the object (e.g., 'Company')", embeddable=True
- )
- plural_noun: str = AirweaveField(
- ..., description="Plural name of the object (e.g., 'Companies')", embeddable=True
- )
- api_slug: str = AirweaveField(..., description="API slug for the object", embeddable=False)
- icon: Optional[str] = AirweaveField(
- None, description="Icon representing this object", embeddable=False
- )
-
- @computed_field(return_type=str)
- def web_url(self) -> str:
- """Link to the object definition inside Attio."""
- return f"https://app.attio.com/objects/{self.api_slug}"
-
-
-class AttioListEntity(BaseEntity):
- """Schema for Attio List."""
+class AttioOriginalEntity(BaseEntity):
+ """Search text derived from one retained Attio original, without provider reads."""
- list_id: str = AirweaveField(..., description="Attio list ID", is_entity_id=True)
- name: str = AirweaveField(..., description="List name", is_name=True, embeddable=True)
- created_at: Optional[datetime] = AirweaveField(
- None, description="When the list was created", is_created_at=True
+ native_id: str = AirweaveField(..., description="Native Attio UUID", is_entity_id=True)
+ original_kind: Literal["object", "record", "list", "entry", "note"] = AirweaveField(
+ ..., description="Native record kind"
)
-
- workspace_id: str = AirweaveField(
- ..., description="ID of the workspace this list belongs to", embeddable=False
+ title: str = AirweaveField(..., description="Original label", is_name=True, embeddable=True)
+ text: str = AirweaveField(..., description="Retained original content", embeddable=True)
+ content_coverage: str = AirweaveField(..., description="Included and omitted content")
+ native_url: str | None = AirweaveField(None, description="Verified retained native URL")
+ created_at: datetime | None = AirweaveField(
+ None, description="Native creation time", is_created_at=True
)
- parent_object: Optional[str] = AirweaveField(
- None, description="Parent object type if applicable", embeddable=True
+ updated_at: datetime | None = AirweaveField(
+ None, description="Native modification time", is_updated_at=True
)
@computed_field(return_type=str)
def web_url(self) -> str:
- """Link to the list inside Attio."""
- return f"https://app.attio.com/lists/{self.list_id}"
-
-
-class AttioRecordEntity(BaseEntity):
- """Schema for Attio Record."""
-
- record_id: str = AirweaveField(..., description="Attio record ID", is_entity_id=True)
- name: str = AirweaveField(..., description="Record display name", is_name=True, embeddable=True)
- created_at: Optional[datetime] = AirweaveField(
- None, description="When the record was created", is_created_at=True
- )
- updated_at: Optional[datetime] = AirweaveField(
- None, description="When the record was last updated", is_updated_at=True
- )
-
- object_id: Optional[str] = AirweaveField(
- None, description="ID/slug of the object this record belongs to", embeddable=False
- )
- list_id: Optional[str] = AirweaveField(
- None, description="ID of the list this record belongs to", embeddable=False
- )
- parent_object_name: Optional[str] = AirweaveField(
- None, description="Name of the parent object/list", embeddable=True
- )
-
- description: Optional[str] = AirweaveField(
- None, description="Description of the record", embeddable=True
- )
- email_addresses: List[Dict[str, Any]] = AirweaveField(
- default_factory=list,
- description="Email addresses associated with this record",
- embeddable=True,
- )
- phone_numbers: List[Dict[str, Any]] = AirweaveField(
- default_factory=list,
- description="Phone numbers associated with this record",
- embeddable=True,
- )
- domains: List[str] = AirweaveField(
- default_factory=list, description="Domain names (for company records)", embeddable=True
- )
- categories: List[str] = AirweaveField(
- default_factory=list, description="Categories/tags for this record", embeddable=True
- )
- attributes: Dict[str, Any] = AirweaveField(
- default_factory=dict,
- description="Custom attributes and their values",
- embeddable=True,
- )
- permalink_url: Optional[str] = AirweaveField(
- None,
- description="URL to view this record in Attio (if provided by API)",
- embeddable=False,
- unhashable=True,
- )
-
- @computed_field(return_type=Optional[str])
- def web_url(self) -> Optional[str]:
- """Best-effort link back to the record inside Attio."""
- if self.permalink_url:
- return self.permalink_url
- if self.object_id:
- return f"https://app.attio.com/objects/{self.object_id}/{self.record_id}"
- if self.list_id:
- return f"https://app.attio.com/lists/{self.list_id}/{self.record_id}"
- return None
-
-
-class AttioNoteEntity(BaseEntity):
- """Schema for Attio Note."""
-
- note_id: str = AirweaveField(..., description="Attio note ID", is_entity_id=True)
- name: str = AirweaveField(
- ..., description="Note title or preview", is_name=True, embeddable=True
- )
- created_at: Optional[datetime] = AirweaveField(
- None, description="When the note was created", is_created_at=True
- )
- updated_at: Optional[datetime] = AirweaveField(
- None, description="When the note was last updated", is_updated_at=True
- )
-
- parent_record_id: str = AirweaveField(
- ..., description="ID of the record this note is attached to", embeddable=False
- )
- parent_object: Optional[str] = AirweaveField(
- None, description="Type of parent object", embeddable=False
- )
- title: Optional[str] = AirweaveField(None, description="Title of the note", embeddable=True)
- content: str = AirweaveField(..., description="Content of the note", embeddable=True)
- format: Optional[str] = AirweaveField(
- None, description="Format of the note (plaintext, markdown, etc.)", embeddable=False
- )
- author: Optional[Dict[str, Any]] = AirweaveField(
- None, description="User who created this note", embeddable=True
- )
- permalink_url: Optional[str] = AirweaveField(
- None,
- description="URL to view this note in Attio (if provided by API)",
- embeddable=False,
- unhashable=True,
- )
-
- @computed_field(return_type=Optional[str])
- def web_url(self) -> Optional[str]:
- """Best-effort link back to the note inside Attio."""
- if self.permalink_url:
- return self.permalink_url
- return f"https://app.attio.com/notes/{self.note_id}"
-
-
-# Note: AttioCommentEntity was removed because the Attio API does not provide
-# a way to fetch comments for notes through their public REST API.
-# Comments are visible in the Attio UI but not accessible via /v2/threads or any other endpoint.
+ """Native navigation requires a verified URL; do not synthesize one."""
+ return self.native_url or ""
diff --git a/backend/airweave/platform/entities/github.py b/backend/airweave/platform/entities/github.py
index 051ad2cf76..7733a3b39b 100644
--- a/backend/airweave/platform/entities/github.py
+++ b/backend/airweave/platform/entities/github.py
@@ -558,3 +558,23 @@ class GitHubFileDeletionEntity(DeletionEntity):
def web_url(self) -> str:
"""Fallback to repository page since deleted files have no direct URLs."""
return f"https://github.com/{self.repo_owner}/{self.repo_name}"
+
+
+class GitHubTextEntity(BaseEntity):
+ """Searchable native GitHub text with its exact resource kind and content coverage."""
+
+ original_id: str = AirweaveField(
+ ..., description="Canonical original identity", is_entity_id=True
+ )
+ title: str = AirweaveField(
+ ..., description="Native title or resource label", is_name=True, embeddable=True
+ )
+ resource_kind: str = AirweaveField(
+ ..., description="Native GitHub resource kind", embeddable=True
+ )
+ repository: str = AirweaveField(..., description="Repository full name", embeddable=True)
+ text: str = AirweaveField(..., description="Retained native text", embeddable=True)
+ state: str | None = AirweaveField(None, description="Native state", embeddable=True)
+ content_coverage: str = AirweaveField(
+ ..., description="Capture content coverage", embeddable=False
+ )
diff --git a/backend/airweave/platform/entities/google_drive.py b/backend/airweave/platform/entities/google_drive.py
index c4ff91a50b..ac6e7e5468 100644
--- a/backend/airweave/platform/entities/google_drive.py
+++ b/backend/airweave/platform/entities/google_drive.py
@@ -23,7 +23,7 @@
from airweave.platform.entities._base import BaseEntity, Breadcrumb, DeletionEntity, FileEntity
from airweave.platform.entities.utils import _determine_file_type_from_mime
-_GOOGLE_EXPORT_MAP: Dict[str, tuple[str, str]] = {
+GOOGLE_EXPORT_FORMATS: Dict[str, tuple[str, str]] = {
"application/vnd.google-apps.document": (
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
".docx",
@@ -190,7 +190,7 @@ def from_api(
file_id = data["id"]
if mime_type.startswith("application/vnd.google-apps."):
- export_mime_type, file_extension = _GOOGLE_EXPORT_MAP.get(
+ export_mime_type, file_extension = GOOGLE_EXPORT_FORMATS.get(
mime_type, ("application/pdf", ".pdf")
)
download_url = (
@@ -303,3 +303,16 @@ def web_url(self) -> str:
if self.drive_id:
return f"https://drive.google.com/drive/folders/{self.drive_id}"
return "https://drive.google.com/drive/my-drive"
+
+
+class GoogleDriveFolderEntity(BaseEntity):
+ """Searchable folder metadata; folders never require a fabricated file body."""
+
+ folder_id: str = AirweaveField(..., description="Native folder ID", is_entity_id=True)
+ title: str = AirweaveField(..., description="Folder name", is_name=True, embeddable=True)
+ description: str = AirweaveField(..., description="Folder description", embeddable=True)
+
+ @computed_field(return_type=str)
+ def web_url(self) -> str:
+ """Direct folder locator."""
+ return f"https://drive.google.com/drive/folders/{self.folder_id}"
diff --git a/backend/airweave/platform/entities/linear.py b/backend/airweave/platform/entities/linear.py
index 3ab98dce2a..f344fad3fd 100644
--- a/backend/airweave/platform/entities/linear.py
+++ b/backend/airweave/platform/entities/linear.py
@@ -555,3 +555,26 @@ class LinearAttachmentEntity(FileEntity):
def web_url(self) -> str:
"""Browser URL for the attachment (falls back to download URL)."""
return self.web_url_value or self.url
+
+
+class LinearLinkedAttachmentEntity(BaseEntity):
+ """Searchable link metadata; never represents retained file content."""
+
+ attachment_id: str = AirweaveField(
+ ..., description="Native attachment identity", is_entity_id=True
+ )
+ title: str = AirweaveField(
+ ..., description="Native attachment title", is_name=True, embeddable=True
+ )
+ issue_identifier: str = AirweaveField(
+ ..., description="Parent issue identifier", embeddable=True
+ )
+ target_url: str = AirweaveField(..., description="Original linked resource", embeddable=False)
+ content_coverage: str = AirweaveField(
+ ..., description="Available content coverage", embeddable=True
+ )
+
+ @computed_field(return_type=str)
+ def web_url(self) -> str:
+ """Return the original link without pretending its content was fetched."""
+ return self.target_url
diff --git a/backend/airweave/platform/entities/notion.py b/backend/airweave/platform/entities/notion.py
index 752b116269..d8f6c99eae 100644
--- a/backend/airweave/platform/entities/notion.py
+++ b/backend/airweave/platform/entities/notion.py
@@ -556,3 +556,29 @@ def web_url(self) -> str:
if self.web_url_value:
return self.web_url_value
return self.url
+
+
+class NotionOriginalEntity(BaseEntity):
+ """Searchable metadata or block-local text from one retained Notion original."""
+
+ native_id: str = AirweaveField(
+ ..., description="Native Notion object or property ID.", is_entity_id=True
+ )
+ original_kind: str = AirweaveField(..., description="Native Notion object kind.")
+ title: str = AirweaveField(
+ ..., description="Retained title or block label.", is_name=True, embeddable=True
+ )
+ text: str = AirweaveField(
+ ..., description="Retained metadata or this block's text only.", embeddable=True
+ )
+ native_url: Optional[str] = AirweaveField(
+ None, description="URL supplied by Notion, when retained."
+ )
+ content_coverage: str = AirweaveField(
+ ..., description="What this projection includes and omits."
+ )
+
+ @computed_field(return_type=str)
+ def web_url(self) -> str:
+ """Preserve native navigation without inventing block or child URLs."""
+ return self.native_url or ""
diff --git a/backend/airweave/platform/entities/outlook_calendar.py b/backend/airweave/platform/entities/outlook_calendar.py
index 7167c050de..ebcd1df46a 100644
--- a/backend/airweave/platform/entities/outlook_calendar.py
+++ b/backend/airweave/platform/entities/outlook_calendar.py
@@ -454,3 +454,69 @@ def web_url(self) -> str:
if self.event_web_url:
return self.event_web_url
return f"https://outlook.office.com/calendar/item/{self.event_id}"
+
+
+class OutlookCalendarMetadataEntity(BaseEntity):
+ """Owned calendar metadata without invented permission defaults."""
+
+ native_id: str = AirweaveField(..., description="Native calendar ID.", is_entity_id=True)
+ title: str = AirweaveField(..., description="Calendar name.", is_name=True, embeddable=True)
+ owner: Optional[str] = AirweaveField(None, description="Calendar owner.", embeddable=True)
+ can_edit: Optional[bool] = AirweaveField(None, description="Native edit permission.")
+ can_view_private_items: Optional[bool] = AirweaveField(
+ None, description="Native private-item permission."
+ )
+ is_default_calendar: Optional[bool] = AirweaveField(None, description="Native default status.")
+
+
+class OutlookCalendarSearchEventEntity(BaseEntity):
+ """Searchable native facts; wall times never gain an inferred UTC offset."""
+
+ native_id: str = AirweaveField(..., description="Native event ID.", is_entity_id=True)
+ title: str = AirweaveField(..., description="Event subject.", is_name=True, embeddable=True)
+ calendar_id: str = AirweaveField(..., description="Native owning calendar ID.")
+ event_type: str = AirweaveField(..., description="Native event kind.", embeddable=True)
+ start: str = AirweaveField(..., description="Native start wall time and zone.", embeddable=True)
+ end: str = AirweaveField(..., description="Native end wall time and zone.", embeddable=True)
+ is_all_day: bool = AirweaveField(..., description="Native all-day status.", embeddable=True)
+ is_cancelled: bool = AirweaveField(
+ ..., description="Native cancellation status.", embeddable=True
+ )
+ organizer: Optional[str] = AirweaveField(None, description="Organizer.", embeddable=True)
+ attendees: List[str] = AirweaveField(
+ default_factory=list,
+ description="Participants with native response status.",
+ embeddable=True,
+ )
+ recurrence: Optional[str] = AirweaveField(
+ None, description="Native recurrence pattern and range.", embeddable=True
+ )
+ series_master_id: Optional[str] = AirweaveField(None, description="Native recurring-series ID.")
+ original_start: Optional[str] = AirweaveField(
+ None, description="Native original occurrence start.", embeddable=True
+ )
+ original_start_timezone: Optional[str] = AirweaveField(
+ None, description="Original native start zone.", embeddable=True
+ )
+ original_end_timezone: Optional[str] = AirweaveField(
+ None, description="Original native end zone.", embeddable=True
+ )
+ locations: List[str] = AirweaveField(
+ default_factory=list, description="Native location names.", embeddable=True
+ )
+ categories: List[str] = AirweaveField(
+ default_factory=list, description="Native categories.", embeddable=True
+ )
+ show_as: Optional[str] = AirweaveField(
+ None, description="Native availability.", embeddable=True
+ )
+ web_link: Optional[str] = AirweaveField(None, description="Native Outlook event URL.")
+
+ @computed_field(return_type=Optional[str])
+ def web_url(self) -> Optional[str]:
+ """Expose only the retained native link to shared search result serialization."""
+ return self.web_link
+
+
+class OutlookCalendarHtmlEventEntity(OutlookCalendarSearchEventEntity, FileEntity):
+ """Same event facts plus a disposable HTML input for the shared converter."""
diff --git a/backend/airweave/platform/entities/outlook_mail.py b/backend/airweave/platform/entities/outlook_mail.py
index bfea72230a..7e7c71e967 100644
--- a/backend/airweave/platform/entities/outlook_mail.py
+++ b/backend/airweave/platform/entities/outlook_mail.py
@@ -94,8 +94,8 @@ class OutlookMessageEntity(EmailEntity):
description="Message ID from Microsoft Graph.",
is_entity_id=True,
)
- folder_name: str = AirweaveField(
- ..., description="Name of the folder containing this message", embeddable=True
+ folder_name: Optional[str] = AirweaveField(
+ None, description="Name of the folder containing this message", embeddable=True
)
subject: str = AirweaveField(
...,
@@ -106,6 +106,17 @@ class OutlookMessageEntity(EmailEntity):
sender: Optional[str] = AirweaveField(
None, description="Email address of the sender", embeddable=True
)
+ folder_id: Optional[str] = AirweaveField(None, description="Native current folder ID")
+ from_address: Optional[str] = AirweaveField(
+ None, description="Author address, distinct from the sending agent", embeddable=True
+ )
+ conversation_id: Optional[str] = AirweaveField(None, description="Native conversation ID")
+ bcc_recipients: List[str] = AirweaveField(
+ default_factory=list, description="BCC recipients", embeddable=True
+ )
+ reply_to: List[str] = AirweaveField(
+ default_factory=list, description="Reply-to recipients", embeddable=True
+ )
to_recipients: List[str] = AirweaveField(
default_factory=list, description="Recipients of the message", embeddable=True
)
@@ -232,7 +243,9 @@ class OutlookAttachmentEntity(FileEntity):
is_name=True,
)
message_id: str = AirweaveField(..., description="ID of the message this attachment belongs to")
- attachment_id: str = AirweaveField(..., description="Outlook's attachment ID")
+ attachment_id: str = AirweaveField(
+ ..., description="Native attachment ID or MIME part path; metadata identifies the scheme"
+ )
content_type: Optional[str] = AirweaveField(None, description="Content type of the attachment")
is_inline: bool = AirweaveField(False, description="Whether this is an inline attachment")
content_id: Optional[str] = AirweaveField(None, description="Content ID for inline attachments")
diff --git a/backend/airweave/platform/entities/slack.py b/backend/airweave/platform/entities/slack.py
index e79b9c2a65..7d2f4617d6 100644
--- a/backend/airweave/platform/entities/slack.py
+++ b/backend/airweave/platform/entities/slack.py
@@ -8,7 +8,7 @@
from pydantic import computed_field
from airweave.platform.entities._airweave_field import AirweaveField
-from airweave.platform.entities._base import BaseEntity, Breadcrumb
+from airweave.platform.entities._base import BaseEntity, Breadcrumb, FileEntity
def _parse_slack_ts(ts_str: Any) -> Optional[datetime]:
@@ -132,3 +132,21 @@ def from_api(
def web_url(self) -> str:
"""Permalink for the Slack message."""
return self.web_url_value or self.url or ""
+
+
+class SlackChannelEntity(BaseEntity):
+ """Captured conversation metadata, without inferring visibility from its name."""
+
+ channel_id: str = AirweaveField(..., description="Native conversation ID", is_entity_id=True)
+ title: str = AirweaveField(..., description="Conversation name", is_name=True, embeddable=True)
+ purpose: str = AirweaveField(..., description="Conversation purpose", embeddable=True)
+ topic: str = AirweaveField(..., description="Conversation topic", embeddable=True)
+
+
+class SlackAttachmentEntity(FileEntity):
+ """Retained original file belonging to one captured Slack message."""
+
+ attachment_key: str = AirweaveField(
+ ..., description="Message-scoped native file identity", is_entity_id=True
+ )
+ filename: str = AirweaveField(..., description="Native filename or file ID", is_name=True)
diff --git a/backend/airweave/platform/entities/stripe.py b/backend/airweave/platform/entities/stripe.py
index 6aece31740..42d7acdbbb 100644
--- a/backend/airweave/platform/entities/stripe.py
+++ b/backend/airweave/platform/entities/stripe.py
@@ -163,8 +163,8 @@ class StripeChargeEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the charge was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="Last activity timestamp for the charge.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="Last activity timestamp for the charge.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -180,12 +180,14 @@ class StripeChargeEntity(BaseEntity):
currency: Optional[str] = AirweaveField(
None, description="Three-letter ISO currency code", embeddable=True
)
- captured: bool = AirweaveField(
- False, description="Whether the charge was captured", embeddable=True
+ captured: Optional[bool] = AirweaveField(
+ None, description="Whether the charge was captured", embeddable=True
)
- paid: bool = AirweaveField(False, description="Whether the charge was paid", embeddable=True)
- refunded: bool = AirweaveField(
- False, description="Whether the charge was refunded", embeddable=True
+ paid: Optional[bool] = AirweaveField(
+ None, description="Whether the charge was paid", embeddable=True
+ )
+ refunded: Optional[bool] = AirweaveField(
+ None, description="Whether the charge was refunded", embeddable=True
)
description: Optional[str] = AirweaveField(
None, description="Arbitrary description of the charge", embeddable=True
@@ -270,8 +272,8 @@ class StripeCustomerEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the customer was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="Timestamp of the latest update.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="Timestamp of the latest update.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -300,8 +302,8 @@ class StripeCustomerEntity(BaseEntity):
description="ID of the default payment source (e.g. card) attached to this customer",
embeddable=False,
)
- delinquent: bool = AirweaveField(
- False, description="Whether the customer has any unpaid/overdue invoices", embeddable=True
+ delinquent: Optional[bool] = AirweaveField(
+ None, description="Whether the customer has any unpaid/overdue invoices", embeddable=True
)
invoice_prefix: Optional[str] = AirweaveField(
None, description="Prefix for the customer's invoices", embeddable=False
@@ -383,8 +385,8 @@ class StripeInvoiceEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the invoice was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="When the invoice was last updated.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="When the invoice was last updated.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -417,8 +419,8 @@ class StripeInvoiceEntity(BaseEntity):
due_date: Optional[datetime] = AirweaveField(
None, description="Date on which payment is due (if applicable)", embeddable=True
)
- paid: bool = AirweaveField(
- False, description="Whether the invoice has been fully paid", embeddable=True
+ paid: Optional[bool] = AirweaveField(
+ None, description="Whether the invoice has been fully paid", embeddable=True
)
currency: Optional[str] = AirweaveField(
None, description="Three-letter ISO currency code (e.g. 'usd')", embeddable=True
@@ -497,8 +499,8 @@ class StripePaymentIntentEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the payment intent was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="Last update timestamp for the payment intent.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="Last update timestamp for the payment intent.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -689,8 +691,8 @@ class StripePayoutEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the payout was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="Last update timestamp for the payout.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="Last update timestamp for the payout.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -852,8 +854,8 @@ class StripeSubscriptionEntity(BaseEntity):
created_time: datetime = AirweaveField(
..., description="When the subscription was created.", is_created_at=True
)
- updated_time: datetime = AirweaveField(
- ..., description="Last update timestamp for the subscription.", is_updated_at=True
+ updated_time: Optional[datetime] = AirweaveField(
+ None, description="Last update timestamp for the subscription.", is_updated_at=True
)
web_url_value: Optional[str] = AirweaveField(
None,
@@ -881,8 +883,8 @@ class StripeSubscriptionEntity(BaseEntity):
description="End of the current billing period for this subscription",
embeddable=True,
)
- cancel_at_period_end: bool = AirweaveField(
- False,
+ cancel_at_period_end: Optional[bool] = AirweaveField(
+ None,
description="Whether the subscription will cancel at the end of the current period",
embeddable=True,
)
diff --git a/backend/airweave/platform/entities/wispr.py b/backend/airweave/platform/entities/wispr.py
new file mode 100644
index 0000000000..f2ab126b9c
--- /dev/null
+++ b/backend/airweave/platform/entities/wispr.py
@@ -0,0 +1,48 @@
+"""Search projection of Wispr's captured meeting representations."""
+
+from datetime import datetime
+
+from pydantic import computed_field
+
+from airweave.platform.entities._airweave_field import AirweaveField
+from airweave.platform.entities._base import BaseEntity
+
+
+class WisprMeetingEntity(BaseEntity):
+ """Selected meeting fields; original response ranges stay in canonical storage."""
+
+ meeting_id: str = AirweaveField(..., description="Native meeting ID", is_entity_id=True)
+ title: str = AirweaveField(..., description="Meeting title", is_name=True, embeddable=True)
+ notes: str = AirweaveField(..., description="Captured markdown notes", embeddable=True)
+ summary: str = AirweaveField(..., description="Provider summary", embeddable=True)
+ transcript: str = AirweaveField(..., description="Captured transcript ranges", embeddable=True)
+ starts_at: datetime | None = AirweaveField(
+ None, description="Meeting start", is_created_at=True
+ )
+ modified_at: datetime | None = AirweaveField(
+ None, description="Provider modification", is_updated_at=True
+ )
+ share_link: str | None = AirweaveField(
+ None, description="Provider notes page", embeddable=False
+ )
+
+ @computed_field(return_type=str)
+ def web_url(self) -> str:
+ """Use only the provider-supplied sharing link."""
+ return self.share_link or ""
+
+
+class WisprNoteEntity(BaseEntity):
+ """Search projection of a retained scratchpad note, independent of meetings."""
+
+ note_id: str = AirweaveField(..., description="Native scratchpad note ID", is_entity_id=True)
+ title: str = AirweaveField(..., description="Note title", is_name=True, embeddable=True)
+ content: str = AirweaveField(..., description="Captured normalized note text", embeddable=True)
+ modified_at: datetime | None = AirweaveField(
+ None, description="Provider modification", is_updated_at=True
+ )
+
+ @computed_field(return_type=str)
+ def web_url(self) -> str:
+ """The scratchpad response supplies no native URL."""
+ return ""
diff --git a/backend/airweave/platform/http_client/airweave_client.py b/backend/airweave/platform/http_client/airweave_client.py
index 820870cd64..15fbde9c5f 100644
--- a/backend/airweave/platform/http_client/airweave_client.py
+++ b/backend/airweave/platform/http_client/airweave_client.py
@@ -154,19 +154,10 @@ async def _log_error_response(self, method: str, url: str, response: httpx.Respo
if not self._logger:
return
- try:
- # Try to read response body - may already be consumed
- response_body = response.text
- except Exception:
- response_body = " "
-
+ # Error bodies and signed URL queries may carry account data or credentials.
self._logger.debug(
- f"[AirweaveHttpClient] HTTP {response.status_code} error\n"
- f" Source: {self._source_short_name}\n"
- f" Method: {method}\n"
- f" URL: {url}\n"
- f" Response Headers: {dict(response.headers)}\n"
- f" Response Body: {response_body}"
+ "[AirweaveHttpClient] HTTP %s source=%s method=%s host=%s",
+ response.status_code, self._source_short_name, method, httpx.URL(url).host,
)
# Mimic httpx.AsyncClient methods
diff --git a/backend/airweave/platform/http_client/bounded_response.py b/backend/airweave/platform/http_client/bounded_response.py
new file mode 100644
index 0000000000..c789595dfb
--- /dev/null
+++ b/backend/airweave/platform/http_client/bounded_response.py
@@ -0,0 +1,22 @@
+"""Bound identity-encoded response bodies while owning iterator closure."""
+
+from contextlib import aclosing
+
+import httpx
+
+from airweave.domains.storage import FileSkippedException
+
+
+async def bounded_response_bytes(response: httpx.Response, maximum: int, *, label: str) -> bytes:
+ """Caller requests identity encoding; the raw bound is then a decoded byte bound too."""
+ if response.is_stream_consumed:
+ if len(response.content) > maximum:
+ raise FileSkippedException(f"{label} exceeds size limit", label)
+ return response.content
+ body = bytearray()
+ async with aclosing(response.stream.__aiter__()) as chunks:
+ async for chunk in chunks:
+ if len(body) + len(chunk) > maximum:
+ raise FileSkippedException(f"{label} exceeds size limit", label)
+ body.extend(chunk)
+ return bytes(body)
diff --git a/backend/airweave/platform/http_client/composio_transport.py b/backend/airweave/platform/http_client/composio_transport.py
new file mode 100644
index 0000000000..77e069da93
--- /dev/null
+++ b/backend/airweave/platform/http_client/composio_transport.py
@@ -0,0 +1,262 @@
+"""Account-bound Composio HTTPX transport; credentials remain with Composio."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator, Iterable
+from contextlib import aclosing
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl, JsonValue
+
+from airweave.domains.auth_provider.exceptions import (
+ AuthProviderAuthError,
+ AuthProviderRateLimitError,
+ AuthProviderServerError,
+)
+from airweave.platform.utils.ssrf import validate_url
+
+
+class BinaryData(BaseModel):
+ """Temporary download supplied by Composio, never a provider credential."""
+
+ model_config = ConfigDict(extra="ignore")
+ url: HttpUrl = Field(repr=False)
+ content_type: str
+ size: int = Field(ge=0)
+ expires_at: str
+
+
+class ProxyResponse(BaseModel):
+ """Documented proxy envelope (binary data takes precedence over JSON data)."""
+
+ model_config = ConfigDict(extra="ignore")
+ status: int = Field(ge=100, le=599)
+ data: JsonValue = None
+ headers: dict[str, str] = Field(default_factory=dict)
+ binary_data: BinaryData | None = None
+
+
+class ComposioProxyError(httpx.TransportError):
+ """Invalid proxy protocol or disallowed request, without sensitive payloads."""
+
+
+class DownloadStream(httpx.AsyncByteStream):
+ """Own the temporary download response until its consumer closes it."""
+
+ def __init__(self, response: httpx.Response, max_bytes: int):
+ """Retain the response and enforce an independent byte ceiling."""
+ self.response = response
+ self.max_bytes = max_bytes
+
+ async def __aiter__(self) -> AsyncIterator[bytes]:
+ """Yield bytes while enforcing the configured download bound."""
+ count = 0
+ async with aclosing(self.response.aiter_bytes()) as chunks:
+ async for chunk in chunks:
+ count += len(chunk)
+ if count > self.max_bytes:
+ raise ComposioProxyError("Proxy download exceeds the file size limit")
+ yield chunk
+
+ async def aclose(self) -> None:
+ """Close the underlying download connection."""
+ await self.response.aclose()
+
+
+class ComposioTransport(httpx.AsyncBaseTransport):
+ """Route only this connection's approved hosts through the fixed proxy API."""
+
+ MAX_BINARY_BYTES = 200 * 1024 * 1024
+
+ def __init__(
+ self,
+ *,
+ api_key: str,
+ connected_account_id: str,
+ allowed_hosts: Iterable[str],
+ client: httpx.AsyncClient | None = None,
+ download_transport: httpx.AsyncBaseTransport | None = None,
+ ) -> None:
+ """Bind one account; injected clients stay owned by their caller."""
+ self._owns_client = client is None
+ self._owns_download_transport = download_transport is None
+ self._client = client or httpx.AsyncClient(timeout=180.0, follow_redirects=False)
+ self._download_transport = download_transport or httpx.AsyncHTTPTransport()
+ self._api_key = api_key
+ self._account_id = connected_account_id
+ self._hosts = frozenset(allowed_hosts)
+ if not self._account_id or not self._hosts:
+ raise ValueError("A bound account and provider hosts are required")
+
+ async def aclose(self) -> None:
+ """Close only clients constructed by this transport."""
+ if self._owns_client:
+ await self._client.aclose()
+ if self._owns_download_transport:
+ await self._download_transport.aclose()
+
+ async def _download(self, data: BinaryData) -> httpx.Response:
+ if data.size > self.MAX_BINARY_BYTES:
+ raise ComposioProxyError("Proxy download exceeds the file size limit")
+ url = str(data.url)
+ for _ in range(6):
+ parsed = httpx.URL(url)
+ if parsed.scheme != "https" or parsed.userinfo or parsed.port not in (None, 443):
+ raise ComposioProxyError("Invalid proxy download URL")
+ validate_url(url)
+ # Construct a new request: never inherit client auth, cookies or provider headers.
+ request = httpx.Request(
+ "GET",
+ url,
+ extensions={
+ "timeout": {name: 180.0 for name in ("connect", "read", "write", "pool")}
+ },
+ )
+ # Use the public transport interface to avoid HTTPX logging the signed URL.
+ # This layer owns redirects and streaming, and sends no ambient credentials.
+ try:
+ response = await self._download_transport.handle_async_request(request)
+ except httpx.RequestError:
+ raise ComposioProxyError("Temporary download request failed") from None
+ if response.is_redirect:
+ location = response.headers.get("location")
+ await response.aclose()
+ if not location:
+ raise ComposioProxyError("Proxy download redirect lacks location")
+ url = str(parsed.join(location))
+ continue
+ if response.is_error:
+ status = response.status_code
+ await response.aclose()
+ raise ComposioProxyError(f"Temporary download returned HTTP {status}")
+ return response
+ raise ComposioProxyError("Too many proxy download redirects")
+
+ @staticmethod
+ def _check_proxy_response(response: httpx.Response, request: httpx.Request) -> None:
+ """Classify proxy failures independently of upstream provider status."""
+ if response.status_code in (401, 403):
+ raise AuthProviderAuthError(
+ "Composio project authentication failed", provider_name="composio"
+ )
+ if response.status_code == 429:
+ try:
+ wait = float(response.headers.get("retry-after", "30"))
+ except ValueError:
+ wait = 30.0
+ raise AuthProviderRateLimitError(provider_name="composio", retry_after=wait)
+ if response.status_code >= 500:
+ raise AuthProviderServerError(
+ provider_name="composio", status_code=response.status_code
+ )
+ if response.status_code == 422:
+ try:
+ error = response.json().get("error", {})
+ except (ValueError, AttributeError):
+ error = {}
+ if isinstance(error, dict) and error.get("code") == 1706:
+ raise AuthProviderAuthError(
+ "Composio connected account requires reauthorization", provider_name="composio"
+ )
+ if not 200 <= response.status_code < 300:
+ raise ComposioProxyError(
+ f"Composio proxy returned HTTP {response.status_code}", request=request
+ )
+
+ async def _proxy_request(self, payload: dict) -> httpx.Response:
+ """Bound encoded attachment JSON before decoding the proxy envelope."""
+ limit = ((self.MAX_BINARY_BYTES + 2) // 3) * 4 + 1024 * 1024
+ content = bytearray()
+ try:
+ async with self._client.stream(
+ "POST",
+ "https://backend.composio.dev/api/v3/tools/execute/proxy",
+ headers={"x-api-key": self._api_key},
+ json=payload,
+ ) as response:
+ async with aclosing(response.aiter_bytes()) as chunks:
+ async for chunk in chunks:
+ if len(content) + len(chunk) > limit:
+ raise ComposioProxyError(
+ "Composio response exceeds the file size limit"
+ )
+ content.extend(chunk)
+ headers = dict(response.headers)
+ for name in ("content-length", "content-encoding", "transfer-encoding"):
+ headers.pop(name, None)
+ return httpx.Response(
+ response.status_code,
+ headers=headers,
+ content=bytes(content),
+ request=response.request,
+ )
+ except ComposioProxyError:
+ raise
+ except httpx.RequestError:
+ raise AuthProviderServerError(
+ "Composio proxy request failed", provider_name="composio"
+ ) from None
+
+ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+ """Forward one provider request and preserve JSON or binary semantics."""
+ url = request.url
+ if (
+ url.scheme != "https"
+ or url.host not in self._hosts
+ or url.port not in (None, 443)
+ or url.userinfo
+ ):
+ raise ComposioProxyError("Provider URL outside the bound source", request=request)
+ parameters = [
+ {"name": name, "value": value, "type": "header"}
+ for name, value in request.headers.multi_items()
+ if name.lower()
+ not in {"authorization", "host", "cookie", "content-length", "x-api-key"}
+ ]
+ payload = {
+ "connected_account_id": self._account_id,
+ "endpoint": str(url),
+ "method": request.method,
+ "parameters": parameters,
+ }
+ body = await request.aread()
+ if body:
+ try:
+ payload["body"] = json.loads(body)
+ except (ValueError, UnicodeDecodeError) as exc:
+ raise ComposioProxyError(
+ "Proxy requires a JSON request body", request=request
+ ) from exc
+ response = await self._proxy_request(payload)
+ self._check_proxy_response(response, request)
+ try:
+ envelope = ProxyResponse.model_validate(response.json())
+ except ValueError:
+ raise ComposioProxyError(
+ "Invalid Composio response envelope", request=request
+ ) from None
+ headers = httpx.Headers(envelope.headers)
+ for name in ("content-length", "content-encoding", "transfer-encoding", "set-cookie"):
+ headers.pop(name, None)
+ if envelope.binary_data is not None and 200 <= envelope.status < 300:
+ data = envelope.binary_data
+ headers["content-type"] = data.content_type
+ if request.method == "HEAD":
+ headers["content-length"] = str(data.size)
+ return httpx.Response(
+ envelope.status, headers=headers, content=b"", request=request
+ )
+ downloaded = await self._download(data)
+ return httpx.Response(
+ envelope.status,
+ headers=headers,
+ stream=DownloadStream(downloaded, self.MAX_BINARY_BYTES),
+ request=request,
+ )
+ if isinstance(envelope.data, str):
+ content = envelope.data.encode("utf-8")
+ else:
+ content = json.dumps(envelope.data, ensure_ascii=False).encode("utf-8")
+ headers["content-type"] = "application/json"
+ return httpx.Response(envelope.status, headers=headers, content=content, request=request)
diff --git a/backend/airweave/platform/sources/retry_helpers.py b/backend/airweave/platform/http_client/retry_helpers.py
similarity index 89%
rename from backend/airweave/platform/sources/retry_helpers.py
rename to backend/airweave/platform/http_client/retry_helpers.py
index 88ca044884..eca7f3a21a 100644
--- a/backend/airweave/platform/sources/retry_helpers.py
+++ b/backend/airweave/platform/http_client/retry_helpers.py
@@ -6,11 +6,13 @@
"""
import logging
+import math
from typing import Callable
import httpx
from tenacity import retry_if_exception, wait_exponential
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.sources.exceptions import (
SourceRateLimitError,
SourceServerError,
@@ -120,16 +122,23 @@ async def _get_with_auth(self, url, params=None): ...
)
-def wait_rate_limit_with_backoff(retry_state) -> float:
+def wait_rate_limit_with_backoff(
+ retry_state, *, max_rate_limit_wait: float | None = 120.0
+) -> float:
"""Wait strategy: Retry-After for rate limits, exponential backoff for transients.
Handles both raw ``httpx.HTTPStatusError`` (429) and domain
- ``SourceRateLimitError`` / ``SourceServerError``.
+ ``SourceRateLimitError`` / ``SourceServerError`` and managed-provider throttling.
+ Set ``max_rate_limit_wait=None`` to honor the entire provider delay.
"""
exception = retry_state.outcome.exception()
- if isinstance(exception, SourceRateLimitError):
- return min(max(exception.retry_after, 1.0), 120.0)
+ def bounded_wait(seconds: float) -> float:
+ seconds = max(seconds, 1.0) if math.isfinite(seconds) else 30.0
+ return min(seconds, max_rate_limit_wait) if max_rate_limit_wait is not None else seconds
+
+ if isinstance(exception, (SourceRateLimitError, AuthProviderRateLimitError)):
+ return bounded_wait(exception.retry_after)
if isinstance(exception, SourceServerError):
return wait_exponential(multiplier=1, min=2, max=30)(retry_state)
@@ -140,7 +149,7 @@ def wait_rate_limit_with_backoff(retry_state) -> float:
try:
wait_seconds = float(retry_after)
wait_seconds = max(wait_seconds, 1.0)
- return min(wait_seconds, 120.0)
+ return bounded_wait(wait_seconds)
except (ValueError, TypeError):
pass
return wait_exponential(multiplier=1, min=2, max=30)(retry_state)
diff --git a/backend/airweave/platform/sources/__init__.py b/backend/airweave/platform/sources/__init__.py
index 0dc48be828..b5edfeb8ca 100644
--- a/backend/airweave/platform/sources/__init__.py
+++ b/backend/airweave/platform/sources/__init__.py
@@ -62,6 +62,7 @@
from .timed import TimedSource
from .todoist import TodoistSource
from .trello import TrelloSource
+from .wispr import WisprSource
from .word import WordSource
from .zendesk import ZendeskSource
from .zoho_crm import ZohoCRMSource
@@ -129,6 +130,7 @@
TimedSource,
TodoistSource,
TrelloSource,
+ WisprSource,
WordSource,
ZendeskSource,
ZoomSource,
diff --git a/backend/airweave/platform/sources/_base.py b/backend/airweave/platform/sources/_base.py
index 054ec2f818..f83c4826db 100644
--- a/backend/airweave/platform/sources/_base.py
+++ b/backend/airweave/platform/sources/_base.py
@@ -19,6 +19,7 @@
if TYPE_CHECKING:
from airweave.domains.access_control.schemas import MembershipTuple
+ from airweave.domains.entities.canonical.page_source import CanonicalPageSource
from pydantic import BaseModel
@@ -41,6 +42,8 @@ class BaseSource:
node_selections) are passed as params to ``generate_entities()``.
"""
+ _capture_page_source: CanonicalPageSource | None = None
+
# Identity (set by @source decorator — required)
is_source: ClassVar[bool] = False
source_name: ClassVar[str] = ""
@@ -104,6 +107,15 @@ def http_client(self) -> AirweaveHttpClient:
"""Pre-built HTTP client with rate limiting."""
return self._http_client
+ @property
+ def capture_page_source(self) -> CanonicalPageSource | None:
+ """The active page adapter; source identity and authentication stay on this object."""
+ from airweave.domains.entities.canonical.page_source import CanonicalPageSource
+
+ if self._capture_page_source is not None:
+ return self._capture_page_source
+ return self if isinstance(self, CanonicalPageSource) else None
+
# ------------------------------------------------------------------
# Auth convenience methods
# ------------------------------------------------------------------
diff --git a/backend/airweave/platform/sources/airtable.py b/backend/airweave/platform/sources/airtable.py
index 671e7d03d0..aba9cef64a 100644
--- a/backend/airweave/platform/sources/airtable.py
+++ b/backend/airweave/platform/sources/airtable.py
@@ -28,12 +28,12 @@
AirtableUserEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
_API = "https://api.airtable.com/v0"
diff --git a/backend/airweave/platform/sources/apollo.py b/backend/airweave/platform/sources/apollo.py
index e15ed1fbea..13b6acd397 100644
--- a/backend/airweave/platform/sources/apollo.py
+++ b/backend/airweave/platform/sources/apollo.py
@@ -26,12 +26,12 @@
ApolloSequenceEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
APOLLO_BASE_URL = "https://api.apollo.io/api/v1"
diff --git a/backend/airweave/platform/sources/asana.py b/backend/airweave/platform/sources/asana.py
index 76317764a5..2a1f407482 100644
--- a/backend/airweave/platform/sources/asana.py
+++ b/backend/airweave/platform/sources/asana.py
@@ -32,12 +32,12 @@
AsanaWorkspaceEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
_API = "https://app.asana.com/api/1.0"
diff --git a/backend/airweave/platform/sources/attio.py b/backend/airweave/platform/sources/attio.py
index 76258f05b7..f0f4e29e48 100644
--- a/backend/airweave/platform/sources/attio.py
+++ b/backend/airweave/platform/sources/attio.py
@@ -1,47 +1,55 @@
-"""Attio source implementation.
-
-Attio is a flexible CRM platform. We extract:
-- Objects (Companies, People, Deals, etc.)
-- Lists (Custom collections)
-- Records (Individual entries in objects/lists)
-- Notes (Attached to records)
-
-Note: Comments are not supported as the Attio API does not expose them via REST API.
-"""
+"""Attio originals with durable page progress and explicit discovery-only coverage."""
from __future__ import annotations
-from typing import Any, AsyncGenerator, Dict, List, Optional
+import hashlib
+import json
+from collections.abc import AsyncGenerator
+from datetime import datetime, timezone
+from functools import partial
+from uuid import UUID
-from tenacity import retry, stop_after_attempt
+from pydantic import JsonValue
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import (
- SourceEntityNotFoundError,
- SourceError,
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import CapturePage
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceAuthError
+from airweave.domains.sources.token_providers.credential import DirectCredentialProvider
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
)
-from airweave.domains.sources.token_providers.protocol import AuthProviderKind, SourceAuthProvider
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.auth import AttioAuthConfig
from airweave.platform.configs.config import AttioConfig
from airweave.platform.decorators import source
-from airweave.platform.entities._base import BaseEntity, Breadcrumb
-from airweave.platform.entities.attio import (
- AttioListEntity,
- AttioNoteEntity,
- AttioObjectEntity,
- AttioRecordEntity,
-)
+from airweave.platform.entities._base import BaseEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.attio_models import (
+ AttioNative,
+ AttioPage,
+ AttioPrincipal,
+ AttioProgress,
+)
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
_API = "https://api.attio.com/v2"
@@ -59,7 +67,20 @@
rate_limit_level=RateLimitLevel.ORG,
)
class AttioSource(BaseSource):
- """Attio source connector — syncs objects, lists, records, and notes."""
+ """Keep CRM records, list memberships and notes as separate native originals.
+
+ Offset pagination is not a snapshot. Every scope is discovery-only, so missing
+ rows never authorize deletion. Exact-read and webhook qualification remains
+ required before this connector can claim complete synchronization.
+ """
+
+ canonical_record_types = ("object", "record", "list", "entry", "note")
+ canonical_container_parents = {"record": "object", "entry": "list", "note": "record"}
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Bind the engine's existing durable cycle to the attested workspace."""
+ return self._cycle_configuration
@classmethod
async def create(
@@ -70,475 +91,213 @@ async def create(
http_client: AirweaveHttpClient,
config: AttioConfig,
) -> AttioSource:
- """Create a new Attio source instance."""
+ """Attest account identity without extracting managed credentials."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
- if auth.provider_kind == AuthProviderKind.CREDENTIAL:
- instance._api_key = auth.credentials.api_key
- else:
- instance._api_key = await auth.get_token()
- return instance
-
- # ------------------------------------------------------------------
- # HTTP
- # ------------------------------------------------------------------
-
- def _headers(self) -> Dict[str, str]:
- """Build request headers with Attio bearer token."""
- return {
- "Authorization": f"Bearer {self._api_key}",
- "Content-Type": "application/json",
- "Accept": "application/json",
+ instance.config = config
+ await instance.validate()
+ fingerprint = {
+ "version": 1,
+ "workspace": str(config.workspace_id),
+ "principal": instance.principal.model_dump(mode="json"),
+ "kinds": cls.canonical_record_types,
}
-
- @retry(
- stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
- reraise=True,
- )
- async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
- """Authenticated GET with retry on 429/5xx/timeout."""
- response = await self.http_client.get(
- url, headers=self._headers(), params=params, timeout=20.0
- )
- raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ instance._cycle_configuration = CycleConfiguration.from_source(
+ fingerprint=hashlib.sha256(
+ json.dumps(fingerprint, sort_keys=True).encode()
+ ).hexdigest(),
+ record_types=cls.canonical_record_types,
+ container_parents=cls.canonical_container_parents,
+ completion_policies={kind: "discovery_only" for kind in cls.canonical_record_types},
)
- return response.json()
+ return instance
+
+ async def _headers(self, *, refresh: bool = False) -> dict[str, str]:
+ if isinstance(self.auth, DirectCredentialProvider):
+ credentials = AttioAuthConfig.model_validate(self.auth.credentials.model_dump())
+ return {"Authorization": f"Bearer {credentials.api_key}"}
+ return await authorization_headers(self.auth, refresh=refresh)
@retry(
stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
reraise=True,
)
- async def _post(self, url: str, json_data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
- """Authenticated POST with retry on 429/5xx/timeout."""
- response = await self.http_client.post(
- url, headers=self._headers(), json=json_data or {}, timeout=20.0
- )
+ async def _request(
+ self,
+ path: str,
+ *,
+ body: dict[str, JsonValue] | None = None,
+ params: dict[str, str | int] | None = None,
+ ) -> dict[str, JsonValue]:
+ headers = await self._headers()
+
+ async def send():
+ if body is not None:
+ return await self.http_client.post(_API + path, headers=headers, json=body)
+ return await self.http_client.get(_API + path, headers=headers, params=params)
+
+ response = await send()
+ if response.status_code == 401 and self.auth.supports_refresh:
+ headers = await self._headers(refresh=True)
+ response = await send()
+ # In particular, 404 merge_in_progress must never become an empty page.
raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ response, source_short_name="attio", token_provider_kind=self.auth.provider_kind
)
return response.json()
- # ------------------------------------------------------------------
- # Pagination
- # ------------------------------------------------------------------
-
- async def _paginate_get(
- self, url: str, data_key: str = "data", limit: int = 100
- ) -> AsyncGenerator[Dict[str, Any], None]:
- """Paginate through a GET endpoint using offset/limit. Yields items."""
- offset = 0
- while True:
- params = {"offset": offset, "limit": limit}
- data = await self._get(url, params)
- items = data.get(data_key, [])
- for item in items:
- yield item
- if len(items) < limit:
- break
- offset += limit
-
- async def _paginate_post(
- self, url: str, data_key: str = "data", limit: int = 100
- ) -> AsyncGenerator[Dict[str, Any], None]:
- """Paginate through a POST query endpoint using offset/limit. Yields items."""
- offset = 0
- while True:
- body = {"offset": offset, "limit": limit}
- data = await self._post(url, body)
- items = data.get(data_key, [])
- for item in items:
- yield item
- if len(items) < limit:
- break
- offset += limit
-
- # ------------------------------------------------------------------
- # Entity generators
- # ------------------------------------------------------------------
-
- async def _generate_objects(self) -> AsyncGenerator[AttioObjectEntity, None]:
- """Generate Attio Object entities (Companies, People, Deals, etc.)."""
- self.logger.info("Fetching Attio objects...")
- async for obj in self._paginate_get(f"{_API}/objects"):
- object_id = obj.get("id", {}).get("object_id")
- if not object_id:
- continue
-
- singular_noun = obj.get("singular_noun", "")
- display_name = singular_noun or object_id
-
- yield AttioObjectEntity(
- object_id=object_id,
- breadcrumbs=[],
- name=display_name,
- created_at=obj.get("created_at"),
- singular_noun=display_name,
- plural_noun=obj.get("plural_noun", display_name + "s"),
- api_slug=obj.get("api_slug", object_id),
- icon=obj.get("icon"),
- )
-
- async def _generate_lists(self) -> AsyncGenerator[AttioListEntity, None]:
- """Generate Attio List entities."""
- self.logger.info("Fetching Attio lists...")
- async for lst in self._paginate_get(f"{_API}/lists"):
- list_id = lst.get("id", {}).get("list_id")
- if not list_id:
- continue
-
- parent_object = lst.get("parent_object")
- if isinstance(parent_object, list):
- parent_object = ", ".join(parent_object) if parent_object else None
-
- yield AttioListEntity(
- list_id=list_id,
- breadcrumbs=[],
- name=lst.get("name", list_id),
- created_at=lst.get("created_at"),
- workspace_id=lst.get("workspace_id", ""),
- parent_object=parent_object,
- )
-
- async def _generate_records_for_object(
- self,
- object_slug: str,
- object_name: str,
- object_breadcrumb: Breadcrumb,
- ) -> AsyncGenerator[AttioRecordEntity, None]:
- """Generate record entities for a specific object."""
- self.logger.debug(f"Fetching records for object: {object_slug}")
- url = f"{_API}/objects/{object_slug}/records/query"
-
- try:
- async for record in self._paginate_post(url):
- entity = self._build_object_record(
- record, object_slug, object_name, object_breadcrumb
- )
- if entity:
- yield entity
- except SourceEntityNotFoundError:
- self.logger.warning(f"Object {object_slug} not found or not accessible, skipping")
-
- async def _generate_records_for_list(
- self,
- list_id: str,
- list_name: str,
- list_breadcrumb: Breadcrumb,
- ) -> AsyncGenerator[AttioRecordEntity, None]:
- """Generate record entities for a specific list."""
- self.logger.debug(f"Fetching records for list: {list_id}")
- url = f"{_API}/lists/{list_id}/records/query"
-
- try:
- async for record in self._paginate_post(url):
- entity = self._build_list_record(record, list_id, list_name, list_breadcrumb)
- if entity:
- yield entity
- except SourceEntityNotFoundError:
- self.logger.debug(f"List {list_id} not found or not accessible, skipping")
-
- async def _generate_notes_for_record(
- self,
- parent_object_or_list_id: str,
- record_id: str,
- record_breadcrumbs: List[Breadcrumb],
- ) -> AsyncGenerator[AttioNoteEntity, None]:
- """Generate note entities for a specific record."""
- url = f"{_API}/notes"
- params = {
- "parent_object": parent_object_or_list_id,
- "parent_record_id": record_id,
- "limit": 50,
- }
-
- try:
- data = await self._get(url, params=params)
- except SourceEntityNotFoundError:
- return
- except SourceError as exc:
- self.logger.warning(f"Error fetching notes for record {record_id}: {exc}")
- return
-
- for note in data.get("data", []):
- note_id = note.get("id", {}).get("note_id")
- if not note_id:
- continue
-
- title = note.get("title")
- content = note.get("content", "")
- if title:
- note_name = title
- else:
- note_name = content[:50] + "..." if len(content) > 50 else content
- if not note_name:
- note_name = f"Note {note_id}"
-
- yield AttioNoteEntity(
- note_id=note_id,
- breadcrumbs=record_breadcrumbs,
- name=note_name,
- created_at=note.get("created_at"),
- updated_at=note.get("updated_at"),
- parent_record_id=record_id,
- parent_object=parent_object_or_list_id,
- title=title,
- content=content,
- format=note.get("format"),
- author=note.get("author"),
- permalink_url=None,
+ async def validate(self) -> None:
+ """HTTP200 with active=false is revoked access, never identity evidence."""
+ principal = AttioPrincipal.model_validate(await self._request("/self"))
+ if not principal.active:
+ raise SourceAuthError(
+ "Attio token is inactive; reconnect the account",
+ source_short_name="attio",
+ status_code=200,
+ token_provider_kind=self.auth.provider_kind,
)
-
- # ------------------------------------------------------------------
- # Record builders (extract attribute parsing from generator loops)
- # ------------------------------------------------------------------
-
- def _build_object_record( # noqa: C901
- self,
- record: Dict[str, Any],
- object_slug: str,
- object_name: str,
- breadcrumb: Breadcrumb,
- ) -> Optional[AttioRecordEntity]:
- """Build a record entity from raw API data for an object."""
- record_id = record.get("id", {}).get("record_id")
- if not record_id or record.get("deleted_at"):
- return None
-
- values = record.get("values", {})
- name = None
- description = None
- email_addresses: List[Dict[str, Any]] = []
- phone_numbers: List[Dict[str, Any]] = []
- domains: List[str] = []
- categories: List[str] = []
- attributes: Dict[str, Any] = {}
-
- for attr_key, attr_values in values.items():
- if not attr_values:
- continue
-
- first_val = attr_values[0] if isinstance(attr_values, list) else attr_values
- key_lower = attr_key.lower()
-
- if "name" in key_lower or "title" in key_lower:
- name = (
- (first_val.get("value") or first_val.get("text"))
- if isinstance(first_val, dict)
- else str(first_val)
- )
- elif "description" in key_lower or "notes" in key_lower:
- description = (
- (first_val.get("value") or first_val.get("text"))
- if isinstance(first_val, dict)
- else str(first_val)
- )
- elif "email" in key_lower:
- email_addresses = [v for v in attr_values if isinstance(v, dict)]
- elif "phone" in key_lower:
- phone_numbers = [v for v in attr_values if isinstance(v, dict)]
- elif "domain" in key_lower:
- domains = [
- v.get("domain") if isinstance(v, dict) else v
- for v in attr_values
- if v is not None
- and (isinstance(v, dict) and v.get("domain") or isinstance(v, str))
- ]
- elif "category" in key_lower or "tag" in key_lower:
- categories = [
- v.get("value") if isinstance(v, dict) else v
- for v in attr_values
- if v is not None
- and (isinstance(v, dict) and v.get("value") or isinstance(v, str))
- ]
- attributes[attr_key] = attr_values
-
- return AttioRecordEntity(
- record_id=record_id,
- breadcrumbs=[breadcrumb],
- name=name or record_id,
- created_at=record.get("created_at"),
- updated_at=record.get("updated_at"),
- object_id=object_slug,
- list_id=None,
- parent_object_name=object_name,
- description=description,
- email_addresses=email_addresses,
- phone_numbers=phone_numbers,
- domains=domains,
- categories=categories,
- attributes=attributes,
- permalink_url=None,
+ if principal.workspace_id != self.config.workspace_id:
+ raise ValueError("Attio workspace identity does not match")
+ self.principal = principal
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Use provider-global UUIDs while retaining the full canonical parent."""
+ if self.canonical_container_parents.get(record_type) != parent.identity.record_type:
+ raise ValueError("Unsupported Attio parent relationship")
+ return CompletedScope(
+ record_type=record_type, container_id=parent.identity.native_id, parent=parent.identity
)
- def _build_list_record( # noqa: C901
+ async def capture_page(
self,
- record: Dict[str, Any],
- list_id: str,
- list_name: str,
- breadcrumb: Breadcrumb,
- ) -> Optional[AttioRecordEntity]:
- """Build a record entity from raw API data for a list."""
- record_id = record.get("id", {}).get("record_id")
- if not record_id or record.get("deleted_at"):
- return None
-
- values = record.get("values", {})
- name = None
- description = None
- attributes: Dict[str, Any] = {}
-
- for attr_key, attr_values in values.items():
- if not attr_values:
- continue
-
- first_val = attr_values[0] if isinstance(attr_values, list) else attr_values
- key_lower = attr_key.lower()
-
- if "name" in key_lower or "title" in key_lower:
- name = (
- (first_val.get("value") or first_val.get("text"))
- if isinstance(first_val, dict)
- else str(first_val)
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Fetch one atomic page, retaining raw native fields and a resumable offset."""
+ progress = AttioProgress.model_validate(continuation.value)
+ kind = scope.record_type
+ limit = 50 if kind == "note" else 500
+ root = kind in {"object", "list"}
+ object_slug: str | None = None
+ if root:
+ if scope.container_id is not None or scope.parent is not None or progress.offset:
+ raise ValueError("Invalid Attio root scope")
+ data = await self._request("/objects" if kind == "object" else "/lists")
+ else:
+ if parent is None or scope != self.child_scope(parent, kind):
+ raise ValueError("Attio child scope requires its exact captured parent")
+ UUID(scope.container_id)
+ if kind == "record":
+ data = await self._request(
+ f"/objects/{scope.container_id}/records/query",
+ body={"limit": limit, "offset": progress.offset},
)
- elif "description" in key_lower:
- description = (
- (first_val.get("value") or first_val.get("text"))
- if isinstance(first_val, dict)
- else str(first_val)
+ elif kind == "entry":
+ data = await self._request(
+ f"/lists/{scope.container_id}/entries/query",
+ body={"limit": limit, "offset": progress.offset},
)
- attributes[attr_key] = attr_values
-
- return AttioRecordEntity(
- record_id=record_id,
- breadcrumbs=[breadcrumb],
- name=name or record_id,
- created_at=record.get("created_at"),
- updated_at=record.get("updated_at"),
- object_id=None,
- list_id=list_id,
- parent_object_name=list_name,
- description=description,
- email_addresses=[],
- phone_numbers=[],
- domains=[],
- categories=[],
- attributes=attributes,
- permalink_url=None,
+ elif kind == "note":
+ object_id = parent.identity.container_id
+ UUID(object_id)
+ object_response = await self._request(f"/objects/{object_id}")
+ native_object = AttioNative.model_validate(object_response.get("data"))
+ if (
+ native_object.id.workspace_id != self.config.workspace_id
+ or str(native_object.id.object_id) != object_id
+ or not native_object.api_slug
+ ):
+ raise ValueError("Attio note parent object identity is invalid")
+ object_slug = native_object.api_slug
+ data = await self._request(
+ "/notes",
+ params={
+ "parent_object": object_id,
+ "parent_record_id": scope.container_id,
+ "limit": limit,
+ "offset": progress.offset,
+ },
+ )
+ else:
+ raise ValueError("Unsupported Attio scope")
+ page = AttioPage.model_validate(data)
+ if len(page.data) > limit:
+ raise ValueError("Attio returned more rows than requested")
+ records = tuple(self._capture(scope, raw, parent, object_slug) for raw in page.data)
+ ids = tuple(record.identity.native_id for record in records)
+ if len(set(ids)) != len(ids) or set(ids).intersection(progress.previous_ids):
+ raise ValueError("Attio offset page repeated identities; restart discovery")
+ return CapturePage(
+ records=records,
+ final=root or len(records) < limit,
+ continuation=ScanContinuation(
+ value=AttioProgress(
+ offset=progress.offset + len(records), previous_ids=ids
+ ).model_dump(mode="json")
+ ),
)
- # ------------------------------------------------------------------
- # Orchestration helpers
- # ------------------------------------------------------------------
-
- async def _yield_record_with_notes(
+ def _capture(
self,
- record: AttioRecordEntity,
- parent_slug_or_id: str,
- parent_breadcrumb: Breadcrumb,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Yield a record entity followed by its notes."""
- yield record
-
- record_breadcrumbs = [
- parent_breadcrumb,
- Breadcrumb(
- entity_id=record.record_id,
- name=record.name,
- entity_type="AttioRecordEntity",
+ scope: CompletedScope,
+ raw: dict[str, JsonValue],
+ parent: SourceRecord | None,
+ object_slug: str | None = None,
+ ) -> CaptureRecord:
+ item = AttioNative.model_validate(raw)
+ if item.id.workspace_id != self.config.workspace_id:
+ raise ValueError("Attio returned a different workspace")
+ kind = scope.record_type
+ identity = {
+ "object": item.id.object_id,
+ "record": item.id.record_id,
+ "list": item.id.list_id,
+ "entry": item.id.entry_id,
+ "note": item.id.note_id,
+ }.get(kind)
+ if identity is None:
+ raise ValueError("Attio omitted the native identity")
+ if kind == "record" and str(item.id.object_id) != scope.container_id:
+ raise ValueError("Attio record belongs to a different object")
+ if kind == "entry" and (
+ str(item.id.list_id) != scope.container_id
+ or item.parent_record_id is None
+ or not item.parent_object
+ ):
+ raise ValueError("Attio entry has an invalid list or record reference")
+ if kind == "note" and (
+ parent is None
+ or str(item.parent_record_id) != scope.container_id
+ or item.parent_object not in {parent.identity.container_id, object_slug}
+ ):
+ raise ValueError("Attio note belongs to a different object or record")
+ return CaptureRecord(
+ identity=RecordIdentity(
+ record_type=kind, native_id=str(identity), container_id=scope.container_id
),
- ]
-
- try:
- async for note in self._generate_notes_for_record(
- parent_slug_or_id, record.record_id, record_breadcrumbs
- ):
- yield note
- except Exception as e:
- self.logger.warning(
- f"Error fetching notes for record {record.record_id}: {e}",
- exc_info=True,
- )
+ parent=scope.parent,
+ payload=raw,
+ source_created_at=item.created_at,
+ observed_at=datetime.now(timezone.utc),
+ )
- # ------------------------------------------------------------------
- # Main entry point
- # ------------------------------------------------------------------
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Offset omissions and ambiguous404 responses cannot certify deletion."""
+ raise ValueError(
+ "Attio discovery cannot confirm absence without qualified provider evidence"
+ )
- async def generate_entities( # noqa: C901
+ async def generate_entities(
self,
*,
cursor: SyncCursor | None = None,
files: FileService | None = None,
node_selections: list[NodeSelectionData] | None = None,
) -> AsyncGenerator[BaseEntity, None]:
- """Generate all entities from Attio."""
- self.logger.info("Starting Attio sync...")
-
- object_map: Dict[str, AttioObjectEntity] = {}
- async for obj in self._generate_objects():
- yield obj
- object_map[obj.object_id] = obj
-
- list_map: Dict[str, AttioListEntity] = {}
- async for lst in self._generate_lists():
- yield lst
- list_map[lst.list_id] = lst
-
- for object_id, obj in object_map.items():
- object_breadcrumb = Breadcrumb(
- entity_id=object_id, name=obj.name, entity_type="AttioObjectEntity"
- )
-
- records: List[AttioRecordEntity] = []
- async for record in self._generate_records_for_object(
- obj.api_slug, obj.singular_noun, object_breadcrumb
- ):
- records.append(record)
-
- async def _object_record_worker(record, _obj=obj, _bc=object_breadcrumb):
- async for entity in self._yield_record_with_notes(record, _obj.api_slug, _bc):
- yield entity
-
- async for entity in self.process_entities_concurrent(
- items=records,
- worker=_object_record_worker,
- batch_size=10,
- preserve_order=False,
- stop_on_error=False,
- ):
- yield entity
-
- for list_id, lst in list_map.items():
- list_breadcrumb = Breadcrumb(
- entity_id=list_id, name=lst.name, entity_type="AttioListEntity"
- )
-
- records = []
- async for record in self._generate_records_for_list(list_id, lst.name, list_breadcrumb):
- records.append(record)
-
- async def _list_record_worker(record, _list_id=list_id, _bc=list_breadcrumb):
- async for entity in self._yield_record_with_notes(record, _list_id, _bc):
- yield entity
-
- async for entity in self.process_entities_concurrent(
- items=records,
- worker=_list_record_worker,
- batch_size=10,
- preserve_order=False,
- stop_on_error=False,
- ):
- yield entity
-
- self.logger.info("Attio sync completed")
-
- async def validate(self) -> None:
- """Verify credentials by pinging the Attio API."""
- await self._get(f"{_API}/objects", params={"limit": 1})
+ """Canonical page capture replaces the legacy lossy entity generator."""
+ raise NotImplementedError("Attio requires canonical page capture")
+ yield # pragma: no cover
diff --git a/backend/airweave/platform/sources/attio_models.py b/backend/airweave/platform/sources/attio_models.py
new file mode 100644
index 0000000000..345c749d59
--- /dev/null
+++ b/backend/airweave/platform/sources/attio_models.py
@@ -0,0 +1,57 @@
+"""Attio native response boundaries and bounded durable page progress."""
+
+from typing import Literal
+from uuid import UUID
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool
+
+
+class AttioModel(BaseModel):
+ """Validate consumed fields while preserving raw payloads separately."""
+
+ model_config = ConfigDict(extra="ignore")
+
+
+class AttioPrincipal(AttioModel):
+ """Native token identity; active is the only provider-guaranteed field."""
+
+ active: StrictBool
+ workspace_id: UUID | None = None
+ token_level: Literal["workspace", "user"] | None = None
+ authorized_by_workspace_member_id: UUID | None = None
+
+
+class AttioIdentity(AttioModel):
+ """Native composite IDs, validated by record kind at acquisition."""
+
+ workspace_id: UUID
+ object_id: UUID | None = None
+ list_id: UUID | None = None
+ record_id: UUID | None = None
+ entry_id: UUID | None = None
+ note_id: UUID | None = None
+
+
+class AttioNative(AttioModel):
+ """Identity and provenance shared by native CRM representations."""
+
+ id: AttioIdentity
+ created_at: AwareDatetime | None = None
+ parent_record_id: UUID | None = None
+ parent_object: str | None = None
+ api_slug: str | None = None
+
+
+class AttioPage(AttioModel):
+ """A bounded native page; overflow fails instead of truncating."""
+
+ data: list[dict[str, JsonValue]] = Field(max_length=500)
+
+
+class AttioProgress(BaseModel):
+ """Durable offset plus the previous page IDs for overlap detection."""
+
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ offset: int = Field(default=0, ge=0)
+ previous_ids: tuple[str, ...] = Field(default=(), max_length=500)
diff --git a/backend/airweave/platform/sources/bitbucket.py b/backend/airweave/platform/sources/bitbucket.py
index 7552afbbd7..6febe44fc7 100644
--- a/backend/airweave/platform/sources/bitbucket.py
+++ b/backend/airweave/platform/sources/bitbucket.py
@@ -27,12 +27,12 @@
BitbucketWorkspaceEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.platform.utils.file_extensions import (
get_language_for_extension,
is_text_file,
diff --git a/backend/airweave/platform/sources/box.py b/backend/airweave/platform/sources/box.py
index bde1d66e6b..0676d97f84 100644
--- a/backend/airweave/platform/sources/box.py
+++ b/backend/airweave/platform/sources/box.py
@@ -27,12 +27,12 @@
BoxUserEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/clickup.py b/backend/airweave/platform/sources/clickup.py
index da6d27aad9..04554c9e7b 100644
--- a/backend/airweave/platform/sources/clickup.py
+++ b/backend/airweave/platform/sources/clickup.py
@@ -31,12 +31,12 @@
_parse_clickup_ts,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/coda.py b/backend/airweave/platform/sources/coda.py
index 41a20391d4..1aa3b5b76b 100644
--- a/backend/airweave/platform/sources/coda.py
+++ b/backend/airweave/platform/sources/coda.py
@@ -32,12 +32,12 @@
_parse_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
CODA_API_BASE = "https://coda.io/apis/v1"
diff --git a/backend/airweave/platform/sources/confluence.py b/backend/airweave/platform/sources/confluence.py
index 8dd486d29e..cd969b8b67 100644
--- a/backend/airweave/platform/sources/confluence.py
+++ b/backend/airweave/platform/sources/confluence.py
@@ -41,12 +41,12 @@
ConfluenceSpaceEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
ATLASSIAN_ACCESSIBLE_RESOURCES_URL = "https://api.atlassian.com/oauth/token/accessible-resources"
diff --git a/backend/airweave/platform/sources/dropbox.py b/backend/airweave/platform/sources/dropbox.py
index 75c01d5b45..6515ee0132 100644
--- a/backend/airweave/platform/sources/dropbox.py
+++ b/backend/airweave/platform/sources/dropbox.py
@@ -27,12 +27,12 @@
DropboxFolderEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/fireflies.py b/backend/airweave/platform/sources/fireflies.py
index 163e3dcec6..b3b6c7fd98 100644
--- a/backend/airweave/platform/sources/fireflies.py
+++ b/backend/airweave/platform/sources/fireflies.py
@@ -26,12 +26,12 @@
from airweave.platform.entities._base import BaseEntity
from airweave.platform.entities.fireflies import FirefliesTranscriptEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
FIREFLIES_GRAPHQL_URL = "https://api.fireflies.ai/graphql"
diff --git a/backend/airweave/platform/sources/freshdesk.py b/backend/airweave/platform/sources/freshdesk.py
index cd8cebcba8..61d0697bfa 100644
--- a/backend/airweave/platform/sources/freshdesk.py
+++ b/backend/airweave/platform/sources/freshdesk.py
@@ -33,12 +33,12 @@
FreshdeskTicketEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
diff --git a/backend/airweave/platform/sources/github.py b/backend/airweave/platform/sources/github.py
index 4124decb0b..80d1410707 100644
--- a/backend/airweave/platform/sources/github.py
+++ b/backend/airweave/platform/sources/github.py
@@ -1,51 +1,67 @@
-r"""GitHub source implementation for syncing repositories, directories, and code files."""
+"""Selected GitHub repositories captured through the existing durable page engine."""
from __future__ import annotations
-import base64
+import hashlib
+import json
import mimetypes
-from pathlib import Path
-from typing import Any, AsyncGenerator, Dict, List, Optional
+from collections.abc import AsyncGenerator
+from datetime import datetime, timezone
+from urllib.parse import quote
-from tenacity import retry, stop_after_attempt
+from pydantic import JsonValue
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import SourceAuthError
-from airweave.domains.sources.token_providers.protocol import (
- AuthProviderKind,
- TokenProviderProtocol,
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ InvalidScanContinuation,
+ ScopeAccessLost,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+ parent_container_key,
)
-from airweave.domains.storage import FileSkippedException
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceAuthError
+from airweave.domains.sources.token_providers.protocol import SourceAuthProvider
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.auth import GitHubAuthConfig
-from airweave.platform.configs.config import GitHubConfig
-from airweave.platform.cursors import GitHubCursor
+from airweave.platform.configs.config import GitHubConfig, GitHubRepositorySelection
from airweave.platform.decorators import source
-from airweave.platform.entities._base import BaseEntity, Breadcrumb
-from airweave.platform.entities.github import (
- GitHubCodeFileEntity,
- GitHubDirectoryEntity,
- GitHubFileDeletionEntity,
- GitHubPRCommentEntity,
- GitHubPullRequestEntity,
- GitHubRepositoryEntity,
-)
+from airweave.platform.entities._base import BaseEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
- retry_if_rate_limit_or_timeout,
- wait_rate_limit_with_backoff,
-)
-from airweave.platform.utils.file_extensions import (
- get_language_for_extension,
- is_text_file,
+from airweave.platform.sources.github_http import GitHubReader, GitHubUnavailable
+from airweave.platform.sources.github_models import (
+ Branch,
+ BranchPayload,
+ Dates,
+ FileProgress,
+ GitCommit,
+ Identity,
+ Issue,
+ IssuePayload,
+ ListProgress,
+ PullRequest,
+ Repository,
+ RepositoryContext,
+ RootProgress,
+ Tree,
+ TreeEntry,
+ TreeFrame,
)
from airweave.schemas.source_connection import AuthenticationMethod
+FIELD_SET_VERSION = 1
+MAX_CODE_BYTES = 10 * 1024 * 1024
+
@source(
name="GitHub",
@@ -55,701 +71,480 @@
auth_config_class=GitHubAuthConfig,
config_class=GitHubConfig,
labels=["Code"],
- supports_continuous=True,
- supports_temporal_relevance=False,
- cursor_class=GitHubCursor,
+ supports_continuous=False,
rate_limit_level=RateLimitLevel.ORG,
)
class GitHubSource(BaseSource):
- """GitHub source connector integrates with the GitHub REST API to extract and synchronize data.
-
- Connects to your GitHub repositories.
-
- It supports syncing repository metadata, directory structures, and code files with
- configurable filtering options for branches and file types.
- """
-
- BASE_URL = "https://api.github.com"
+ """Full observed sweeps; no claim of an upstream transactional snapshot or event log."""
+
+ canonical_record_types = (
+ "repository",
+ "branch",
+ "file",
+ "issue",
+ "comment",
+ "review",
+ "review_comment",
+ "pull_file",
+ )
+ canonical_container_parents = {
+ "branch": "repository",
+ "file": "branch",
+ "issue": "repository",
+ "comment": "issue",
+ "review": "issue",
+ "review_comment": "issue",
+ "pull_file": "issue",
+ }
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Bind selections and field policy to existing cycle restart semantics."""
+ return self._cycle_configuration
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: GitHubConfig,
) -> GitHubSource:
- """Create a new source instance with authentication."""
+ """Do not fetch or extract managed provider credentials."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
- if auth.provider_kind == AuthProviderKind.CREDENTIAL:
- instance._personal_access_token = auth.credentials.personal_access_token
- else:
- instance._personal_access_token = await auth.get_token()
- instance._repo_name = config.repo_name
- instance._branch = config.branch or None
- instance._sync_pull_requests = config.sync_pull_requests
- instance._max_file_size = 10 * 1024 * 1024
- return instance
-
- def get_default_cursor_field(self) -> Optional[str]:
- """Get the default cursor field for GitHub source."""
- return "last_repository_pushed_at"
-
- def validate_cursor_field(self, cursor_field: str) -> None:
- """Validate if the given cursor field is valid for GitHub.
-
- Args:
- cursor_field: The cursor field to validate
-
- Raises:
- ValueError: If the cursor field is invalid
- """
- valid_field = self.get_default_cursor_field()
- if cursor_field != valid_field:
- error_msg = (
- f"Invalid cursor field '{cursor_field}' for GitHub source. "
- f"GitHub requires '{valid_field}' as the cursor field. "
- f"GitHub tracks repository changes using push timestamps, not entity fields. "
- f"Please use the default cursor field or omit it entirely."
- )
- self.logger.warning(error_msg)
- raise ValueError(error_msg)
-
- @retry(
- stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
- reraise=True,
- )
- async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
- """Make authenticated API request using Personal Access Token."""
- headers = {
- "Authorization": f"token {self._personal_access_token}",
- "Accept": "application/vnd.github.v3+json",
- "X-GitHub-Api-Version": "2022-11-28",
- }
- response = await self.http_client.get(url, headers=headers, params=params)
-
- if response.status_code == 401 and self.auth.supports_refresh:
- new_token = await self.auth.force_refresh()
- headers["Authorization"] = f"token {new_token}"
- response = await self.http_client.get(url, headers=headers, params=params)
-
- raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ instance.config = config
+ instance.reader = GitHubReader(http_client, auth)
+ instance.selections = {item.repository_id: item for item in config.repositories}
+ fingerprint = {"config": config.model_dump(mode="json"), "fields": FIELD_SET_VERSION}
+ instance._cycle_configuration = CycleConfiguration.from_source(
+ fingerprint=hashlib.sha256(
+ json.dumps(fingerprint, sort_keys=True).encode()
+ ).hexdigest(),
+ record_types=cls.canonical_record_types,
+ container_parents=cls.canonical_container_parents,
)
- return response.json()
-
- async def _get_paginated_results(
- self, url: str, params: Optional[Dict[str, Any]] = None
- ) -> List[Dict[str, Any]]:
- """Get all pages of results from a paginated GitHub API endpoint."""
- if params is None:
- params = {}
-
- params["per_page"] = 100
-
- all_results = []
- page = 1
-
- while True:
- params["page"] = page
- headers = {
- "Authorization": f"token {self._personal_access_token}",
- "Accept": "application/vnd.github.v3+json",
- "X-GitHub-Api-Version": "2022-11-28",
- }
-
- response = await self.http_client.get(url, headers=headers, params=params)
-
- if response.status_code == 401 and self.auth.supports_refresh:
- new_token = await self.auth.force_refresh()
- headers["Authorization"] = f"token {new_token}"
- response = await self.http_client.get(url, headers=headers, params=params)
+ return instance
- raise_for_status(
- response,
- source_short_name=self.short_name,
+ async def validate(self) -> None:
+ """Verify authentication independently from permission to every selected repository."""
+ principal = Identity.model_validate(await self.reader.object("/user"))
+ if principal.id != self.config.expected_user_id:
+ raise SourceAuthError(
+ "GitHub authenticated user does not match the configured account",
+ source_short_name="github",
+ status_code=200,
token_provider_kind=self.auth.provider_kind,
)
- results = response.json()
- if not results:
- break
-
- all_results.extend(results)
-
- link_header = response.headers.get("Link", "")
- if 'rel="next"' not in link_header:
- break
-
- page += 1
-
- return all_results
-
- def _detect_language_from_extension(self, file_path: str) -> str:
- """Detect programming language from file extension."""
- ext = Path(file_path).suffix.lower()
- return get_language_for_extension(ext)
-
- def _check_repository_updates(
- self, repo_name: str, last_pushed_at: str, current_pushed_at: str
- ) -> bool:
- """Check if repository has been updated since last sync."""
- if last_pushed_at:
- self.logger.debug(
- f"Repository {repo_name} pushed_at: {last_pushed_at} -> {current_pushed_at}"
- )
- has_updates = current_pushed_at > last_pushed_at
- if has_updates:
- self.logger.debug(f"Repository {repo_name} has new commits since last sync")
- else:
- self.logger.debug(f"Repository {repo_name} has no new commits since last sync")
- return has_updates
- else:
- self.logger.debug(f"First sync for repository {repo_name}")
- return True
-
- async def _get_repository_info(
- self, repo_name: str, cursor: SyncCursor | None = None
- ) -> GitHubRepositoryEntity:
- """Get repository information with cursor support."""
- url = f"{self.BASE_URL}/repos/{repo_name}"
- repo_data = await self._get(url)
-
- cursor_data = cursor.data if cursor else {}
-
- last_pushed_at = cursor_data.get("last_repository_pushed_at")
- current_pushed_at = repo_data["pushed_at"]
-
- self._check_repository_updates(repo_name, last_pushed_at, current_pushed_at)
-
- if cursor:
- cursor.update(
- last_repository_pushed_at=current_pushed_at,
- repo_name=repo_name,
- branch=self._branch,
- )
-
- return GitHubRepositoryEntity.from_api(repo_data)
-
- async def _traverse_repository(
- self,
- repo_name: str,
- branch: str,
- cursor: SyncCursor | None = None,
- files: FileService | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Traverse repository contents using DFS."""
- repo_entity = await self._get_repository_info(repo_name, cursor=cursor)
- yield repo_entity
-
- owner, repo = repo_name.split("/")
-
- repo_breadcrumb = Breadcrumb(
- entity_id=str(repo_entity.repo_id),
- name=repo_entity.name,
- entity_type=GitHubRepositoryEntity.__name__,
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Use one bounded, stable parent identity policy at every hierarchy depth."""
+ if self.canonical_container_parents.get(record_type) != parent.identity.record_type:
+ raise ValueError("GitHub child scope does not belong to this parent")
+ return CompletedScope(
+ record_type=record_type,
+ container_id=parent_container_key(parent.identity),
+ parent=parent.identity,
)
- processed_paths = set()
- processed_files = set()
-
- async for entity in self._traverse_directory(
- repo_name,
- "",
- [repo_breadcrumb],
- owner,
- repo,
- branch,
- processed_paths,
- processed_files,
- files=files,
- ):
- yield entity
-
- async def _traverse_repository_incremental(
- self,
- repo_name: str,
- branch: str,
- since_timestamp: str,
- cursor: SyncCursor | None = None,
- files: FileService | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Traverse repository contents incrementally using commits since last sync."""
- repo_entity = await self._get_repository_info(repo_name, cursor=cursor)
- yield repo_entity
-
- owner, repo = repo_name.split("/")
-
- repo_breadcrumb = Breadcrumb(
- entity_id=str(repo_entity.repo_id),
- name=repo_entity.name,
- entity_type=GitHubRepositoryEntity.__name__,
+ @staticmethod
+ def _record(
+ scope: CompletedScope, native_id: str, payload: dict[str, JsonValue]
+ ) -> CaptureRecord:
+ native = payload.get("issue", payload.get("native", payload))
+ dates = Dates.model_validate(native)
+ body = native.get("body") if isinstance(native, dict) else None
+ linked_attachment = isinstance(body, str) and any(
+ host in body
+ for host in ("github.com/user-attachments/", "user-images.githubusercontent.com/")
+ )
+ return CaptureRecord(
+ identity=RecordIdentity(
+ record_type=scope.record_type, native_id=native_id, container_id=scope.container_id
+ ),
+ parent=scope.parent,
+ payload=payload,
+ payload_schema_version=FIELD_SET_VERSION,
+ completeness="partial" if linked_attachment else "complete",
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=dates.created_at,
+ source_updated_at=dates.updated_at,
)
- commits = await self._get_commits_since(repo_name, since_timestamp, branch)
-
- if not commits:
- self.logger.debug(f"No new commits found since {since_timestamp}")
- return
-
- self.logger.debug(f"Found {len(commits)} new commits since {since_timestamp}")
-
- processed_files_set = set()
-
- for commit in commits:
- commit_sha = commit["sha"]
- commit_message = commit["commit"]["message"]
-
- self.logger.debug(f"Processing commit {commit_sha[:8]}: {commit_message}")
-
- changed_files = await self._get_commit_files(repo_name, commit_sha)
-
- for file_info in changed_files:
- file_path = file_info["filename"]
-
- if file_path in processed_files_set:
- continue
-
- processed_files_set.add(file_path)
-
- if file_info["status"] == "removed":
- self.logger.debug(f"Processing deleted file: {file_path}")
- deletion_entity = GitHubFileDeletionEntity(
- breadcrumbs=[],
- full_path=f"{repo_name}/{file_path}",
- deletion_label=f"Deleted file {file_path}",
- file_path=file_path,
- repo_name=repo,
- repo_owner=owner,
- branch=branch,
- deletion_status="removed",
- )
- yield deletion_entity
- continue
-
- try:
- async for entity in self._process_changed_file(
- repo_name, file_path, owner, repo, branch, repo_breadcrumb, files=files
- ):
- yield entity
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing changed file {file_path}: {e}")
-
- async def _get_commits_since(
- self, repo_name: str, since_timestamp: str, branch: str
- ) -> List[Dict[str, Any]]:
- """Get commits since a specific timestamp."""
- url = f"{self.BASE_URL}/repos/{repo_name}/commits"
- params = {
- "since": since_timestamp,
- "sha": branch,
- "per_page": 100,
- }
-
- commits = await self._get_paginated_results(url, params)
- self.logger.debug(f"Retrieved {len(commits)} commits since {since_timestamp}")
- return commits
+ @staticmethod
+ def _page(
+ records: list[CaptureRecord],
+ progress: RootProgress | ListProgress | FileProgress,
+ *,
+ final: bool,
+ ) -> CapturePage:
+ return CapturePage(
+ records=tuple(records),
+ final=final,
+ continuation=ScanContinuation(value=progress.model_dump(mode="json")),
+ )
- async def _get_commit_files(self, repo_name: str, commit_sha: str) -> List[Dict[str, Any]]:
- """Get files changed in a specific commit."""
- url = f"{self.BASE_URL}/repos/{repo_name}/commits/{commit_sha}"
+ @staticmethod
+ def _context(repository: Repository) -> RepositoryContext:
+ return RepositoryContext(
+ repository_id=repository.id,
+ owner_id=repository.owner.id,
+ full_name=repository.full_name,
+ )
+ async def _repository(
+ self, selection: GitHubRepositorySelection
+ ) -> dict[str, JsonValue] | None:
try:
- commit_data = await self._get(url)
- files_list = commit_data.get("files", [])
- self.logger.debug(f"Commit {commit_sha[:8]} changed {len(files_list)} files")
- return files_list
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error getting files for commit {commit_sha}: {e}")
- return []
-
- async def _process_changed_file(
+ payload = await self.reader.object("/repos/" + selection.full_name, redirects=True)
+ except GitHubUnavailable:
+ return None
+ repository = Repository.model_validate(payload)
+ # Numeric owner selection is an authorization constraint, not repository identity.
+ if repository.id != selection.repository_id or repository.owner.id != selection.owner_id:
+ return None
+ return payload
+
+ async def capture_page(
self,
- repo_name: str,
- file_path: str,
- owner: str,
- repo: str,
- branch: str,
- repo_breadcrumb: Breadcrumb,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
*,
- files: FileService | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Process a single changed file."""
- path_parts = file_path.split("/")
- breadcrumbs = [repo_breadcrumb]
-
- current_path = ""
- for _i, part in enumerate(path_parts[:-1]):
- current_path = f"{current_path}/{part}" if current_path else part
- dir_breadcrumb = Breadcrumb(
- entity_id=f"{repo_name}/{current_path}",
- name=part,
- entity_type=GitHubDirectoryEntity.__name__,
- )
- breadcrumbs.append(dir_breadcrumb)
-
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Provider I/O happens outside SQL; the driver commits returned progress atomically."""
+ if scope.record_type == "repository":
+ if parent is not None or scope.container_id is not None:
+ raise ValueError("GitHub repository inventory must be a root scope")
+ return await self._repositories(scope, continuation)
+ if parent is None or scope != self.child_scope(parent, scope.record_type):
+ raise ValueError("GitHub page requires its exact captured parent")
+ file_progress = None
+ if scope.record_type == "file":
+ file_progress = FileProgress.model_validate(continuation.value)
+ branch = Branch.model_validate(BranchPayload.model_validate(parent.payload).branch)
+ if file_progress.initialized and (
+ file_progress.ref != branch.name or file_progress.commit_sha != branch.commit.sha
+ ):
+ raise InvalidScanContinuation(
+ "GitHub selected branch changed since the pinned scan"
+ )
try:
- file_url = f"{self.BASE_URL}/repos/{repo_name}/contents/{file_path}?ref={branch}"
- file_data = await self._get(file_url)
-
- file_size = file_data.get("size", 0)
- if file_size > self._max_file_size:
- self.logger.debug(f"Skipping large file: {file_path} ({file_size} bytes)")
- return
-
- content_url = file_data["download_url"]
- content_response = await self.http_client.get(content_url)
- content_response.raise_for_status()
- content_text = content_response.text
-
- if not self._is_text_content(content_text):
- self.logger.debug(f"Skipping binary file: {file_path}")
- return
-
- language = self._detect_language_from_extension(file_path)
- line_count = content_text.count("\n") + 1
- mime_type = mimetypes.guess_type(file_path)[0] or "text/plain"
- file_type = mime_type.split("/")[0] if "/" in mime_type else "file"
-
- file_entity = GitHubCodeFileEntity(
- breadcrumbs=breadcrumbs,
- full_path=f"{repo_name}/{file_path}",
- name=path_parts[-1],
- branch=branch,
- url=file_data.get("download_url") or file_data.get("html_url", ""),
- size=file_size,
- file_type=file_type,
- mime_type=mime_type,
- local_path=None,
- repo_name=repo,
- path_in_repo=file_path,
- repo_owner=owner,
- language=language,
- commit_id=file_data["sha"],
- html_url=file_data.get("html_url"),
- sha=file_data["sha"],
- line_count=line_count,
- is_binary=False,
- )
-
- if files:
- await files.save_bytes(
- entity=file_entity,
- content=content_text.encode("utf-8"),
- filename_with_extension=file_path,
- logger=self.logger,
+ await self._verify_parent_repository(parent)
+ if scope.record_type == "branch":
+ return await self._branches(scope, parent)
+ if scope.record_type == "file":
+ assert file_progress is not None
+ return await self._files(scope, file_progress, parent, files)
+ if scope.record_type == "issue":
+ return await self._issues(scope, continuation, parent)
+ return await self._conversation(scope, continuation, parent)
+ except GitHubUnavailable as error:
+ raise ScopeAccessLost(
+ "GitHub parent is no longer readable", removal_reason="scope_removed"
+ ) from error
+
+ async def _verify_parent_repository(self, parent: SourceRecord) -> None:
+ """A captured path is routing context, never continuing authorization after rename/reuse."""
+ if parent.identity.record_type == "repository":
+ context = self._context(Repository.model_validate(parent.payload))
+ elif parent.identity.record_type == "branch":
+ context = BranchPayload.model_validate(parent.payload).repository
+ else:
+ context = IssuePayload.model_validate(parent.payload).repository
+ current = Repository.model_validate(await self.reader.object("/repos/" + context.full_name))
+ if current.id != context.repository_id or current.owner.id != context.owner_id:
+ raise GitHubUnavailable("GitHub repository no longer matches the authorized selection")
+ if parent.identity.record_type == "issue":
+ issue = Issue.model_validate(IssuePayload.model_validate(parent.payload).issue)
+ current_issue = Issue.model_validate(
+ await self.reader.object(
+ f"/repos/{context.full_name}/issues/{issue.number}", redirects=True
)
-
- if not file_entity.local_path:
- raise ValueError(f"Save failed - no local path set for {file_entity.name}")
-
- yield file_entity
-
- except FileSkippedException as e:
- self.logger.debug(f"Skipping file: {e.reason}")
-
- except SourceAuthError:
- raise
-
- except Exception as e:
- self.logger.warning(f"Error processing changed file {file_path}: {e}")
-
- def _is_text_content(self, content: str) -> bool:
- """Check if content appears to be text."""
- if "\x00" in content:
- return False
- printable_ratio = sum(1 for c in content[:1000] if c.isprintable() or c.isspace()) / min(
- len(content), 1000
+ )
+ if current_issue.id != issue.id or current_issue.repository_url != (
+ f"https://api.github.com/repos/{context.full_name}"
+ ):
+ raise GitHubUnavailable("GitHub issue no longer matches the captured parent")
+
+ async def _repositories(
+ self, scope: CompletedScope, continuation: ScanContinuation
+ ) -> CapturePage:
+ progress = RootProgress.model_validate(continuation.value)
+ if progress.repository_index > len(self.config.repositories):
+ raise InvalidScanContinuation("GitHub repository selection changed")
+ records = []
+ if progress.repository_index < len(self.config.repositories):
+ selection = self.config.repositories[progress.repository_index]
+ payload = await self._repository(selection)
+ if payload is not None:
+ records.append(self._record(scope, str(selection.repository_id), payload))
+ progress.repository_index += 1
+ return self._page(
+ records, progress, final=progress.repository_index == len(self.config.repositories)
)
- return printable_ratio > 0.7
-
- async def _traverse_directory( # noqa: C901
- self,
- repo_name: str,
- path: str,
- breadcrumbs: List[Breadcrumb],
- owner: str,
- repo: str,
- branch: str,
- processed_paths: set,
- processed_files: set,
- *,
- files: FileService | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Recursively traverse a directory using DFS."""
- if path in processed_paths:
- return
-
- processed_paths.add(path)
-
- url = f"{self.BASE_URL}/repos/{repo_name}/contents/{path}"
- params = {"ref": branch}
+ async def _branches(self, scope: CompletedScope, parent: SourceRecord) -> CapturePage:
+ repository = Repository.model_validate(parent.payload)
+ selection = self.selections[repository.id]
+ if not self.config.include_code:
+ return self._page([], RootProgress(), final=True)
+ ref = selection.ref or repository.default_branch
try:
- contents = await self._get(url, params)
+ payload = await self.reader.object(
+ f"/repos/{repository.full_name}/branches/{quote(ref, safe='')}"
+ )
+ except GitHubUnavailable:
+ if await self._repository(selection) is None:
+ raise
+ # A missing branch is not proof that repository issues also became private.
+ return self._page([], RootProgress(), final=True)
+ branch = Branch.model_validate(payload)
+ if branch.name != ref:
+ raise ValueError("GitHub returned a different selected branch")
+ body = BranchPayload(branch=payload, repository=self._context(repository))
+ return self._page(
+ [self._record(scope, branch.name, body.model_dump(mode="json"))],
+ RootProgress(),
+ final=True,
+ )
- if isinstance(contents, List):
- items = contents
+ async def _issues(
+ self, scope: CompletedScope, continuation: ScanContinuation, parent: SourceRecord
+ ) -> CapturePage:
+ progress = ListProgress.model_validate(continuation.value)
+ if not self.config.include_conversations:
+ return self._page([], progress, final=True)
+ repository = Repository.model_validate(parent.payload)
+ path = f"/repos/{repository.full_name}"
+ payloads, final = await self.reader.page(
+ path + "/issues", progress.page, extra="state=all&sort=created&direction=asc"
+ )
+ records = []
+ for payload in payloads:
+ issue = Issue.model_validate(payload)
+ if issue.repository_url != "https://api.github.com" + path:
+ raise ValueError("GitHub issue belongs to a different repository")
+ detail = None
+ if issue.pull_request is not None:
+ detail = await self.reader.object(f"{path}/pulls/{issue.number}")
+ if PullRequest.model_validate(detail).number != issue.number:
+ raise ValueError("GitHub returned a different pull request")
+ body = IssuePayload(
+ issue=payload, repository=self._context(repository), pull_request_detail=detail
+ )
+ records.append(self._record(scope, str(issue.id), body.model_dump(mode="json")))
+ progress.page += 1
+ progress.count += len(records)
+ return self._page(records, progress, final=final)
+
+ async def _conversation(
+ self, scope: CompletedScope, continuation: ScanContinuation, parent: SourceRecord
+ ) -> CapturePage:
+ progress = ListProgress.model_validate(continuation.value)
+ body = IssuePayload.model_validate(parent.payload)
+ issue = Issue.model_validate(body.issue)
+ base = f"/repos/{body.repository.full_name}"
+ if scope.record_type != "comment" and issue.pull_request is None:
+ return self._page([], progress, final=True)
+ suffix = {
+ "comment": f"issues/{issue.number}/comments",
+ "review": f"pulls/{issue.number}/reviews",
+ "review_comment": f"pulls/{issue.number}/comments",
+ "pull_file": f"pulls/{issue.number}/files",
+ }[scope.record_type]
+ before = None
+ if scope.record_type == "pull_file":
+ before = PullRequest.model_validate(
+ await self.reader.object(f"{base}/pulls/{issue.number}")
+ )
+ if progress.head_sha is not None and (
+ progress.head_sha != before.head.sha or progress.base_sha != before.base.sha
+ ):
+ raise InvalidScanContinuation("GitHub pull request changed while listing files")
+ progress.head_sha, progress.base_sha = before.head.sha, before.base.sha
+ payloads, final = await self.reader.page(f"{base}/{suffix}", progress.page)
+ if before is not None:
+ await self._verify_pull_files(
+ base, issue.number, before, progress.count + len(payloads), final
+ )
+ records = []
+ for payload in payloads:
+ if scope.record_type == "pull_file":
+ native_id = payload.get("filename")
+ if not isinstance(native_id, str) or not native_id:
+ raise ValueError("GitHub pull request file has no native path")
else:
- items = [contents]
-
- for item in items:
- item_path = item["path"]
- item_type = item["type"]
-
- if item_type == "dir":
- dir_entity = GitHubDirectoryEntity(
- breadcrumbs=breadcrumbs.copy(),
- full_path=f"{repo_name}/{item_path}",
- name=Path(item_path).name,
- path=item_path,
- repo_name=repo,
- repo_owner=owner,
- branch=branch,
- )
-
- dir_breadcrumb = Breadcrumb(
- entity_id=dir_entity.full_path,
- name=dir_entity.name,
- entity_type=GitHubDirectoryEntity.__name__,
- )
-
- yield dir_entity
-
- dir_breadcrumbs = breadcrumbs.copy() + [dir_breadcrumb]
-
- async for child_entity in self._traverse_directory(
- repo_name,
- item_path,
- dir_breadcrumbs,
- owner,
- repo,
- branch,
- processed_paths,
- processed_files,
- files=files,
- ):
- yield child_entity
-
- elif item_type == "file":
- if item_path in processed_files:
- continue
-
- processed_files.add(item_path)
-
- async for file_entity in self._process_file(
- repo_name, item_path, item, breadcrumbs, owner, repo, branch, files=files
- ):
- yield file_entity
-
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error traversing path {path}: {str(e)}")
-
- async def _process_file( # noqa: C901
+ native_id = str(Identity.model_validate(payload).id)
+ captured = {
+ "native": payload,
+ "repository": body.repository.model_dump(mode="json"),
+ "issue_number": issue.number,
+ }
+ if before is not None:
+ captured.update({"head_sha": before.head.sha, "base_sha": before.base.sha})
+ record = self._record(scope, native_id, captured)
+ # PR file API is metadata/patch coverage; original bytes live in selected branch files.
+ if scope.record_type == "pull_file":
+ record = record.model_copy(update={"completeness": "metadata_only"})
+ records.append(record)
+ progress.page += 1
+ progress.count += len(records)
+ return self._page(records, progress, final=final)
+
+ async def _verify_pull_files(
+ self, base: str, number: int, before: PullRequest, count: int, final: bool
+ ) -> None:
+ after = PullRequest.model_validate(await self.reader.object(f"{base}/pulls/{number}"))
+ if (after.head.sha, after.base.sha) != (before.head.sha, before.base.sha):
+ raise InvalidScanContinuation("GitHub pull request changed while listing files")
+ if count >= 3000 and (not final or before.changed_files > count):
+ raise ValueError("GitHub pull request file listing exceeds its complete API window")
+ if final and count != before.changed_files:
+ raise ValueError("GitHub pull request file count was not completely enumerated")
+
+ async def _files(
self,
- repo_name: str,
- item_path: str,
- item: Dict[str, Any],
- breadcrumbs: List[Breadcrumb],
- owner: str,
- repo: str,
- branch: str,
- *,
- files: FileService | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Process a file item and create file entities."""
+ scope: CompletedScope,
+ progress: FileProgress,
+ parent: SourceRecord,
+ files: FileService,
+ ) -> CapturePage:
+ body = BranchPayload.model_validate(parent.payload)
+ branch = Branch.model_validate(body.branch)
+ base = f"/repos/{body.repository.full_name}/git"
try:
- file_url = f"{self.BASE_URL}/repos/{repo_name}/contents/{item_path}"
- file_data = await self._get(file_url, {"ref": branch})
- file_size = file_data.get("size", 0)
-
- content_sample = None
- content_text = None
- if file_data.get("encoding") == "base64" and file_data.get("content"):
- try:
- content_sample = base64.b64decode(file_data["content"])
- content_text = content_sample.decode("utf-8", errors="replace")
- except Exception:
- pass
-
- if is_text_file(item_path, file_size, content_sample):
- language = self._detect_language_from_extension(item_path)
- file_name = Path(item_path).name
-
- line_count = 0
- if content_text:
- try:
- line_count = content_text.count("\n") + 1
- except Exception as e:
- self.logger.warning(f"Error counting lines for {item_path}: {str(e)}")
-
- mime_type = mimetypes.guess_type(item_path)[0] or "text/plain"
- file_type = mime_type.split("/")[0] if "/" in mime_type else "file"
-
- file_entity = GitHubCodeFileEntity(
- breadcrumbs=breadcrumbs.copy(),
- full_path=f"{repo_name}/{item_path}",
- name=file_name,
- branch=branch,
- url=file_data.get("download_url") or file_data["html_url"],
- size=file_size,
- file_type=file_type,
- mime_type=mime_type,
- local_path=None,
- repo_name=repo,
- path_in_repo=item_path,
- repo_owner=owner,
- language=language,
- commit_id=file_data["sha"],
- html_url=file_data.get("html_url"),
- sha=file_data["sha"],
- line_count=line_count,
- is_binary=False,
- )
-
- if files and content_text:
- await files.save_bytes(
- entity=file_entity,
- content=content_text.encode("utf-8"),
- filename_with_extension=item_path,
- logger=self.logger,
+ progress = await self._pin_code(progress, branch, base)
+ records = []
+ stack = list(progress.stack)
+ # One tree response per page bounds work for enormous empty directory trees.
+ if stack:
+ frame = stack[-1]
+ tree = Tree.model_validate(await self.reader.object(f"{base}/trees/{frame.sha}"))
+ if tree.sha != frame.sha or tree.truncated or frame.offset > len(tree.tree):
+ raise ValueError("GitHub tree cannot be completely and consistently enumerated")
+ while frame.offset < len(tree.tree) and len(records) < 1:
+ raw = tree.tree[frame.offset]
+ entry = TreeEntry.model_validate(raw)
+ if entry.path in {".", ".."} or "/" in entry.path or "\x00" in entry.path:
+ raise ValueError("GitHub tree entry is not a single native path component")
+ frame.offset += 1
+ if entry.type == "tree":
+ stack.append(TreeFrame(sha=entry.sha, component=entry.path))
+ break
+ path = "/".join(
+ [item.component for item in stack if item.component] + [entry.path]
)
-
- if not file_entity.local_path:
- raise ValueError(f"Save failed - no local path set for {file_entity.name}")
-
- yield file_entity
- except FileSkippedException as e:
- self.logger.debug(f"Skipping file: {e.reason}")
-
- except SourceAuthError:
- raise
-
- except Exception as e:
- self.logger.warning(f"Error processing file {item_path}: {str(e)}")
-
- def _get_cursor_timestamp(self, cursor: SyncCursor | None) -> str:
- """Get last repository pushed timestamp from cursor."""
- cursor_data = cursor.data if cursor else {}
- last_pushed_at = cursor_data.get("last_repository_pushed_at")
-
- if last_pushed_at:
- self.logger.debug(f"Incremental sync from cursor: {last_pushed_at}")
- else:
- self.logger.debug("Full sync (no cursor)")
-
- return last_pushed_at
-
- async def _verify_branch(self, branch: str) -> None:
- """Verify that the specified branch exists in the repository."""
- if self._branch:
- branches_url = f"{self.BASE_URL}/repos/{self._repo_name}/branches"
- branches_data = await self._get_paginated_results(branches_url)
- branch_names = [b["name"] for b in branches_data]
-
- if branch not in branch_names:
- available_branches = ", ".join(branch_names)
- raise ValueError(
- f"Branch '{branch}' not found in repository '{self._repo_name}'. "
- f"Available branches: {available_branches}"
- )
-
- async def _fetch_merged_pull_requests(
- self,
- repo_name: str,
- repo_breadcrumb: Breadcrumb,
- since: Optional[str] = None,
- cursor: SyncCursor | None = None,
- ) -> AsyncGenerator[BaseEntity, None]:
- """Fetch merged pull requests and their review comments."""
- owner, repo = repo_name.split("/")
- url = f"{self.BASE_URL}/repos/{repo_name}/pulls"
- params = {"state": "closed", "sort": "updated", "direction": "desc"}
-
- all_prs = await self._get_paginated_results(url, params)
-
- latest_updated_at = since or ""
-
- for pr_data in all_prs:
- if not pr_data.get("merged_at"):
- continue
-
- pr_updated = pr_data["updated_at"]
-
- if since and pr_updated <= since:
- continue
-
- if pr_updated > latest_updated_at:
- latest_updated_at = pr_updated
-
- pr_number = pr_data["number"]
-
- pr_breadcrumb = Breadcrumb(
- entity_id=f"{repo_name}#{pr_number}",
- name=f"PR #{pr_number}",
- entity_type=GitHubPullRequestEntity.__name__,
+ payload = {
+ "entry": raw,
+ "path": path,
+ "ref": progress.ref,
+ "commit_sha": progress.commit_sha,
+ "tree_sha": progress.tree_sha,
+ "repository": body.repository.model_dump(mode="json"),
+ }
+ record = self._record(scope, path, payload)
+ record = await self._retain_blob(record, entry, base, path, files)
+ records.append(record)
+ if stack[-1] is frame and frame.offset == len(tree.tree):
+ stack.pop()
+ progress.stack = tuple(stack)
+ return self._page(records, progress, final=not stack)
+ except GitHubUnavailable as error:
+ # An unreadable branch withdraws this code scope. A still-readable branch
+ # with a missing immutable snapshot requires restart, never absence cleanup.
+ await self.reader.object(
+ f"/repos/{body.repository.full_name}/branches/{quote(branch.name, safe='')}"
)
+ raise InvalidScanContinuation("GitHub pinned code snapshot is unavailable") from error
- files_data = await self._get_paginated_results(
- f"{self.BASE_URL}/repos/{repo_name}/pulls/{pr_number}/files",
- {},
- )
- changed_paths = [f["filename"] for f in files_data if f.get("filename")]
+ async def _pin_code(self, progress: FileProgress, branch: Branch, base: str) -> FileProgress:
+ if progress.initialized:
+ return progress
+ commit = GitCommit.model_validate(
+ await self.reader.object(f"{base}/commits/{branch.commit.sha}")
+ )
+ if commit.sha != branch.commit.sha:
+ raise ValueError("GitHub returned a different pinned commit")
+ return FileProgress(
+ ref=branch.name,
+ commit_sha=commit.sha,
+ tree_sha=commit.tree.sha,
+ stack=(TreeFrame(sha=commit.tree.sha, component=""),),
+ initialized=True,
+ )
- pr_entity = GitHubPullRequestEntity.from_api(
- pr_data,
- repo_name=repo,
- repo_owner=owner,
- changed_files_list=changed_paths or None,
- breadcrumbs=[repo_breadcrumb],
- )
- yield pr_entity
+ async def _retain_blob(
+ self, record: CaptureRecord, entry: TreeEntry, base: str, path: str, files: FileService
+ ) -> CaptureRecord:
+ if entry.type == "commit":
+ return record.model_copy(update={"completeness": "metadata_only"})
+ if entry.size is None or entry.size > MAX_CODE_BYTES:
+ return record.model_copy(update={"completeness": "partial"})
+ content = await self.reader.blob(
+ f"{base}/blobs/{entry.sha}", max_bytes=min(entry.size, MAX_CODE_BYTES)
+ )
+ if len(content) != entry.size or len(content) > MAX_CODE_BYTES:
+ raise ValueError("GitHub blob size differs from its immutable tree entry")
+ digest = hashlib.sha1(f"blob {len(content)}\0".encode() + content).hexdigest()
+ if digest != entry.sha:
+ raise ValueError("GitHub blob does not match its immutable object identity")
+ blob = await files.store_canonical_blob(content, media_type=mimetypes.guess_type(path)[0])
+ lfs_pointer = content.startswith(b"version https://git-lfs.github.com/spec/v1\n")
+ return record.model_copy(
+ update={"blobs": (blob,), "completeness": "partial" if lfs_pointer else "complete"}
+ )
- async for comment_entity in self._fetch_pr_review_comments(
- repo_name, pr_number, owner, repo, [repo_breadcrumb, pr_breadcrumb]
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Confirm omitted inventory owners before descendants lose visibility."""
+ kind = record.identity.record_type
+ if kind == "repository":
+ selection = self.selections.get(int(record.identity.native_id))
+ if selection is None or await self._repository(selection) is None:
+ return
+ elif kind == "branch":
+ if await self._branch_absent(record):
+ return
+ elif kind == "issue":
+ if not self.config.include_conversations:
+ return
+ body = IssuePayload.model_validate(record.payload)
+ issue = Issue.model_validate(body.issue)
+ try:
+ current = Issue.model_validate(
+ await self.reader.object(
+ f"/repos/{body.repository.full_name}/issues/{issue.number}", redirects=True
+ )
+ )
+ except GitHubUnavailable:
+ return
+ if current.id != issue.id or current.repository_url != (
+ f"https://api.github.com/repos/{body.repository.full_name}"
):
- yield comment_entity
-
- if cursor and latest_updated_at and latest_updated_at != since:
- cursor.update(last_pr_updated_at=latest_updated_at)
+ return
+ else:
+ raise ValueError("GitHub absence confirmation is only valid for inventory owners")
+ raise ValueError("GitHub omitted an inventory record that remains readable in selection")
- async def _fetch_pr_review_comments(
- self,
- repo_name: str,
- pr_number: int,
- owner: str,
- repo: str,
- breadcrumbs: List[Breadcrumb],
- ) -> AsyncGenerator[BaseEntity, None]:
- """Fetch review comments (inline code comments) for a single PR."""
- url = f"{self.BASE_URL}/repos/{repo_name}/pulls/{pr_number}/comments"
+ async def _branch_absent(self, record: SourceRecord) -> bool:
+ body = BranchPayload.model_validate(record.payload)
+ selection = self.selections.get(body.repository.repository_id)
+ if selection is None or not self.config.include_code:
+ return True
+ repository = await self._repository(selection)
+ if repository is None:
+ return True
+ ref = selection.ref or Repository.model_validate(repository).default_branch
+ if ref != record.identity.native_id:
+ return True
try:
- comments = await self._get_paginated_results(url)
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error fetching review comments for PR #{pr_number}: {e}")
- return
-
- for comment in comments:
- yield GitHubPRCommentEntity.from_api(
- comment,
- repo_name=repo,
- repo_owner=owner,
- pr_number=pr_number,
- breadcrumbs=breadcrumbs.copy(),
+ await self.reader.object(
+ f"/repos/{body.repository.full_name}/branches/{quote(ref, safe='')}"
)
+ except GitHubUnavailable:
+ return True
+ return False
async def generate_entities(
self,
@@ -758,79 +553,6 @@ async def generate_entities(
files: FileService | None = None,
node_selections: list[NodeSelectionData] | None = None,
) -> AsyncGenerator[BaseEntity, None]:
- """Generate entities from GitHub repository with incremental support.
-
- Yields:
- Repository, directory, file, pull request, and PR comment entities
- """
- if not self._repo_name:
- raise ValueError("Repository name must be specified")
-
- last_pushed_at = self._get_cursor_timestamp(cursor)
-
- repo_url = f"{self.BASE_URL}/repos/{self._repo_name}"
- repo_data = await self._get(repo_url)
- current_pushed_at = repo_data["pushed_at"]
-
- branch = self._branch if self._branch else repo_data["default_branch"]
-
- await self._verify_branch(branch)
-
- should_sync = True
- if last_pushed_at and current_pushed_at <= last_pushed_at:
- self.logger.debug(
- f"Repository {self._repo_name} has no new commits since last sync, "
- "skipping file traversal"
- )
- should_sync = False
-
- if should_sync:
- self.logger.debug(f"Using branch: {branch} for repo {self._repo_name}")
-
- if last_pushed_at:
- self.logger.debug(f"Performing INCREMENTAL sync - changes since {last_pushed_at}")
- async for entity in self._traverse_repository_incremental(
- self._repo_name, branch, last_pushed_at, cursor=cursor, files=files
- ):
- yield entity
- else:
- self.logger.debug("Performing FULL sync - no previous cursor data")
- async for entity in self._traverse_repository(
- self._repo_name, branch, cursor=cursor, files=files
- ):
- yield entity
- else:
- repo_entity = await self._get_repository_info(self._repo_name, cursor=cursor)
- yield repo_entity
-
- if self._sync_pull_requests:
- cursor_data = cursor.data if cursor else {}
- last_pr_updated = cursor_data.get("last_pr_updated_at") or None
-
- repo_breadcrumb = Breadcrumb(
- entity_id=str(repo_data["id"]),
- name=repo_data["name"],
- entity_type=GitHubRepositoryEntity.__name__,
- )
-
- if last_pr_updated:
- self.logger.debug(f"Incremental PR sync - changes since {last_pr_updated}")
- else:
- self.logger.debug("Full PR sync - no previous cursor data")
-
- async for entity in self._fetch_merged_pull_requests(
- self._repo_name, repo_breadcrumb, since=last_pr_updated, cursor=cursor
- ):
- yield entity
-
- async def validate(self) -> None:
- """Verify GitHub PAT and optional repo/branch access via the same path as sync.
-
- Uses :meth:`_get` (Airweave client + ``raise_for_status``) so failures surface as
- the same domain errors as the rest of the connector. Relies on :meth:`create` for
- ``_personal_access_token``, ``_repo_name``, and ``_branch``.
- """
- if self._repo_name:
- await self._get(f"{self.BASE_URL}/repos/{self._repo_name}")
- if self._branch:
- await self._get(f"{self.BASE_URL}/repos/{self._repo_name}/branches/{self._branch}")
+ """Replace the legacy lossy traversal; the registry chooses canonical capture."""
+ raise NotImplementedError("GitHub requires canonical page capture")
+ yield # pragma: no cover
diff --git a/backend/airweave/platform/sources/github_http.py b/backend/airweave/platform/sources/github_http.py
new file mode 100644
index 0000000000..90df4e86e7
--- /dev/null
+++ b/backend/airweave/platform/sources/github_http.py
@@ -0,0 +1,217 @@
+"""GitHub read transport: bounded trusted redirects and strict native pagination."""
+
+from __future__ import annotations
+
+import time
+from contextlib import aclosing
+from functools import partial
+from urllib.parse import parse_qs, urlsplit
+
+import httpx
+from pydantic import JsonValue, TypeAdapter
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
+
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
+from airweave.domains.sources.exceptions import (
+ SourceAuthError,
+ SourceError,
+ SourceRateLimitError,
+ SourceServerError,
+)
+from airweave.domains.sources.token_providers.credential import DirectCredentialProvider
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
+from airweave.platform.configs.auth import GitHubAuthConfig
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.http_client.retry_helpers import (
+ retry_if_rate_limit_or_timeout,
+ wait_rate_limit_with_backoff,
+)
+
+BASE_URL = "https://api.github.com"
+OBJECT = TypeAdapter(dict[str, JsonValue])
+PAGE = TypeAdapter(list[dict[str, JsonValue]])
+
+
+class GitHubUnavailable(Exception):
+ """A non-rate-limited 403/404; never proof of upstream deletion."""
+
+
+def trusted_url(url: str) -> bool:
+ """Never send managed credentials to provider-supplied arbitrary locations."""
+ parsed = urlsplit(url)
+ return (
+ parsed.scheme == "https"
+ and parsed.netloc == "api.github.com"
+ and not parsed.fragment
+ and not parsed.username
+ and not parsed.password
+ )
+
+
+class GitHubReader:
+ """Existing shared client and auth own network access; this adds GitHub semantics only."""
+
+ def __init__(self, client: AirweaveHttpClient, auth: SourceAuthProvider):
+ """Reuse the source-scoped native transport and credential owner."""
+ self.client = client
+ self.auth = auth
+
+ @retry(
+ stop=stop_after_attempt(5),
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
+ reraise=True,
+ )
+ async def get(self, path: str, *, raw: bool = False, redirects: bool = False) -> httpx.Response:
+ """Read a native endpoint without exposing provider error bodies in logs."""
+ url = BASE_URL + path
+ refreshed = False
+ for _ in range(4):
+ if not trusted_url(url):
+ raise SourceError(
+ "GitHub returned an untrusted redirect", source_short_name="github"
+ )
+ headers = await self._headers(refresh=refreshed, raw=raw)
+ response = await self.client.get(url, headers=headers, follow_redirects=False)
+ if response.status_code == 401 and self.auth.supports_refresh and not refreshed:
+ refreshed = True
+ continue
+ if response.status_code in {301, 302, 307, 308}:
+ if not redirects:
+ raise SourceError(
+ "GitHub captured route changed; refresh its parent inventory",
+ source_short_name="github",
+ )
+ url = response.headers.get("Location", "")
+ continue
+ self._status(response)
+ return response
+ raise SourceError(
+ "GitHub redirect or authentication retry did not resolve", source_short_name="github"
+ )
+
+ async def _headers(self, *, refresh: bool = False, raw: bool = False) -> dict[str, str]:
+ if isinstance(self.auth, DirectCredentialProvider):
+ credentials = GitHubAuthConfig.model_validate(self.auth.credentials)
+ headers = {"Authorization": f"Bearer {credentials.personal_access_token}"}
+ else:
+ headers = await authorization_headers(self.auth, refresh=refresh)
+ headers.update(
+ {
+ "Accept": "application/vnd.github.raw+json"
+ if raw
+ else "application/vnd.github+json",
+ "X-GitHub-Api-Version": "2022-11-28",
+ }
+ )
+ return headers
+
+ @retry(
+ stop=stop_after_attempt(5),
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
+ reraise=True,
+ )
+ async def blob(self, path: str, *, max_bytes: int) -> bytes:
+ """Bound bytes while the existing HTTP transport streams the native blob."""
+ for refresh in (False, True):
+ headers = await self._headers(refresh=refresh, raw=True)
+ headers["Accept-Encoding"] = "identity"
+ async with self.client.stream(
+ "GET", BASE_URL + path, headers=headers, follow_redirects=False
+ ) as response:
+ if response.status_code == 401 and not refresh and self.auth.supports_refresh:
+ continue
+ if response.status_code != 200:
+ error_body = bytearray()
+ async with aclosing(response.aiter_raw()) as chunks:
+ async for chunk in chunks:
+ error_body.extend(chunk[: 65536 - len(error_body)])
+ if len(error_body) >= 65536:
+ break
+ # Status handling sees only a bounded error envelope, never logs its body.
+ buffered = httpx.Response(
+ response.status_code, headers=response.headers, content=bytes(error_body)
+ )
+ self._status(buffered)
+ if response.headers.get("content-encoding", "identity") != "identity":
+ raise ValueError("GitHub ignored the requested identity blob encoding")
+ length = response.headers.get("content-length")
+ if length is not None and int(length) > max_bytes:
+ raise ValueError("GitHub blob exceeds its retained byte bound")
+ content = bytearray()
+ async with aclosing(response.aiter_raw()) as chunks:
+ async for chunk in chunks:
+ if len(content) + len(chunk) > max_bytes:
+ raise ValueError("GitHub blob exceeds its retained byte bound")
+ content.extend(chunk)
+ return bytes(content)
+ raise SourceError(
+ "GitHub authorization refresh did not resolve", source_short_name="github"
+ )
+
+ def _status(self, response: httpx.Response) -> None:
+ status = response.status_code
+ if status in {403, 429}:
+ message = ""
+ if "json" in response.headers.get("content-type", ""):
+ value = OBJECT.validate_python(response.json()).get("message")
+ message = value.lower() if isinstance(value, str) else ""
+ if (
+ status == 429
+ or "retry-after" in response.headers
+ or response.headers.get("x-ratelimit-remaining") == "0"
+ or "rate limit" in message
+ ):
+ try:
+ delay = float(response.headers.get("retry-after", "0"))
+ if delay <= 0:
+ delay = float(response.headers.get("x-ratelimit-reset", "0")) - time.time()
+ except ValueError:
+ delay = 60
+ raise SourceRateLimitError(source_short_name="github", retry_after=max(60, delay))
+ if status == 401:
+ raise SourceAuthError(
+ "GitHub authorization is unavailable",
+ source_short_name="github",
+ status_code=401,
+ token_provider_kind=self.auth.provider_kind,
+ )
+ if status in {403, 404}:
+ raise GitHubUnavailable("GitHub resource is outside the currently readable scope")
+ if status >= 500:
+ raise SourceServerError(
+ "GitHub is temporarily unavailable", source_short_name="github", status_code=status
+ )
+ if status != 200:
+ raise SourceError(f"GitHub read failed with HTTP {status}", source_short_name="github")
+
+ async def object(self, path: str, *, redirects: bool = False) -> dict[str, JsonValue]:
+ """Validate only the JSON boundary, retaining every provider field."""
+ return OBJECT.validate_python((await self.get(path, redirects=redirects)).json())
+
+ async def page(
+ self, path: str, page: int, *, extra: str = ""
+ ) -> tuple[list[dict[str, JsonValue]], bool]:
+ """Follow only the exact next page of this fixed endpoint/query."""
+ query = f"per_page=100&page={page}" + ("&" + extra if extra else "")
+ response = await self.get(path + "?" + query)
+ records = PAGE.validate_python(response.json())
+ if len(records) > 100:
+ raise ValueError("GitHub exceeded its requested page size")
+ following = response.links.get("next", {}).get("url")
+ if following is None:
+ return records, True
+ parsed = urlsplit(following)
+ expected = parse_qs(query)
+ expected["page"] = [str(page + 1)]
+ # Redirected repository names are accepted only by root identity resolution;
+ # mutable child routes cannot quietly switch their captured parent context.
+ if not trusted_url(following) or parsed.path != path or parse_qs(parsed.query) != expected:
+ raise ValueError("GitHub returned a non-advancing or out-of-scope page link")
+ if not records:
+ raise ValueError("GitHub returned an empty page with a continuation")
+ return records, False
diff --git a/backend/airweave/platform/sources/github_models.py b/backend/airweave/platform/sources/github_models.py
new file mode 100644
index 0000000000..192e1a6807
--- /dev/null
+++ b/backend/airweave/platform/sources/github_models.py
@@ -0,0 +1,187 @@
+"""Native GitHub validation and bounded provider progress, not a second capture state owner."""
+
+from typing import Annotated, Literal
+
+from pydantic import (
+ AwareDatetime,
+ BaseModel,
+ ConfigDict,
+ Field,
+ JsonValue,
+ StrictBool,
+ model_validator,
+)
+
+Sha = Annotated[str, Field(pattern=r"^[a-f0-9]{40}$")]
+NativeId = Annotated[int, Field(gt=0, strict=True)]
+
+
+class NativeModel(BaseModel):
+ """Validate routing fields while retaining untouched responses separately."""
+
+ model_config = ConfigDict(extra="ignore")
+
+
+class Identity(NativeModel):
+ """Native numeric object identity."""
+
+ id: NativeId
+
+
+class Repository(Identity):
+ """Selected repository routing and authorization fields."""
+
+ full_name: str = Field(pattern=r"^[a-zA-Z0-9_-]+/[a-zA-Z0-9_.-]+$")
+ owner: Identity
+ default_branch: str
+
+
+class Issue(Identity):
+ """Native issue routing; pull requests also appear in this endpoint."""
+
+ number: int = Field(gt=0, strict=True)
+ repository_url: str
+ pull_request: dict[str, JsonValue] | None = None
+
+
+class CommitSide(NativeModel):
+ """Native immutable Git commit identity."""
+
+ sha: Sha
+
+
+class PullRequest(Identity):
+ """Pinned diff comparison and completeness count."""
+
+ number: int = Field(gt=0, strict=True)
+ head: CommitSide
+ base: CommitSide
+ changed_files: int = Field(ge=0)
+
+
+class TreeEntry(NativeModel):
+ """Native nonrecursive Git tree entry."""
+
+ path: str = Field(min_length=1, max_length=4096)
+ mode: Literal["100644", "100755", "040000", "160000", "120000"]
+ type: Literal["blob", "tree", "commit"]
+ sha: Sha
+ size: int | None = Field(default=None, ge=0)
+
+ @model_validator(mode="after")
+ def coherent_type(self) -> "TreeEntry":
+ """Git mode determines whether this is a blob, subtree, or submodule reference."""
+ expected = (
+ "tree" if self.mode == "040000" else "commit" if self.mode == "160000" else "blob"
+ )
+ if self.type != expected:
+ raise ValueError("GitHub tree entry type does not match its native mode")
+ return self
+
+
+class Tree(NativeModel):
+ """Native Git tree response with explicit truncation signal."""
+
+ sha: Sha
+ tree: list[dict[str, JsonValue]]
+ truncated: StrictBool
+
+
+class GitCommit(NativeModel):
+ """Native commit and its immutable root tree."""
+
+ sha: Sha
+ tree: CommitSide
+
+
+class Branch(NativeModel):
+ """Native selected branch and current commit."""
+
+ name: str = Field(min_length=1, max_length=1024)
+ commit: CommitSide
+
+
+class RepositoryContext(BaseModel):
+ """Source provenance needed to route child reads; never substitutes for native JSON."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ repository_id: NativeId
+ owner_id: NativeId
+ full_name: str = Field(pattern=r"^[a-zA-Z0-9_-]+/[a-zA-Z0-9_.-]+$")
+
+
+class IssuePayload(BaseModel):
+ """Unmodified issue and PR originals plus verified repository provenance."""
+
+ model_config = ConfigDict(extra="forbid")
+ issue: dict[str, JsonValue]
+ repository: RepositoryContext
+ pull_request_detail: dict[str, JsonValue] | None = None
+
+
+class BranchPayload(BaseModel):
+ """Unmodified branch original plus verified repository provenance."""
+
+ model_config = ConfigDict(extra="forbid")
+ branch: dict[str, JsonValue]
+ repository: RepositoryContext
+
+
+class RootProgress(BaseModel):
+ """Position in the fixed authorized repository selection."""
+
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ repository_index: int = Field(default=0, ge=0)
+
+
+class ListProgress(BaseModel):
+ """Native page position and optional pinned PR diff identity."""
+
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ page: int = Field(default=1, ge=1)
+ count: int = Field(default=0, ge=0)
+ head_sha: Sha | None = None
+ base_sha: Sha | None = None
+
+
+class TreeFrame(BaseModel):
+ """A DFS frame holds one path component, avoiding repeated full paths at every depth."""
+
+ model_config = ConfigDict(extra="forbid")
+ sha: Sha
+ component: str = Field(max_length=4096)
+ offset: int = Field(default=0, ge=0)
+
+
+class FileProgress(BaseModel):
+ """Pinned immutable tree plus bounded DFS position, validated again by ScanContinuation."""
+
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ ref: str | None = None
+ commit_sha: Sha | None = None
+ tree_sha: Sha | None = None
+ stack: tuple[TreeFrame, ...] = ()
+ initialized: StrictBool = False
+
+ @model_validator(mode="after")
+ def coherent_snapshot(self) -> "FileProgress":
+ """Malformed progress must fail before provider I/O or absence reconciliation."""
+ pinned = (self.ref, self.commit_sha, self.tree_sha)
+ if self.initialized:
+ if any(value is None for value in pinned):
+ raise ValueError("Initialized GitHub code progress requires a complete snapshot")
+ if self.stack and (self.stack[0].sha != self.tree_sha or self.stack[0].component):
+ raise ValueError("GitHub DFS root does not match the pinned tree")
+ elif any(value is not None for value in pinned) or self.stack:
+ raise ValueError("Uninitialized GitHub code progress cannot contain a snapshot")
+ return self
+
+
+class Dates(NativeModel):
+ """Optional upstream timestamps, never substituted with local observation time."""
+
+ created_at: AwareDatetime | None = None
+ updated_at: AwareDatetime | None = None
diff --git a/backend/airweave/platform/sources/gitlab.py b/backend/airweave/platform/sources/gitlab.py
index 9620c8a828..dac0953bad 100644
--- a/backend/airweave/platform/sources/gitlab.py
+++ b/backend/airweave/platform/sources/gitlab.py
@@ -31,12 +31,12 @@
_require_gl_datetime,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.platform.utils.file_extensions import (
get_language_for_extension,
is_text_file,
diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py
index 57819cc3bc..8acb1765ff 100644
--- a/backend/airweave/platform/sources/gmail.py
+++ b/backend/airweave/platform/sources/gmail.py
@@ -1,27 +1,48 @@
-"""Gmail source implementation for syncing email threads, messages, and attachments.
+"""Gmail originals through the canonical page engine and native history checkpoints.
-Uses concurrent/batching processing for optimal performance:
- * Thread detail fetch + per-thread processing
- * Per-thread message processing
- * Per-message attachment fetch & processing
- * Incremental history message-detail fetch
+The upstream entity extractor remains available for legacy consumers; owned capture
+uses one fenced message/blob authority and independent derived projection.
"""
from __future__ import annotations
import asyncio
import base64
+import hashlib
+import json
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional, Set
import httpx
-from tenacity import retry, stop_after_attempt
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from tenacity import retry, retry_if_exception, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import (
+ AuthProviderRateLimitError,
+ AuthProviderServerError,
+)
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import SourceAuthError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidScanContinuation,
+)
+from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import (
+ SourceAuthError,
+ SourceError,
+ SourceRateLimitError,
+ SourceServerError,
+)
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage import FileSkippedException
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
@@ -36,28 +57,62 @@
GmailThreadEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.bounded_response import bounded_response_bytes
+from airweave.platform.http_client.retry_helpers import (
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.gmail_capture import GmailCapture
+from airweave.platform.sources.gmail_errors import GmailThrottleError, raise_gmail_throttle
+from airweave.platform.sources.gmail_pages import GmailPages, invalid_page_token
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.platform.utils.filename_utils import safe_filename
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
def _should_retry_gmail_request(exception: Exception) -> bool:
"""Custom retry condition that excludes 404 errors but includes 429 and timeouts."""
+ if isinstance(exception, GmailThrottleError):
+ return exception.retry_after is None or exception.retry_after <= 120
+ if isinstance(exception, (SourceRateLimitError, AuthProviderRateLimitError)):
+ # The shared wait caps at 120s. Defer a longer provider minimum instead
+ # of retrying before it or extending this request's bounded wait.
+ return exception.retry_after <= 120
if isinstance(exception, httpx.HTTPStatusError):
- if exception.response.status_code == 404:
- return False
if exception.response.status_code == 429:
- return True
- return True
- if isinstance(exception, (httpx.ConnectTimeout, httpx.ReadTimeout)):
+ try:
+ return float(exception.response.headers.get("Retry-After", "0")) <= 120
+ except ValueError:
+ return True # Preserve the shared helper's fallback for invalid headers.
+ return exception.response.status_code >= 500
+ if isinstance(
+ exception,
+ (
+ httpx.ConnectTimeout,
+ httpx.ReadTimeout,
+ SourceServerError,
+ AuthProviderServerError,
+ ),
+ ):
return True
return False
+def _wait_gmail_request(retry_state) -> float:
+ """Honor native quota timing; otherwise reuse the existing bounded backoff."""
+ error = retry_state.outcome.exception()
+ if isinstance(error, GmailThrottleError) and error.retry_after is not None:
+ return max(1.0, error.retry_after)
+ return wait_rate_limit_with_backoff(retry_state)
+
+
+class GmailProfile(BaseModel):
+ """Native mailbox identity; unrelated profile fields remain provider-owned."""
+
+ model_config = ConfigDict(extra="ignore")
+ emailAddress: str = Field(pattern=r"^[^\s@]+@[^\s@]+$")
+
+
@source(
name="Gmail",
short_name="gmail",
@@ -83,6 +138,90 @@ class GmailSource(BaseSource):
It supports syncing email threads, individual messages, and file attachments.
"""
+ canonical_record_types = ("message",)
+
+ canonical_container_parents = {}
+ capture_config: GmailConfig
+ _verified_mailbox: str | None = None
+
+ def _require_mailbox(self) -> None:
+ """Direct capture calls cannot bypass native identity verification."""
+ expected = self.capture_config.expected_mailbox
+ if expected is None or self._verified_mailbox != expected.casefold():
+ raise ValueError("Owned Gmail capture requires an attested mailbox identity")
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Bind the actual query and typed filters; preserve every existing filtered scope."""
+ self._require_mailbox()
+ query = self._build_gmail_query()
+ fingerprint = hashlib.sha256(
+ json.dumps(
+ {
+ "capture_version": 2,
+ "query": query,
+ "filters": self.capture_config.model_dump(mode="json"),
+ },
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+ return CycleConfiguration(
+ fingerprint=fingerprint,
+ parents={"message": (None,)},
+ known_object_validation=() if query else ("message",),
+ )
+
+ def _canonical_capture(self, files: FileService | None = None) -> GmailCapture:
+ self._require_mailbox()
+ return GmailCapture(
+ self._get_capture_json,
+ self._build_gmail_query(),
+ files=files,
+ attachment_get=self._get_mime_json,
+ )
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Plan full mailbox capture or changes from verified canonical progress."""
+ return await GmailPages(self._canonical_capture()).prepare(
+ previous, self.capture_cycle_configuration
+ )
+
+ def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation:
+ """Resume the cycle's original mailbox boundary, never a fresh profile value."""
+ self._require_mailbox()
+ return GmailPages.initial(cycle)
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Native messages enter the same page transaction as their continuation."""
+ if scope != CompletedScope(record_type="message") or parent is not None:
+ raise ValueError("Gmail requires its independent message scope")
+ return await GmailPages(self._canonical_capture(files)).page(continuation)
+
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Exact current GET proves unfiltered mailbox membership, including known omissions."""
+ if (
+ self._build_gmail_query()
+ or record.identity.record_type != "message"
+ or record.parent is not None
+ ):
+ raise ValueError("Exact Gmail omission refresh requires unfiltered mailbox messages")
+ return await self._canonical_capture(files).message(record.identity.native_id)
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Mail messages have no canonical child scopes."""
+ raise ValueError("Gmail has no child capture scopes")
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Unfiltered omissions require exact hydration; filtered absence describes only scope."""
+ raise ValueError("Gmail omission must use declared exact known-object validation")
+
# -----------------------
# Construction / Config
# -----------------------
@@ -90,7 +229,7 @@ class GmailSource(BaseSource):
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: GmailConfig,
@@ -98,6 +237,7 @@ async def create(
"""Create a new Gmail source instance."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.capture_config = config
config_dict = config.model_dump() if config else {}
instance.batch_size = int(config_dict.get("batch_size", 30))
instance.max_queue_size = int(config_dict.get("max_queue_size", 200))
@@ -111,6 +251,8 @@ async def create(
"excluded_categories", ["promotions", "social"]
)
instance.gmail_query = config_dict.get("gmail_query")
+ if config.expected_mailbox is not None:
+ await instance.validate()
return instance
@@ -230,18 +372,16 @@ def _message_matches_label_filters(self, message_data: Dict) -> bool:
async def _authed_headers(self) -> Dict[str, str]:
"""Build Authorization headers with a fresh token."""
- token = await self.auth.get_token()
- return {"Authorization": f"Bearer {token}"}
+ return await authorization_headers(self.auth)
async def _refresh_and_get_headers(self) -> Dict[str, str]:
"""Force-refresh the token and return updated headers."""
- new_token = await self.auth.force_refresh()
- return {"Authorization": f"Bearer {new_token}"}
+ return await authorization_headers(self.auth, refresh=True)
@retry(
stop=stop_after_attempt(5),
- retry=_should_retry_gmail_request,
- wait=wait_rate_limit_with_backoff,
+ retry=retry_if_exception(_should_retry_gmail_request),
+ wait=_wait_gmail_request,
reraise=True,
)
async def _get(self, url: str, params: Optional[dict] = None) -> dict:
@@ -259,11 +399,9 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict:
response = await self.http_client.get(url, headers=headers, params=params)
if response.status_code == 429:
- self.logger.warning(
- f"Got 429 Rate Limited from Gmail API. Headers: {response.headers}. "
- f"Body: {response.text}."
- )
+ self.logger.warning("Gmail rate limit reached")
+ raise_gmail_throttle(response)
raise_for_status(
response,
source_short_name=self.short_name,
@@ -274,6 +412,68 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict:
self.logger.debug(f"Response data keys: {list(data.keys())}")
return data
+ @retry(
+ stop=stop_after_attempt(5),
+ retry=retry_if_exception(_should_retry_gmail_request),
+ wait=_wait_gmail_request,
+ reraise=True,
+ )
+ async def _get_capture_json(self, url: str, params: Optional[dict] = None) -> dict:
+ """Bound canonical JSON responses while streaming; native failures remain visible."""
+ return await self._get_mime_json(url, max_bytes=32 * 1024 * 1024, params=params)
+
+ async def _get_mime_json(
+ self, url: str, *, max_bytes: int, params: Optional[dict] = None
+ ) -> dict:
+ """Read bounded attachment JSON through the managed source HTTP transport."""
+ headers = {**await self._authed_headers(), "Accept-Encoding": "identity"}
+ for attempt in range(2):
+ async with self.http_client.stream(
+ "GET", url, headers=headers, params=params
+ ) as response:
+ if response.status_code == 401 and attempt == 0 and self.auth.supports_refresh:
+ headers = {
+ **await self._refresh_and_get_headers(),
+ "Accept-Encoding": "identity",
+ }
+ continue
+ if response.headers.get("content-encoding", "identity").lower() != "identity":
+ raise SourceError(
+ "Gmail did not honor uncompressed bounded response request",
+ source_short_name="gmail",
+ )
+ maximum = max_bytes if response.is_success else min(max_bytes, 65536)
+ body = await self._bounded_response_bytes(response, maximum)
+ buffered = httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ request=response.request,
+ content=bytes(body),
+ )
+ if response.status_code == 400 and params and params.get("pageToken"):
+ try:
+ native_error = json.loads(body)
+ except (ValueError, UnicodeError):
+ native_error = None
+ if isinstance(native_error, dict) and invalid_page_token(native_error):
+ raise InvalidScanContinuation("Gmail rejected its saved page token")
+ raise_gmail_throttle(buffered)
+ raise_for_status(
+ buffered,
+ source_short_name=self.short_name,
+ token_provider_kind=self.auth.provider_kind,
+ )
+ value = json.loads(body)
+ if not isinstance(value, dict):
+ raise ValueError("Gmail response must be an object")
+ return value
+ raise AssertionError("Unreachable Gmail refresh state")
+
+ @staticmethod
+ async def _bounded_response_bytes(response: httpx.Response, maximum: int) -> bytes:
+ """Own the raw iterator; identity encoding makes this a decoded-content bound too."""
+ return await bounded_response_bytes(response, maximum, label="Gmail response")
+
# -----------------------
# Cursor helper
# -----------------------
@@ -894,5 +1094,18 @@ async def generate_entities(
raise
async def validate(self) -> None:
- """Validate credentials by pinging the Gmail user profile."""
- await self._get("https://gmail.googleapis.com/gmail/v1/users/me/profile")
+ """Re-attest the native mailbox before permitting owned capture."""
+ self._verified_mailbox = None
+ raw = await self._get_capture_json(
+ "https://gmail.googleapis.com/gmail/v1/users/me/profile",
+ params={"fields": "emailAddress"},
+ )
+ try:
+ profile = GmailProfile.model_validate(raw)
+ except ValidationError:
+ raise ValueError("Gmail returned an invalid mailbox identity") from None
+ expected = self.capture_config.expected_mailbox
+ if expected is not None:
+ if profile.emailAddress.casefold() != expected.casefold():
+ raise ValueError("Gmail mailbox identity does not match the trusted binding")
+ self._verified_mailbox = expected.casefold()
diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py
new file mode 100644
index 0000000000..b0d39d0332
--- /dev/null
+++ b/backend/airweave/platform/sources/gmail_capture.py
@@ -0,0 +1,270 @@
+"""Original Gmail message capture, independent of derived thread/search entities."""
+
+import base64
+import hashlib
+import json
+from collections.abc import Awaitable, Callable, Iterator
+from datetime import datetime, timezone
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field
+
+from airweave.domains.entities.canonical.page_source import InvalidScanContinuation
+from airweave.domains.entities.canonical.requests import (
+ BlobReference,
+ CaptureRecord,
+ RecordIdentity,
+)
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError, SourceError
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+
+BASE = "https://gmail.googleapis.com/gmail/v1/users/me"
+Get = Callable[..., Awaitable[dict]]
+
+
+class _HistoryMessage(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(min_length=1)
+
+
+class _HistoryChange(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ message: _HistoryMessage
+
+
+class _HistoryEntry(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ messages: list[_HistoryMessage] = Field(default_factory=list)
+ messagesAdded: list[_HistoryChange] = Field(default_factory=list)
+ messagesDeleted: list[_HistoryChange] = Field(default_factory=list)
+ labelsAdded: list[_HistoryChange] = Field(default_factory=list)
+ labelsRemoved: list[_HistoryChange] = Field(default_factory=list)
+
+
+class _HistoryResponse(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ history: list[_HistoryEntry] = Field(default_factory=list)
+ historyId: str = Field(min_length=1)
+ nextPageToken: str | None = Field(default=None, min_length=1)
+
+
+class HistoryPage(BaseModel):
+ """One provider page; its mailbox boundary is not a committed checkpoint."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ message_ids: tuple[str, ...]
+ history_id: str
+ next_page_token: str | None
+ fingerprint: str
+
+
+class HydratedMessages(BaseModel):
+ """Bounded current observations; callers atomically commit records and offset."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ records: tuple[CaptureRecord, ...] = Field(max_length=500)
+ next_offset: int
+ complete: bool
+
+
+def parse_history_page(raw: dict) -> HistoryPage:
+ """Normalize only one page, preserving first occurrence order across all event arrays."""
+ response = _HistoryResponse.model_validate(raw)
+ affected: dict[str, None] = {}
+ for entry in response.history:
+ for message in entry.messages:
+ affected[message.id] = None
+ for changes in (
+ entry.messagesAdded,
+ entry.messagesDeleted,
+ entry.labelsAdded,
+ entry.labelsRemoved,
+ ):
+ for change in changes:
+ affected[change.message.id] = None
+ ids = tuple(affected)
+ # Intermediate mailbox historyId can advance without changing this page.
+ # Terminal pages bind the candidate checkpoint to the exact hydrated response.
+ # Include native events: a later edit can leave deduplicated IDs unchanged.
+ fingerprint = hashlib.sha256(
+ json.dumps(
+ [
+ raw.get("history", []),
+ response.nextPageToken,
+ response.historyId if response.nextPageToken is None else None,
+ ],
+ sort_keys=True,
+ separators=(",", ":"),
+ allow_nan=False,
+ ).encode()
+ ).hexdigest()
+ return HistoryPage(
+ message_ids=ids,
+ history_id=response.historyId,
+ next_page_token=response.nextPageToken,
+ fingerprint=fingerprint,
+ )
+
+
+def _external_parts(part: dict, path: str = "/payload") -> Iterator[tuple[dict, str]]:
+ """Keep a JSON Pointer to each external MIME body in untouched provider JSON."""
+ if part.get("body", {}).get("attachmentId"):
+ yield part, path + "/body"
+ for index, child in enumerate(part.get("parts", [])):
+ yield from _external_parts(child, f"{path}/parts/{index}")
+
+
+class GmailCapture:
+ """Fetch raw provider messages with explicit reconciliation and cursor barriers."""
+
+ def __init__(
+ self,
+ get: Get,
+ query: str | None,
+ *,
+ files: FileService | None = None,
+ attachment_get: Get | None = None,
+ ):
+ """Reuse the source managed HTTP/auth transport and exact provider query."""
+ self.get = get
+ self.query = query
+ self.files = files
+ self.attachment_get = attachment_get
+
+ async def message(self, message_id: str) -> CaptureRecord:
+ """Observe current state; a disappeared message is a provider deletion."""
+ identity = RecordIdentity(record_type="message", native_id=message_id)
+ try:
+ payload = await self.get(f"{BASE}/messages/{message_id}", params={"format": "full"})
+ except (httpx.HTTPStatusError, SourceEntityNotFoundError) as exc:
+ if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code != 404:
+ raise
+ return CaptureRecord(
+ identity=identity,
+ payload={"id": message_id},
+ kind="delete",
+ removal_reason="provider_deleted",
+ observed_at=datetime.now(timezone.utc),
+ )
+ if payload.get("id") != message_id:
+ raise ValueError("Gmail message response identity mismatch")
+ blobs, complete = await self.capture_bodies(message_id, payload)
+ created = payload.get("internalDate")
+ return CaptureRecord(
+ identity=identity,
+ payload=payload,
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=(
+ datetime.fromtimestamp(int(created) / 1000, timezone.utc)
+ if created is not None
+ else None
+ ),
+ completeness="complete" if complete else "partial",
+ blobs=blobs,
+ )
+
+ async def capture_bodies(
+ self,
+ message_id: str,
+ payload: dict,
+ ) -> tuple[tuple[BlobReference, ...], bool]:
+ """Persist MIME bytes before their record; only deliberate size skips are partial."""
+ blobs = []
+ complete = True
+ for part, path in _external_parts(payload.get("payload", {})):
+ if self.files is None or self.attachment_get is None:
+ complete = False
+ continue
+ try:
+ content = await self.body_bytes(message_id, part["body"])
+ blob = await self.files.store_canonical_blob(
+ content, media_type=part.get("mimeType")
+ )
+ except FileSkippedException:
+ complete = False
+ continue
+ blobs.append(
+ BlobReference.model_validate(
+ {
+ **blob.model_dump(),
+ "source_path": path,
+ "filename": part.get("filename") or None,
+ }
+ )
+ )
+ return tuple(blobs), complete
+
+ async def body_bytes(self, message_id: str, body: dict) -> bytes:
+ """Validate declared, encoded and decoded lengths before accepting provider bytes."""
+ maximum = self.files.MAX_FILE_SIZE_BYTES
+ if int(body.get("size", 0)) > maximum:
+ raise FileSkippedException("MIME body exceeds size limit", "MIME body")
+ encoded_limit = ((maximum + 2) // 3) * 4
+ result = await self.attachment_get(
+ f"{BASE}/messages/{message_id}/attachments/{body['attachmentId']}",
+ max_bytes=encoded_limit + 4096,
+ )
+ encoded = result["data"]
+ if len(encoded) > encoded_limit:
+ raise FileSkippedException("Encoded MIME body exceeds size limit", "MIME body")
+ content = base64.b64decode(
+ encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True
+ )
+ if len(content) > maximum:
+ raise FileSkippedException("Decoded MIME body exceeds size limit", "MIME body")
+ if len(content) != int(result["size"]) or len(content) != int(body["size"]):
+ raise ValueError("Gmail MIME body size does not match provider metadata")
+ return content
+
+ async def history_page(
+ self, boundary: str, token: str | None = None, *, max_results: int = 500
+ ) -> HistoryPage:
+ """Fetch one history page; boundaries stay opaque and no cursor is advanced."""
+ if not boundary or not 1 <= max_results <= 500:
+ raise ValueError("History requires a boundary and a page size from 1 to 500")
+ params = {"startHistoryId": boundary, "maxResults": max_results}
+ if token is not None:
+ params["pageToken"] = token
+ return parse_history_page(await self.get(f"{BASE}/history", params=params))
+
+ async def hydrate_history_page(
+ self,
+ page: HistoryPage,
+ *,
+ offset: int = 0,
+ limit: int = 500,
+ expected_fingerprint: str | None = None,
+ ) -> HydratedMessages:
+ """Reuse exact current reads, never apply stale history deletion/label payloads."""
+ if expected_fingerprint is not None and page.fingerprint != expected_fingerprint:
+ raise InvalidScanContinuation("Gmail history page changed before offset replay")
+ if offset and expected_fingerprint is None:
+ raise ValueError("Resuming Gmail history hydration requires its page fingerprint")
+ if not 0 <= offset <= len(page.message_ids) or not 1 <= limit <= 500:
+ raise ValueError("Invalid Gmail history hydration offset or batch size")
+ return await self.hydrate_messages(page.message_ids, offset=offset, limit=limit)
+
+ async def hydrate_messages(
+ self, message_ids: tuple[str, ...], *, offset: int = 0, limit: int = 50
+ ) -> HydratedMessages:
+ """Stop at complete originals; bounded byte targets never truncate native content."""
+ if not 0 <= offset <= len(message_ids) or not 1 <= limit <= 500:
+ raise ValueError("Invalid Gmail message hydration boundary")
+ records = []
+ total_bytes = 0
+ for mid in message_ids[offset : offset + limit]:
+ record = await self.message(mid)
+ size = len(record.model_dump_json().encode())
+ if size > 32 * 1024 * 1024:
+ raise SourceError(
+ "Gmail original exceeds the 32 MiB record limit", source_short_name="gmail"
+ )
+ records.append(record)
+ total_bytes += size
+ if total_bytes >= 8 * 1024 * 1024:
+ break
+ end = offset + len(records)
+ return HydratedMessages(
+ records=tuple(records), next_offset=end, complete=end == len(message_ids)
+ )
diff --git a/backend/airweave/platform/sources/gmail_errors.py b/backend/airweave/platform/sources/gmail_errors.py
new file mode 100644
index 0000000000..11742952be
--- /dev/null
+++ b/backend/airweave/platform/sources/gmail_errors.py
@@ -0,0 +1,75 @@
+"""Gmail's structured 403 quota reasons, separate from permission denial."""
+
+import math
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+from airweave.domains.sources.exceptions import SourceServerError
+
+
+class _Reason(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ domain: str
+ reason: str
+
+
+class _Error(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ code: int
+ errors: list[_Reason] = Field(min_length=1)
+
+
+class _Envelope(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ error: _Error
+
+
+class GmailThrottleError(SourceServerError):
+ """Native 403 throttle; unknown provider delay remains unknown."""
+
+ status_code = 403
+
+ def __init__(self, retry_after: float | None):
+ """Retain only safe native timing, never provider bodies or request URLs."""
+ self.retry_after = retry_after
+ super().__init__(
+ "Gmail rate limit exceeded (403); capture remains resumable",
+ source_short_name="gmail",
+ status_code=403,
+ )
+
+
+def _retry_after(value: str | None) -> float | None:
+ if value is None:
+ return None
+ try:
+ seconds = float(value)
+ except ValueError:
+ try:
+ date = parsedate_to_datetime(value)
+ if date.tzinfo is None:
+ return None
+ seconds = (date - datetime.now(timezone.utc)).total_seconds()
+ except (ValueError, TypeError, OverflowError):
+ return None
+ return max(0.0, seconds) if math.isfinite(seconds) else None
+
+
+def raise_gmail_throttle(response: httpx.Response) -> None:
+ """Retry only documented transient quota reasons, not mixed or unknown errors."""
+ if response.status_code != 403:
+ return
+ try:
+ envelope = _Envelope.model_validate_json(response.content)
+ except ValidationError:
+ return
+ if envelope.error.code != 403 or not all(
+ item.domain == "usageLimits"
+ and item.reason in {"rateLimitExceeded", "userRateLimitExceeded"}
+ for item in envelope.error.errors
+ ):
+ return
+ raise GmailThrottleError(_retry_after(response.headers.get("Retry-After")))
diff --git a/backend/airweave/platform/sources/gmail_pages.py b/backend/airweave/platform/sources/gmail_pages.py
new file mode 100644
index 0000000000..1620c43048
--- /dev/null
+++ b/backend/airweave/platform/sources/gmail_pages.py
@@ -0,0 +1,268 @@
+"""Gmail native page acquisition; the canonical engine owns every durable acknowledgement."""
+
+import hashlib
+from typing import Annotated, Literal
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
+
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CycleConfiguration,
+ ProviderCheckpoint,
+)
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidCaptureCheckpoint,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError
+from airweave.platform.sources.gmail_capture import BASE, GmailCapture
+
+PAGE_SIZE = 50
+MessageID = Annotated[str, Field(min_length=1, max_length=128)]
+
+
+class _MessageID(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: MessageID
+
+
+class _ListPage(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ messages: list[_MessageID] = Field(default_factory=list, max_length=PAGE_SIZE)
+ nextPageToken: str | None = Field(default=None, min_length=1, max_length=8192)
+
+
+class _ParameterError(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ location: str | None = None
+ locationType: str | None = None
+ reason: str | None = None
+
+
+class _ErrorBody(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ errors: list[_ParameterError] = Field(default_factory=list)
+
+
+class _ErrorEnvelope(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ error: _ErrorBody
+
+
+def invalid_page_token(raw: dict) -> bool:
+ """Only a native parameter-specific error authorizes page-token recovery."""
+ try:
+ body = _ErrorEnvelope.model_validate(raw)
+ except ValidationError:
+ return False
+ return any(
+ item.location == "pageToken"
+ and item.locationType == "parameter"
+ and item.reason in {"badRequest", "invalidArgument"}
+ for item in body.error.errors
+ )
+
+
+class _Boundary(BaseModel):
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ history_id: str = Field(min_length=1, max_length=1024)
+
+
+class _Progress(BaseModel):
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ version: Literal[1] = 1
+ mode: Literal["full", "changes"]
+ phase: Literal["list", "history", "done"]
+ history_boundary: str | None = Field(default=None, max_length=1024)
+ list_token: str | None = Field(default=None, max_length=8192)
+ listing_complete: bool = False
+ pending_ids: tuple[MessageID, ...] = Field(default=(), max_length=PAGE_SIZE)
+ history_token: str | None = Field(default=None, max_length=8192)
+ history_fingerprint: str | None = None
+ history_offset: int = Field(default=0, ge=0)
+ list_tokens: tuple[str, ...] = Field(default=(), max_length=128)
+ history_tokens: tuple[str, ...] = Field(default=(), max_length=128)
+
+ @model_validator(mode="after")
+ def valid_phase(self):
+ if self.phase == "history" and not self.history_boundary:
+ raise ValueError("History traversal requires its immutable start boundary")
+ if self.history_offset and not self.history_fingerprint:
+ raise ValueError("History offset requires an exact replay page fingerprint")
+ if self.pending_ids and self.phase != "list":
+ raise ValueError("Only a baseline list page may retain pending native IDs")
+ if len(set(self.pending_ids)) != len(self.pending_ids):
+ raise ValueError("Pending Gmail identities must be unique")
+ if self.listing_complete and self.list_token is not None:
+ raise ValueError("Completed listing cannot retain a next-page token")
+ if self.mode == "changes" and (not self.listing_complete or self.phase == "list"):
+ raise ValueError("Changes cannot enumerate a mailbox inventory")
+ return self
+
+ def continuation(self) -> ScanContinuation:
+ validated = type(self).model_validate(self.model_dump())
+ return ScanContinuation(value=validated.model_dump(mode="json"))
+
+
+def _next_token(token: str | None, recent: tuple[str, ...]) -> tuple[str, ...]:
+ """Detect recent cursor loops without imposing a total page ceiling."""
+ if token is None:
+ return recent
+ digest = hashlib.sha256(token.encode()).hexdigest()
+ if digest in recent:
+ raise ValueError("Gmail pagination repeated a recent cursor")
+ return (*recent[-127:], digest)
+
+
+class GmailPages:
+ """Small source-local acquisition functions; no cursor writes, SQL, or background workers."""
+
+ def __init__(self, capture: GmailCapture):
+ """Reuse native message/MIME acquisition and the source transport."""
+ self.capture = capture
+
+ async def prepare(
+ self, previous: CaptureCycle | None, configuration: CycleConfiguration
+ ) -> CapturePlan:
+ """A legacy cursor never substitutes for compatible canonical full-capture evidence."""
+ if self.capture.query:
+ return CapturePlan()
+ if (
+ previous is not None
+ and previous.configuration == configuration
+ and (
+ previous.last_full_capture is not None and previous.promoted_checkpoint is not None
+ )
+ ):
+ return CapturePlan(
+ mode="changes", starting_checkpoint=previous.promoted_checkpoint.checkpoint
+ )
+ raw = await self.capture.get(f"{BASE}/profile")
+ boundary = _Boundary(history_id=raw["historyId"])
+ return CapturePlan(starting_checkpoint=ProviderCheckpoint(value=boundary.model_dump()))
+
+ @staticmethod
+ def initial(cycle: CaptureCycle) -> ScanContinuation:
+ """Only the persisted cycle decides the first page and history boundary."""
+ boundary = (
+ _Boundary.model_validate(cycle.starting_checkpoint.value).history_id
+ if cycle.starting_checkpoint
+ else None
+ )
+ return _Progress(
+ mode=cycle.mode,
+ phase="history" if cycle.mode == "changes" else "list",
+ history_boundary=boundary,
+ listing_complete=cycle.mode == "changes",
+ ).continuation()
+
+ async def page(self, continuation: ScanContinuation) -> CapturePage:
+ """Return one bounded fully hydrated batch, without acknowledging it."""
+ state = _Progress.model_validate(continuation.value)
+ if state.phase == "done":
+ raise ValueError("A completed Gmail page must not be fetched again")
+ if bool(self.capture.query) != (state.history_boundary is None):
+ raise ValueError("Gmail continuation belongs to a different query mode")
+ return await self._list(state) if state.phase == "list" else await self._history(state)
+
+ async def _list(self, state: _Progress) -> CapturePage:
+ if state.pending_ids:
+ return await self._hydrate_list(state)
+ params = {"maxResults": PAGE_SIZE, "includeSpamTrash": "true"}
+ if self.capture.query:
+ params["q"] = self.capture.query
+ if state.list_token is not None:
+ params["pageToken"] = state.list_token
+ raw = await self.capture.get(f"{BASE}/messages", params=params)
+ listing = _ListPage.model_validate(raw)
+ tokens = _next_token(listing.nextPageToken, state.list_tokens)
+ identities = tuple(item.id for item in listing.messages)
+ if len(set(identities)) != len(identities):
+ raise ValueError("Gmail list response duplicated a native message identity")
+ state = state.model_copy(
+ update={
+ "list_token": listing.nextPageToken,
+ "listing_complete": listing.nextPageToken is None,
+ "list_tokens": tokens,
+ "pending_ids": identities,
+ }
+ )
+ return await self._hydrate_list(state)
+
+ async def _hydrate_list(self, state: _Progress) -> CapturePage:
+ batch = await self.capture.hydrate_messages(state.pending_ids)
+ pending = state.pending_ids[batch.next_offset :]
+ final = bool(self.capture.query) and state.listing_complete and not pending
+ following = state.model_copy(
+ update={
+ "pending_ids": pending,
+ "phase": "list"
+ if pending
+ else "done"
+ if final
+ else "list"
+ if self.capture.query
+ else "history",
+ }
+ )
+ return CapturePage(
+ records=batch.records, continuation=following.continuation(), final=final
+ )
+
+ async def _history(self, state: _Progress) -> CapturePage:
+ try:
+ page = await self.capture.history_page(
+ state.history_boundary, state.history_token, max_results=100
+ )
+ except (SourceEntityNotFoundError, httpx.HTTPStatusError) as error:
+ if isinstance(error, httpx.HTTPStatusError) and error.response.status_code != 404:
+ raise
+ raise InvalidCaptureCheckpoint(
+ "Gmail history boundary expired; full capture required"
+ ) from error
+ batch = await self.capture.hydrate_history_page(
+ page,
+ offset=state.history_offset,
+ limit=PAGE_SIZE,
+ expected_fingerprint=state.history_fingerprint,
+ )
+ if not batch.complete:
+ following = state.model_copy(
+ update={
+ "history_offset": batch.next_offset,
+ "history_fingerprint": page.fingerprint,
+ }
+ )
+ return CapturePage(records=batch.records, continuation=following.continuation())
+ if page.next_page_token is not None:
+ following = state.model_copy(
+ update={
+ "history_offset": 0,
+ "history_fingerprint": None,
+ "history_token": page.next_page_token,
+ "history_tokens": _next_token(page.next_page_token, state.history_tokens),
+ }
+ )
+ return CapturePage(records=batch.records, continuation=following.continuation())
+ final = state.listing_complete
+ following = state.model_copy(
+ update={
+ "phase": "done" if final else "list",
+ "history_boundary": page.history_id,
+ "history_token": None,
+ "history_fingerprint": None,
+ "history_offset": 0,
+ "history_tokens": (),
+ }
+ )
+ return CapturePage(
+ records=batch.records,
+ continuation=following.continuation(),
+ final=final,
+ provider_checkpoint=ProviderCheckpoint(value={"history_id": page.history_id})
+ if final
+ else None,
+ )
diff --git a/backend/airweave/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py
index 86d9f17d3a..eb8cc4f86f 100644
--- a/backend/airweave/platform/sources/google_calendar.py
+++ b/backend/airweave/platform/sources/google_calendar.py
@@ -1,51 +1,51 @@
-"""Google Calendar source implementation.
-
-Retrieves data from a user's Google Calendar (read-only mode):
- - CalendarList entries (the user's list of calendars)
- - Each underlying Calendar resource
- - Events belonging to each Calendar
- - (Optionally) Free/Busy data for each Calendar
-
-Follows the same structure and pattern as other connector implementations
-(e.g., Gmail, Asana, Todoist, HubSpot). The entity schemas are defined in
-entities/google_calendar.py.
-
-Reference:
- https://developers.google.com/calendar/api/v3/reference
-
-Now supports two flows:
- - Non-batching / sequential (default): preserves original behavior.
- - Batching / concurrent (opt-in): gated by `batch_generation` config and uses the
- bounded-concurrency driver in BaseSource across all major I/O points:
- * Per-calendar Calendar resource fetch
- * Per-calendar event listing (still sequential within a calendar due to pagination)
- * Per-calendar Free/Busy fetch
-
-Config (all optional, shown with defaults):
- {
- "batch_generation": False, # enable/disable concurrent generation
- "batch_size": 30, # max concurrent workers (calendars processed in parallel)
- "max_queue_size": 200, # backpressure queue size
- "preserve_order": False, # maintain calendar order when yielding results
- "stop_on_error": False # cancel all on first error
- }
+"""Google Calendar originals with durable per-calendar capture and observed range reads.
+
+Canonical capture uses the shared page engine: complete calendar membership,
+per-calendar full or incremental raw events, and a fixed expanded-instance window.
+The engine owns every page acknowledgement and completed scope publication.
+
+The upstream BaseSource entity interface remains implemented below; it is not the
+canonical cursor authority. Configured native calendar IDs and occurrence windows
+are validated by GoogleCalendarConfig.
"""
from __future__ import annotations
+import json
import urllib.parse
from datetime import datetime, timedelta
from typing import Any, AsyncGenerator, Dict, List, Optional
-from tenacity import retry, stop_after_attempt
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from tenacity import retry, retry_if_exception, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import (
+ AuthProviderRateLimitError,
+ AuthProviderServerError,
+)
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidScanContinuation,
+)
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation, ScanState
+from airweave.domains.entities.canonical.scope_execution import ScopePlan
+from airweave.domains.sources.exceptions import SourceError, SourceRateLimitError
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.config import GoogleCalendarConfig
+from airweave.platform.cursors.google_calendar import GoogleCalendarCursor
from airweave.platform.decorators import source
from airweave.platform.entities._base import BaseEntity, Breadcrumb
from airweave.platform.entities.google_calendar import (
@@ -55,15 +55,39 @@
GoogleCalendarListEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.bounded_response import bounded_response_bytes
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
+ should_retry_on_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.records.calendar_pages import CalendarPages, invalid_page_token
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
+def _retry_capture(exception: BaseException) -> bool:
+ """Do not retry earlier than a provider's minimum beyond this request wait budget."""
+ if isinstance(exception, (SourceRateLimitError, AuthProviderRateLimitError)):
+ return exception.retry_after <= 120
+ if isinstance(exception, AuthProviderServerError):
+ return True
+ if isinstance(exception, httpx.HTTPStatusError) and exception.response.status_code == 429:
+ try:
+ return float(exception.response.headers.get("Retry-After", "0")) <= 120
+ except ValueError:
+ return True
+ return should_retry_on_rate_limit_or_timeout(exception)
+
+
+class CalendarPrincipal(BaseModel):
+ """Primary calendar resource identity, not a Google OIDC subject."""
+
+ model_config = ConfigDict(extra="ignore")
+ id: str = Field(min_length=1, pattern=r"^\S+$")
+
+
@source(
name="Google Calendar",
short_name="google_calendar",
@@ -77,7 +101,8 @@
auth_config_class=None,
config_class=GoogleCalendarConfig,
labels=["Productivity", "Calendar"],
- supports_continuous=False,
+ supports_continuous=True,
+ cursor_class=GoogleCalendarCursor,
rate_limit_level=RateLimitLevel.ORG,
)
class GoogleCalendarSource(BaseSource):
@@ -89,6 +114,77 @@ class GoogleCalendarSource(BaseSource):
Google Calendar scheduling information for productivity and time management insights.
"""
+ _verified_primary_calendar_id: str | None = None
+
+ def _require_principal(self) -> None:
+ """Guard direct capture and resume even outside shared lifecycle validation."""
+ expected = self.calendar_config.expected_primary_calendar_id
+ if expected is None or self._verified_primary_calendar_id != expected:
+ raise ValueError(
+ "Owned Calendar capture requires an attested primary calendar identity"
+ )
+
+ def _calendar_pages(self) -> CalendarPages:
+ """Keep native page logic behind the same attestation gate."""
+ self._require_principal()
+ return CalendarPages(self._get_capture_json, self.calendar_config)
+
+ canonical_record_types = ("calendar", "event", "event_occurrence")
+ canonical_container_parents = {"event": "calendar", "event_occurrence": "calendar"}
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Calendar membership owns exact native event and occurrence scopes."""
+ return self._calendar_pages().configuration
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Resolve the expansion window once per persisted cycle."""
+ return self._calendar_pages().prepare()
+
+ async def prepare_scope(
+ self,
+ scope: CompletedScope,
+ cycle: CaptureCycle,
+ previous: ScanState | None,
+ *,
+ parent: SourceRecord | None,
+ force_full: bool,
+ ) -> ScopePlan:
+ """Choose native delta only from compatible completed exact-scope evidence."""
+ return self._calendar_pages().scope_plan(
+ scope, cycle, previous, parent=parent, force_full=force_full
+ )
+
+ def initial_scope_continuation(
+ self, scope: CompletedScope, cycle: CaptureCycle, plan: ScopePlan
+ ) -> ScanContinuation:
+ """Provider progress starts from the immutable engine-attested scope plan."""
+ self._require_principal()
+ return CalendarPages.initial(plan)
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Native calendar IDs remain the stable container namespace."""
+ return CompletedScope(
+ record_type=record_type, container_id=parent.identity.native_id, parent=parent.identity
+ )
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Originals and pagination commit together through the production page driver."""
+ return await self._calendar_pages().page(scope, continuation)
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """An accessible omitted calendar prevents premature membership reconciliation."""
+ if record.identity.record_type != "calendar" or record.parent is not None:
+ raise ValueError("Calendar omission must identify root membership")
+ await self._calendar_pages().confirm_member_absent(record.identity.native_id)
+
# -----------------------
# Construction / Config
# -----------------------
@@ -96,7 +192,7 @@ class GoogleCalendarSource(BaseSource):
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: GoogleCalendarConfig,
@@ -104,12 +200,15 @@ async def create(
"""Create a new Google Calendar source instance."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.calendar_config = config
config_dict = config.model_dump() if config else {}
instance.batch_generation = bool(config_dict.get("batch_generation", False))
instance.batch_size = int(config_dict.get("batch_size", 30))
instance.max_queue_size = int(config_dict.get("max_queue_size", 200))
instance.preserve_order = bool(config_dict.get("preserve_order", False))
instance.stop_on_error = bool(config_dict.get("stop_on_error", False))
+ if config.expected_primary_calendar_id is not None:
+ await instance.validate()
return instance
@@ -119,13 +218,61 @@ async def create(
async def _authed_headers(self) -> Dict[str, str]:
"""Build Authorization headers with a fresh token."""
- token = await self.auth.get_token()
- return {"Authorization": f"Bearer {token}"}
+ return await authorization_headers(self.auth)
async def _refresh_and_get_headers(self) -> Dict[str, str]:
"""Force-refresh the token and return updated headers."""
- new_token = await self.auth.force_refresh()
- return {"Authorization": f"Bearer {new_token}"}
+ return await authorization_headers(self.auth, refresh=True)
+
+ @retry(
+ stop=stop_after_attempt(5),
+ retry=retry_if_exception(_retry_capture),
+ wait=wait_rate_limit_with_backoff,
+ reraise=True,
+ )
+ async def _get_capture_json(self, url: str, params: Optional[dict] = None) -> dict:
+ """Bound native capture responses, preserving auth refresh and typed errors."""
+ headers = {**await self._authed_headers(), "Accept-Encoding": "identity"}
+ for attempt in range(2):
+ async with self.http_client.stream(
+ "GET", url, headers=headers, params=params
+ ) as response:
+ if response.status_code == 401 and attempt == 0 and self.auth.supports_refresh:
+ headers = {
+ **await self._refresh_and_get_headers(),
+ "Accept-Encoding": "identity",
+ }
+ continue
+ if response.headers.get("content-encoding", "identity").lower() != "identity":
+ raise SourceError(
+ "Calendar did not honor identity encoding",
+ source_short_name="google_calendar",
+ )
+ maximum = 32 * 1024 * 1024 if response.is_success else 65536
+ body = await bounded_response_bytes(response, maximum, label="Calendar response")
+ buffered = httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ request=response.request,
+ content=body,
+ )
+ if response.status_code == 400 and params and params.get("pageToken"):
+ try:
+ native_error = json.loads(body)
+ except (ValueError, UnicodeError):
+ native_error = None
+ if isinstance(native_error, dict) and invalid_page_token(native_error):
+ raise InvalidScanContinuation("Calendar rejected its saved page token")
+ raise_for_status(
+ buffered,
+ source_short_name=self.short_name,
+ token_provider_kind=self.auth.provider_kind,
+ )
+ value = json.loads(body)
+ if not isinstance(value, dict):
+ raise ValueError("Calendar response must be an object")
+ return value
+ raise AssertionError("Unreachable Calendar refresh state")
@retry(
stop=stop_after_attempt(5),
@@ -368,8 +515,23 @@ async def generate_entities(
yield entity
async def validate(self) -> None:
- """Validate credentials by pinging the Calendar API calendarList endpoint."""
- await self._get(
- "https://www.googleapis.com/calendar/v3/users/me/calendarList",
- params={"maxResults": "1"},
+ """Verify the bound primary calendar without inferring a Google subject."""
+ self._verified_primary_calendar_id = None
+ expected = self.calendar_config.expected_primary_calendar_id
+ if expected is None:
+ # Legacy extraction retains its original credential check.
+ await self._get(
+ "https://www.googleapis.com/calendar/v3/users/me/calendarList",
+ params={"maxResults": "1"},
+ )
+ return
+ raw = await self._get_capture_json(
+ "https://www.googleapis.com/calendar/v3/calendars/primary", params={"fields": "id"}
)
+ try:
+ principal = CalendarPrincipal.model_validate(raw)
+ except ValidationError:
+ raise ValueError("Calendar returned an invalid primary identity") from None
+ if principal.id != expected:
+ raise ValueError("Calendar primary identity does not match the trusted binding")
+ self._verified_primary_calendar_id = expected
diff --git a/backend/airweave/platform/sources/google_docs.py b/backend/airweave/platform/sources/google_docs.py
index f7063f7f3e..4eeec1f0ec 100644
--- a/backend/airweave/platform/sources/google_docs.py
+++ b/backend/airweave/platform/sources/google_docs.py
@@ -34,12 +34,12 @@
from airweave.platform.entities._base import BaseEntity
from airweave.platform.entities.google_docs import GoogleDocsDocumentEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py
index 18f3eefb58..0f0783dd14 100644
--- a/backend/airweave/platform/sources/google_drive.py
+++ b/backend/airweave/platform/sources/google_drive.py
@@ -16,17 +16,32 @@
from __future__ import annotations
+import hashlib
+import json
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional
import httpx
+from pydantic import BaseModel, ConfigDict, Field, StrictBool, ValidationError
from tenacity import retry, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
from airweave.domains.browse_tree.types import NodeSelectionData
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidScanContinuation,
+)
+from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
from airweave.domains.sources.exceptions import SourceAuthError, SourceError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage import FileSkippedException
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
@@ -41,15 +56,32 @@
_parse_drive_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.records.google_drive_content import capture_file_content
+from airweave.platform.sources.records.google_drive_pages import DrivePages, rejected_listing_token
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
+class DrivePrincipal(BaseModel):
+ """Native stable identity; display labels and email are not binding authority."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ permissionId: str = Field(min_length=1, pattern=r"^\S+$")
+ me: StrictBool
+
+
+class DriveAbout(BaseModel):
+ """Only the requested native principal response is interpreted."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ user: DrivePrincipal
+
+
@source(
name="Google Drive",
short_name="google_drive",
@@ -77,11 +109,108 @@ class GoogleDriveSource(BaseSource):
while maintaining proper organization and access permissions.
"""
+ canonical_record_types = ("file",)
+
+ canonical_container_parents = {}
+ expected_permission_id: str | None = None
+ capture_native_sheets: bool = False
+ _verified_permission_id: str | None = None
+
+ def _require_principal(self) -> None:
+ """Direct capture and resume cannot bypass native account attestation."""
+ if (
+ self.expected_permission_id is None
+ or self._verified_permission_id != self.expected_permission_id
+ ):
+ raise ValueError("Owned Drive capture requires an attested permission identity")
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """One unfiltered accessible corpus; legacy path selections remain unsupported."""
+ self._require_principal()
+ if self.include_patterns:
+ raise ValueError("Drive path selection is not yet supported by canonical capture")
+ return CycleConfiguration(
+ fingerprint=hashlib.sha256(
+ json.dumps(
+ {
+ "capture_version": 3,
+ "permission_id": self.expected_permission_id,
+ "scope": (
+ "drive-all-accessible-workspace-native-v2"
+ if self.capture_native_sheets
+ else "drive-all-accessible-docs-native"
+ ),
+ },
+ sort_keys=True,
+ ).encode()
+ ).hexdigest(),
+ parents={"file": (None,)},
+ known_object_validation=("file",),
+ )
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Only engine-attested full capture can authorize native changes."""
+ return await DrivePages(self._get).prepare(previous, self.capture_cycle_configuration)
+
+ def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation:
+ """Restore the persisted starting boundary."""
+ self._require_principal()
+ if cycle.configuration != self.capture_cycle_configuration:
+ raise ValueError("Drive cycle does not match the trusted capture configuration")
+ return DrivePages.initial(cycle)
+
+ async def _capture_body(self, record: CaptureRecord, files: FileService) -> CaptureRecord:
+ return await capture_file_content(
+ record,
+ files=files,
+ get=self._get,
+ client=self.http_client,
+ auth=self.auth,
+ logger=self.logger,
+ capture_native_sheets=self.capture_native_sheets,
+ )
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Commit one version-checked file without retaining a native response queue."""
+ self._require_principal()
+ if scope != CompletedScope(record_type="file") or parent is not None:
+ raise ValueError("Drive requires its independent file scope")
+
+ async def hydrate(record: CaptureRecord) -> CaptureRecord:
+ return await self._capture_body(record, files)
+
+ return await DrivePages(self._get).page(continuation, hydrate)
+
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Known list omissions require an exact accessible-file read."""
+ self._require_principal()
+ if record.identity.record_type != "file" or record.parent is not None:
+ raise ValueError("Drive omission must identify an independent file")
+ return await self._capture_body(
+ await DrivePages(self._get).current(record.identity.native_id), files
+ )
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Folders are native metadata, not proof of inherited access."""
+ raise ValueError("Drive has no canonical child scopes")
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Every omitted file uses exact known-object refresh."""
+ raise ValueError("Drive absence requires exact known-object refresh")
+
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: GoogleDriveConfig,
@@ -89,11 +218,15 @@ async def create(
"""Create a new Google Drive source instance."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
instance.include_patterns = config.include_patterns if config else []
+ instance.expected_permission_id = config.expected_permission_id if config else None
+ instance.capture_native_sheets = config.capture_native_sheets if config else False
instance.batch_size = 30
instance.batch_generation = True
instance.max_queue_size = 200
instance.preserve_order = False
instance.stop_on_error = False
+ if instance.expected_permission_id is not None:
+ await instance.validate()
return instance
@staticmethod
@@ -102,11 +235,22 @@ def _parse_datetime(value: Optional[str]) -> Optional[datetime]:
return _parse_drive_dt(value)
async def validate(self) -> None:
- """Validate credentials by pinging the shared drives list."""
- await self._get(
- "https://www.googleapis.com/drive/v3/drives",
- params={"pageSize": "1"},
+ """Re-attest the native principal; failure clears any previous attestation."""
+ self._verified_permission_id = None
+ raw = await self._get(
+ "https://www.googleapis.com/drive/v3/about",
+ params={"fields": "user(permissionId,me)"},
)
+ try:
+ principal = DriveAbout.model_validate(raw).user
+ except ValidationError:
+ raise ValueError("Drive returned an invalid principal identity") from None
+ if not principal.me:
+ raise ValueError("Drive returned a principal that is not the authenticated user")
+ if self.expected_permission_id is not None:
+ if principal.permissionId != self.expected_permission_id:
+ raise ValueError("Drive principal identity does not match the trusted binding")
+ self._verified_permission_id = principal.permissionId
@retry(
stop=stop_after_attempt(5),
@@ -123,15 +267,22 @@ async def _get(self, url: str, params: Optional[Dict] = None) -> Dict:
Max 5 attempts with intelligent wait strategy.
"""
- token = await self.auth.get_token()
- headers = {"Authorization": f"Bearer {token}"}
+ headers = await authorization_headers(self.auth)
response = await self.http_client.get(url, headers=headers, params=params, timeout=30.0)
if response.status_code == 401 and self.auth.supports_refresh:
- new_token = await self.auth.force_refresh()
- headers = {"Authorization": f"Bearer {new_token}"}
+ headers = await authorization_headers(self.auth, refresh=True)
response = await self.http_client.get(url, headers=headers, params=params, timeout=30.0)
+ if (
+ response.status_code == 400
+ and url == "https://www.googleapis.com/drive/v3/files"
+ and params
+ and params.get("pageToken")
+ and rejected_listing_token(response.json())
+ ):
+ raise InvalidScanContinuation("Drive rejected its saved inventory page token")
+
raise_for_status(
response,
source_short_name=self.short_name,
diff --git a/backend/airweave/platform/sources/google_slides.py b/backend/airweave/platform/sources/google_slides.py
index 1b6c8c155f..86feff4b26 100644
--- a/backend/airweave/platform/sources/google_slides.py
+++ b/backend/airweave/platform/sources/google_slides.py
@@ -38,12 +38,12 @@
GoogleSlidesPresentationEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/http_helpers.py b/backend/airweave/platform/sources/http_helpers.py
index 593940f351..02edda11c6 100644
--- a/backend/airweave/platform/sources/http_helpers.py
+++ b/backend/airweave/platform/sources/http_helpers.py
@@ -27,6 +27,7 @@
SourceEntityForbiddenError,
SourceEntityNotFoundError,
SourceError,
+ SourceGoneError,
SourceRateLimitError,
SourceServerError,
)
@@ -107,6 +108,8 @@ def _handle_redirect(
def _handle_fallback(
resp: httpx.Response, status: int, detail: str, ctx: str, sn: str, _tpk: AuthProviderKind
) -> None:
+ if status == 410:
+ raise SourceGoneError(f"Resource or continuation expired (410){ctx}", source_short_name=sn)
if status >= 500:
raise SourceServerError(
f"Server error ({status}){ctx}: {detail}",
@@ -157,7 +160,7 @@ def translate_httpx_error(
def _extract_detail(response: httpx.Response, max_length: int = 200) -> str:
- """Pull a human-readable snippet from the response body."""
+ """Use buffered detail only; status translation must never consume a stream."""
try:
data = response.json()
for key in ("message", "error_description", "error", "detail", "errors"):
@@ -165,6 +168,8 @@ def _extract_detail(response: httpx.Response, max_length: int = 200) -> str:
if val:
return str(val)[:max_length]
return str(data)[:max_length]
+ except httpx.ResponseNotRead:
+ return f"(unread response body, status {response.status_code})"
except Exception:
text = response.text[:max_length] if response.text else ""
return text or f"(no body, status {response.status_code})"
diff --git a/backend/airweave/platform/sources/hubspot.py b/backend/airweave/platform/sources/hubspot.py
index c56551a82c..430e5999ed 100644
--- a/backend/airweave/platform/sources/hubspot.py
+++ b/backend/airweave/platform/sources/hubspot.py
@@ -26,12 +26,12 @@
parse_hubspot_datetime,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/intercom.py b/backend/airweave/platform/sources/intercom.py
index a8120e1062..f52b017286 100644
--- a/backend/airweave/platform/sources/intercom.py
+++ b/backend/airweave/platform/sources/intercom.py
@@ -31,12 +31,12 @@
_unwrap_list,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
API_BASE = "https://api.intercom.io"
diff --git a/backend/airweave/platform/sources/jira.py b/backend/airweave/platform/sources/jira.py
index 660be3caa6..e52a3aabea 100644
--- a/backend/airweave/platform/sources/jira.py
+++ b/backend/airweave/platform/sources/jira.py
@@ -33,12 +33,12 @@
ZephyrTestPlanEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
ATLASSIAN_ACCESSIBLE_RESOURCES_URL = "https://api.atlassian.com/oauth/token/accessible-resources"
diff --git a/backend/airweave/platform/sources/linear.py b/backend/airweave/platform/sources/linear.py
index bd7fe76f98..5879cd04ce 100644
--- a/backend/airweave/platform/sources/linear.py
+++ b/backend/airweave/platform/sources/linear.py
@@ -1,72 +1,111 @@
-"""Linear source implementation for Airweave platform."""
+"""Selected-team Linear originals, with engine-owned durable full-scope pagination."""
from __future__ import annotations
-import mimetypes
-import os
-import re
+import hashlib
+import json
import time
+from collections.abc import AsyncGenerator
+from contextlib import aclosing
from datetime import datetime, timezone
-from typing import AsyncGenerator, Dict, List, Optional, Union
-from uuid import uuid4
+from functools import partial
+from pathlib import Path
+from typing import Literal
+from urllib.parse import urlsplit
+from uuid import UUID
import httpx
-from tenacity import retry, stop_after_attempt
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import (
- SourceAuthError,
- SourceError,
- SourceRateLimitError,
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import CapturePage, ScopeAccessLost
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceError, SourceRateLimitError
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
)
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
-from airweave.domains.storage import FileSkippedException
from airweave.domains.storage.file_service import FileService
+from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.config import LinearConfig
-from airweave.platform.cursors.linear import LinearCursor
from airweave.platform.decorators import source
-from airweave.platform.entities._base import BaseEntity, Breadcrumb
-from airweave.platform.entities.linear import (
- LinearAttachmentEntity,
- LinearCommentEntity,
- LinearIssueEntity,
- LinearProjectEntity,
- LinearTeamEntity,
- LinearUserEntity,
-)
+from airweave.platform.entities._base import BaseEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
-_GRAPHQL_URL = "https://api.linear.app/graphql"
+_URL = "https://api.linear.app/graphql"
+_QUERIES = Path(__file__).with_name("linear_queries")
+_QUERY_NAMES = ("identity", "teams", "issues", "comments", "attachments", "issue-membership")
+QUERIES = {name: (_QUERIES / f"{name}.graphql").read_text() for name in _QUERY_NAMES}
+FIELD_SET_VERSION = 3
-def _is_graphql_rate_limited(body: Dict) -> bool:
- """Check if a Linear GraphQL response contains a RATELIMITED error."""
- for err in body.get("errors", []):
- if err.get("extensions", {}).get("code") == "RATELIMITED":
- return True
- return False
+class _Model(BaseModel):
+ model_config = ConfigDict(extra="ignore")
-def _parse_reset_header(response: httpx.Response, default: float = 60.0) -> float:
- """Compute seconds-until-reset from ``X-RateLimit-Requests-Reset`` (epoch ms)."""
- raw = response.headers.get("X-RateLimit-Requests-Reset")
- if raw:
- try:
- reset_epoch_s = int(raw) / 1000.0
- return max(1.0, reset_epoch_s - time.time())
- except (ValueError, TypeError):
- pass
- return default
+class _Identity(_Model):
+ id: UUID
+
+
+class _PageInfo(_Model):
+ hasNextPage: StrictBool
+ endCursor: str | None
+
+
+class _Connection(_Model):
+ nodes: list[dict[str, JsonValue]] = Field(max_length=50)
+ pageInfo: _PageInfo
+
+
+class _Issue(_Identity):
+ team: _Identity
+
+
+class _Child(_Identity):
+ issue: _Identity
+
+
+class _Dates(_Model):
+ createdAt: AwareDatetime
+ updatedAt: AwareDatetime
+
+
+class _ErrorExtensions(_Model):
+ code: str | None = None
+
+
+class _GraphError(_Model):
+ extensions: _ErrorExtensions = Field(default_factory=_ErrorExtensions)
+
+
+class _Envelope(_Model):
+ data: dict[str, JsonValue] | None = None
+ errors: list[_GraphError] = Field(default_factory=list)
+
+
+class _Progress(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ after: str | None = Field(default=None, min_length=1)
@source(
@@ -81,534 +120,292 @@ def _parse_reset_header(response: httpx.Response, default: float = 60.0) -> floa
auth_config_class=None,
config_class=LinearConfig,
labels=["Project Management"],
- supports_continuous=True,
- cursor_class=LinearCursor,
+ supports_continuous=False,
rate_limit_level=RateLimitLevel.ORG,
)
class LinearSource(BaseSource):
- """Linear source connector — syncs teams, projects, users, issues, comments, attachments."""
+ """Capture selected issue originals and independently paginated child originals."""
+
+ canonical_record_types = ("issue", "comment", "attachment")
+ canonical_container_parents = {"comment": "issue", "attachment": "issue"}
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Expose canonical capability on the class, with instance-specific scope binding."""
+ return self._cycle_configuration
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: LinearConfig,
) -> LinearSource:
- """Create a new Linear source instance."""
- return cls(auth=auth, logger=logger, http_client=http_client)
-
- # ------------------------------------------------------------------
- # HTTP
- # ------------------------------------------------------------------
+ """Attest workspace and selected-team access before constructing a capture cycle."""
+ instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.config = config
+ instance.team_ids = frozenset(str(value) for value in config.team_ids)
+ fingerprint = {
+ "workspace": str(config.workspace_id),
+ "teams": sorted(instance.team_ids),
+ "field_set": FIELD_SET_VERSION,
+ "include_archived": True,
+ "queries": QUERIES,
+ "kinds": cls.canonical_record_types,
+ }
+ instance._cycle_configuration = CycleConfiguration.from_source(
+ fingerprint=hashlib.sha256(
+ json.dumps(fingerprint, sort_keys=True).encode()
+ ).hexdigest(),
+ record_types=cls.canonical_record_types,
+ container_parents=cls.canonical_container_parents,
+ )
+ await instance.validate()
+ return instance
@retry(
stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
reraise=True,
)
- async def _post(self, query: str) -> Dict:
- """Authenticated GraphQL POST with 401 refresh and raise_for_status.
-
- Linear returns rate-limit errors as HTTP 400 with a GraphQL
- ``RATELIMITED`` error code (not HTTP 429). We detect this before
- ``raise_for_status`` so the tenacity retry fires correctly.
- """
- token = await self.auth.get_token()
- headers = {
- "Content-Type": "application/json",
- "Authorization": f"Bearer {token}",
- }
- response = await self.http_client.post(_GRAPHQL_URL, headers=headers, json={"query": query})
-
+ async def _query(self, name: str, variables: dict[str, JsonValue]) -> dict[str, JsonValue]:
+ headers = await authorization_headers(self.auth)
+ response = await self.http_client.post(
+ _URL, headers=headers, json={"query": QUERIES[name], "variables": variables}
+ )
if response.status_code == 401 and self.auth.supports_refresh:
- new_token = await self.auth.force_refresh()
- headers["Authorization"] = f"Bearer {new_token}"
+ headers = await authorization_headers(self.auth, refresh=True)
response = await self.http_client.post(
- _GRAPHQL_URL, headers=headers, json={"query": query}
- )
-
- body = response.json()
-
- if response.status_code == 400 and _is_graphql_rate_limited(body):
- retry_after = _parse_reset_header(response)
- raise SourceRateLimitError(
- retry_after=retry_after,
- source_short_name=self.short_name,
- message=f"Linear RATELIMITED (400). Retry after {retry_after:.0f}s",
+ _URL, headers=headers, json={"query": QUERIES[name], "variables": variables}
)
-
+ # Inspect only the documented GraphQL rate-limit response before HTTP translation.
+ if response.status_code in {200, 400}:
+ envelope = _Envelope.model_validate(response.json())
+ if any(error.extensions.code == "RATELIMITED" for error in envelope.errors):
+ raw_reset = response.headers.get("X-RateLimit-Requests-Reset")
+ try:
+ delay = max(1.0, float(raw_reset) / 1000 - time.time()) if raw_reset else 60.0
+ except ValueError:
+ delay = 60.0
+ raise SourceRateLimitError(source_short_name="linear", retry_after=delay)
raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ response, source_short_name="linear", token_provider_kind=self.auth.provider_kind
)
-
- return body
-
- # ------------------------------------------------------------------
- # Pagination
- # ------------------------------------------------------------------
-
- async def _paginated_query(
- self,
- query_template: str,
- process_node_func,
- page_size: int = 50,
- entity_type: str = "items",
- ) -> AsyncGenerator:
- """Paginate through Linear GraphQL, yielding entities from process_node_func.
-
- SourceAuthError propagates immediately (abort sync).
- Other errors propagate to the caller (typically _generate_entities_safe).
- """
- has_next_page = True
- after_cursor = None
- items_processed = 0
-
- while has_next_page:
- pagination = f"first: {page_size}"
- if after_cursor:
- pagination += f', after: "{after_cursor}"'
-
- query = query_template.format(pagination=pagination)
- response = await self._post(query)
-
- data = response.get("data", {})
- collection_key = next(iter(data.keys()), None)
-
- if not collection_key:
- self.logger.warning(f"Unexpected response structure: {response}")
- break
-
- collection_data = data[collection_key]
- nodes = collection_data.get("nodes", [])
-
- batch_count = len(nodes)
- items_processed += batch_count
- self.logger.debug(
- f"Processing batch of {batch_count} {entity_type} (total: {items_processed})"
+ envelope = _Envelope.model_validate(response.json())
+ if envelope.errors or envelope.data is None:
+ # Never log potentially sensitive provider error messages or accept partial data.
+ raise SourceError(
+ "Linear GraphQL did not return complete data", source_short_name="linear"
)
+ return envelope.data
- for node in nodes:
- async for entity in process_node_func(node):
- if entity:
- yield entity
-
- page_info = collection_data.get("pageInfo", {})
- has_next_page = page_info.get("hasNextPage", False)
- after_cursor = page_info.get("endCursor")
-
- if not nodes or not has_next_page:
+ async def validate(self) -> None:
+ """Verify native workspace identity and current discovery of configured teams."""
+ data = await self._query("identity", {})
+ workspace = _Identity.model_validate(data.get("organization"))
+ self.viewer_id = _Identity.model_validate(data.get("viewer")).id
+ if workspace.id != self.config.workspace_id:
+ raise ValueError("Linear workspace does not match configured workspace")
+ found: set[str] = set()
+ after = None
+ seen: set[str] = set()
+ while True:
+ data = await self._query("teams", {"first": 50, "after": after})
+ page = _Connection.model_validate(data.get("teams"))
+ found.update(str(_Identity.model_validate(node).id) for node in page.nodes)
+ after = self._following(page, after)
+ if after is None:
break
+ if after in seen:
+ raise ValueError("Linear team discovery repeated a cursor")
+ seen.add(after)
+ if not self.team_ids <= found:
+ raise ValueError("One or more selected Linear teams are unavailable")
+
+ @staticmethod
+ def _following(page: _Connection, current: str | None) -> str | None:
+ if not page.pageInfo.hasNextPage:
+ return None
+ following = page.pageInfo.endCursor
+ if not following or following == current:
+ raise ValueError("Linear returned a missing or non-advancing continuation")
+ return following
+
+ def _check_issue(self, payload: JsonValue, expected: str | None = None) -> _Issue:
+ issue = _Issue.model_validate(payload)
+ if expected is not None and str(issue.id) != expected:
+ raise ValueError("Linear returned the wrong issue identity")
+ if str(issue.team.id) not in self.team_ids:
+ raise ScopeAccessLost(
+ "Linear issue moved outside selected teams", removal_reason="scope_removed"
+ )
+ return issue
- # ------------------------------------------------------------------
- # Entity generators
- # ------------------------------------------------------------------
-
- async def _generate_team_entities(
- self, since: Optional[str] = None
- ) -> AsyncGenerator[LinearTeamEntity, None]:
- """Generate entities for all teams in the workspace."""
- filter_clause = f'filter: {{{{ updatedAt: {{{{ gte: "{since}" }}}} }}}}, ' if since else ""
- query_template = (
- """
- {{
- teams("""
- + filter_clause
- + """{pagination}) {{
- nodes {{
- id
- name
- key
- description
- color
- icon
- private
- timezone
- createdAt
- updatedAt
- parent {{
- id
- name
- }}
- issueCount
- }}
- pageInfo {{
- hasNextPage
- endCursor
- }}
- }}
- }}
- """
- )
-
- async def process_team(team):
- self.logger.debug(f"Processing team: {team.get('name')} ({team.get('key')})")
- yield LinearTeamEntity.from_api(team)
-
- async for entity in self._paginated_query(
- query_template, process_team, entity_type="teams"
- ):
- yield entity
-
- async def _generate_project_entities(
- self, since: Optional[str] = None
- ) -> AsyncGenerator[LinearProjectEntity, None]:
- """Generate entities for all projects in the workspace."""
- filter_clause = f'filter: {{{{ updatedAt: {{{{ gte: "{since}" }}}} }}}}, ' if since else ""
- query_template = (
- """
- {{
- projects("""
- + filter_clause
- + """{pagination}) {{
- nodes {{
- id
- name
- slugId
- description
- priority
- startDate
- targetDate
- state
- createdAt
- updatedAt
- completedAt
- startedAt
- progress
- teams {{
- nodes {{
- id
- name
- }}
- }}
- lead {{
- name
- }}
- }}
- pageInfo {{
- hasNextPage
- endCursor
- }}
- }}
- }}
- """
- )
-
- async def process_project(project):
- self.logger.debug(f"Processing project: {project.get('name')}")
- yield LinearProjectEntity.from_api(project)
-
- async for entity in self._paginated_query(
- query_template, process_project, entity_type="projects"
- ):
- yield entity
-
- async def _generate_user_entities(
- self, since: Optional[str] = None
- ) -> AsyncGenerator[LinearUserEntity, None]:
- """Generate entities for all users in the workspace."""
- filter_clause = f'filter: {{{{ updatedAt: {{{{ gte: "{since}" }}}} }}}}, ' if since else ""
- query_template = (
- """
- {{
- users("""
- + filter_clause
- + """{pagination}) {{
- nodes {{
- id
- name
- displayName
- email
- avatarUrl
- description
- timezone
- active
- admin
- guest
- lastSeen
- statusEmoji
- statusLabel
- statusUntilAt
- createdIssueCount
- createdAt
- updatedAt
- teams {{
- nodes {{
- id
- name
- key
- }}
- }}
- }}
- pageInfo {{
- hasNextPage
- endCursor
- }}
- }}
- }}
- """
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Keep the established native container identity for this flat provider."""
+ return CompletedScope(
+ record_type=record_type, container_id=parent.identity.native_id, parent=parent.identity
)
- async def process_user(user):
- self.logger.debug(f"Processing user: {user.get('name')} ({user.get('displayName')})")
- yield LinearUserEntity.from_api(user)
-
- async for entity in self._paginated_query(
- query_template, process_user, entity_type="users"
- ):
- yield entity
-
- # ------------------------------------------------------------------
- # Issues, comments, and attachments
- # ------------------------------------------------------------------
-
- async def _process_issue_comments(
+ async def capture_page(
self,
- comments: List[Dict],
- issue: LinearIssueEntity,
- issue_breadcrumbs: List[Breadcrumb],
- ) -> AsyncGenerator[LinearCommentEntity, None]:
- """Yield comment entities for an issue."""
- for comment in comments:
- if not comment.get("body", "").strip():
- continue
- yield LinearCommentEntity.from_api(
- comment,
- issue_id=issue.issue_id,
- issue_identifier=issue.identifier,
- breadcrumbs=issue_breadcrumbs,
- team_id=issue.team_id,
- team_name=issue.team_name,
- project_id=issue.project_id,
- project_name=issue.project_name,
- )
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Fetch exactly one whole-scope page; engine commits it and its continuation."""
+ progress = _Progress.model_validate(continuation.value)
+ variables: dict[str, JsonValue] = {"first": 50, "after": progress.after}
+ if scope.record_type == "issue" and scope.container_id is None:
+ variables["teamIds"] = sorted(self.team_ids)
+ data = await self._query("issues", variables)
+ page = _Connection.model_validate(data.get("issues"))
+ for node in page.nodes:
+ issue = _Issue.model_validate(node)
+ if str(issue.team.id) not in self.team_ids:
+ raise ValueError("Linear root page contains an unselected team")
+ elif scope.record_type in {"comment", "attachment"} and scope.container_id:
+ UUID(scope.container_id)
+ variables["issueId"] = scope.container_id
+ name = "comments" if scope.record_type == "comment" else "attachments"
+ data = await self._query(name, variables)
+ issue_payload = self._single_issue(data, scope.container_id)
+ if issue_payload is None:
+ raise ScopeAccessLost(
+ "Linear issue is unavailable in the current accessible scope",
+ removal_reason="scope_removed",
+ )
+ self._check_issue(issue_payload, scope.container_id)
+ issue_data = _ChildConnection.model_validate(issue_payload)
+ page = issue_data.comments if name == "comments" else issue_data.attachments
+ if page is None:
+ raise ValueError("Linear omitted the requested child connection")
+ for node in page.nodes:
+ if str(_Child.model_validate(node).issue.id) != scope.container_id:
+ raise ValueError("Linear child belongs to a different issue")
+ else:
+ raise ValueError("Unsupported Linear scope")
+ following = self._following(page, progress.after)
+ return CapturePage(
+ records=tuple([await self._retain(scope, node, files) for node in page.nodes]),
+ continuation=ScanContinuation(value=_Progress(after=following).model_dump()),
+ final=following is None,
+ )
- async def _download_description_attachment(
+ async def _retain(
self,
- entity: LinearAttachmentEntity,
+ scope: CompletedScope,
+ payload: dict[str, JsonValue],
files: FileService,
- ) -> LinearAttachmentEntity | None:
- """Download an inline attachment via FileService. Returns None on expected skips.
-
- 401 after refresh propagates (token is dead → abort sync).
- Infrastructure failures (IOError, OSError) propagate.
- Other HTTP errors skip the file.
- """
- try:
- await files.download_from_url(
- entity=entity,
- client=self.http_client,
- auth=self.auth,
- logger=self.logger,
+ ) -> CaptureRecord:
+ record = self._capture(scope, payload)
+ if scope.record_type != "attachment":
+ return record
+ url = payload.get("url")
+ if not isinstance(url, str) or not self._is_upload(url):
+ return record
+ headers = await authorization_headers(self.auth)
+ async with self.http_client.stream(
+ "GET",
+ url,
+ headers=headers,
+ follow_redirects=False,
+ timeout=httpx.Timeout(180.0, read=540.0),
+ ) as response:
+ # Redirects are errors, never permission to forward auth to another destination.
+ raise_for_status(
+ response, source_short_name="linear", token_provider_kind=self.auth.provider_kind
)
- if not entity.local_path:
- self.logger.warning(f"Download produced no local path for {entity.name}")
- return None
- return entity
- except FileSkippedException as e:
- self.logger.debug(f"Skipping attachment {entity.attachment_id}: {e.reason}")
- return None
- except httpx.HTTPStatusError as e:
- if e.response.status_code == 401:
- raise
- self.logger.warning(
- f"HTTP {e.response.status_code} downloading attachment {entity.attachment_id}: {e}"
+ length = response.headers.get("Content-Length")
+ if length is not None and int(length) > MAX_FILE_SIZE_BYTES:
+ return record
+ content = bytearray()
+ async with aclosing(response.aiter_bytes()) as chunks:
+ async for chunk in chunks:
+ if len(content) + len(chunk) > MAX_FILE_SIZE_BYTES:
+ return record
+ content.extend(chunk)
+ media_type = (
+ response.headers.get("Content-Type", "application/octet-stream")
+ .split(";", 1)[0]
+ .strip()
)
- return None
-
- async def _generate_attachment_entities_from_description(
- self,
- issue_id: str,
- issue_identifier: str,
- issue_description: str,
- breadcrumbs: List[Breadcrumb],
- files: FileService,
- ) -> AsyncGenerator[LinearAttachmentEntity, None]:
- """Extract and download attachments from markdown links in issue descriptions."""
- if not issue_description:
- return
-
- markdown_link_pattern = r"\[([^\]]+)\]\(([^)]+)\)"
- matches = re.findall(markdown_link_pattern, issue_description)
-
- self.logger.debug(
- f"Found {len(matches)} potential attachments in description "
- f"for issue {issue_identifier}"
+ ref = await files.store_canonical_blob(bytes(content), media_type=media_type)
+ return record.model_copy(
+ update={
+ "blobs": (ref.model_copy(update={"source_path": "/url"}),),
+ "completeness": "complete",
+ }
)
- for file_name, url in matches:
- if "uploads.linear.app" not in url:
- continue
-
- self.logger.debug(f"Processing attachment from description: {file_name} - URL: {url}")
-
- attachment_id = str(uuid4())
- mime_type = mimetypes.guess_type(file_name)[0]
- if mime_type and "/" in mime_type:
- file_type = mime_type.split("/")[0]
- else:
- ext = os.path.splitext(file_name)[1].lower().lstrip(".")
- file_type = ext if ext else "file"
-
- attachment_entity = LinearAttachmentEntity(
- entity_id=attachment_id,
- breadcrumbs=breadcrumbs.copy(),
- name=file_name,
- created_at=None,
- updated_at=None,
- url=url,
- size=0,
- file_type=file_type,
- mime_type=mime_type or "application/octet-stream",
- local_path=None,
- attachment_id=attachment_id,
- issue_id=issue_id,
- issue_identifier=issue_identifier,
- title=file_name,
- subtitle="Extracted from issue description",
- source={"type": "description_link"},
- web_url_value=url,
+ @staticmethod
+ def _is_upload(url: str) -> bool:
+ try:
+ parsed = urlsplit(url)
+ return (
+ parsed.scheme == "https"
+ and parsed.hostname == "uploads.linear.app"
+ and parsed.port in {None, 443}
+ and parsed.username is None
+ and parsed.password is None
+ and bool(parsed.path.strip("/"))
+ and not parsed.fragment
)
-
- downloaded = await self._download_description_attachment(attachment_entity, files)
- if downloaded:
- yield downloaded
-
- async def _generate_issue_entities(
- self,
- since: Optional[str] = None,
- files: FileService | None = None,
- ) -> AsyncGenerator[
- Union[LinearIssueEntity, LinearCommentEntity, LinearAttachmentEntity], None
- ]:
- """Generate entities for all issues, their comments, and their attachments."""
- updated_filter = f', updatedAt: {{{{ gte: "{since}" }}}}' if since else ""
- query_template = (
- """
- {{
- issues(filter: {{ archivedAt: {{ null: true }}"""
- + updated_filter
- + """ }}, {pagination}) {{
- nodes {{
- id
- identifier
- title
- description
- priority
- completedAt
- createdAt
- updatedAt
- dueDate
- archivedAt
- state {{
- name
- }}
- team {{
- id
- name
- }}
- project {{
- id
- name
- }}
- assignee {{
- name
- }}
- comments {{
- nodes {{
- id
- body
- createdAt
- updatedAt
- user {{
- id
- name
- }}
- }}
- }}
- }}
- pageInfo {{
- hasNextPage
- endCursor
- }}
- }}
- }}
- """
+ except ValueError:
+ return False
+
+ @staticmethod
+ def _capture(scope: CompletedScope, payload: dict[str, JsonValue]) -> CaptureRecord:
+ identity = _Identity.model_validate(payload)
+ dates = _Dates.model_validate(payload)
+ text = payload.get("description") if scope.record_type == "issue" else payload.get("body")
+ # Retain links and native text; no claim that embedded or attachment bytes are stored.
+ partial_content = scope.record_type == "attachment" or (
+ isinstance(text, str) and "uploads.linear.app" in text
+ )
+ return CaptureRecord(
+ identity=RecordIdentity(
+ record_type=scope.record_type,
+ native_id=str(identity.id),
+ container_id=scope.container_id,
+ ),
+ payload=payload,
+ payload_schema_version=FIELD_SET_VERSION,
+ completeness="partial" if partial_content else "complete",
+ source_created_at=dates.createdAt,
+ source_updated_at=dates.updatedAt,
+ observed_at=datetime.now(timezone.utc),
)
- async def process_issue(data):
- if data.get("archivedAt"):
- self.logger.warning(
- f"Archived issue {data.get('identifier')} passed GraphQL filter — skipping"
- )
- return
-
- self.logger.debug(f"Processing issue: {data.get('identifier')} — '{data.get('title')}'")
-
- issue = LinearIssueEntity.from_api(data)
- yield issue
-
- issue_breadcrumbs = issue.breadcrumbs + [
- Breadcrumb(
- entity_id=issue.issue_id,
- name=issue.title,
- entity_type=LinearIssueEntity.__name__,
- )
- ]
-
- comments = data.get("comments", {}).get("nodes", [])
- self.logger.debug(f"Processing {len(comments)} comments for issue {issue.identifier}")
- async for comment in self._process_issue_comments(comments, issue, issue_breadcrumbs):
- yield comment
-
- if issue.description and files:
- async for attachment in self._generate_attachment_entities_from_description(
- issue.issue_id,
- issue.identifier,
- issue.description,
- issue_breadcrumbs,
- files,
- ):
- yield attachment
-
- async for entity in self._paginated_query(
- query_template, process_issue, entity_type="issues"
- ):
- yield entity
-
- # ------------------------------------------------------------------
- # Error isolation
- # ------------------------------------------------------------------
-
- async def _generate_entities_safe(
- self,
- generator: AsyncGenerator,
- entity_type: str,
- ) -> AsyncGenerator:
- """Yield entities with error isolation per entity type.
-
- SourceAuthError propagates (abort sync — credentials dead).
- Other exceptions are logged and stop that entity type only.
- """
- try:
- self.logger.debug(f"Starting {entity_type} entity generation")
- async for entity in generator:
- yield entity
- except SourceAuthError:
- raise
- except SourceError as exc:
- self.logger.warning(f"Failed to generate {entity_type} entities: {exc}")
- except Exception as exc:
- self.logger.warning(f"Unexpected error generating {entity_type} entities: {exc}")
-
- # ------------------------------------------------------------------
- # Main entry point
- # ------------------------------------------------------------------
+ @staticmethod
+ def _single_issue(data: dict[str, JsonValue], native_id: str) -> dict[str, JsonValue] | None:
+ """Interpret only a successful exhausted exact-ID connection, never failed lookups."""
+ page = _Connection.model_validate(data.get("issues"))
+ if page.pageInfo.hasNextPage or len(page.nodes) > 1:
+ raise ValueError("Linear exact-ID membership response is not complete and unique")
+ if not page.nodes:
+ return None
+ node = page.nodes[0]
+ if str(_Identity.model_validate(node).id) != native_id:
+ raise ValueError("Linear returned the wrong issue identity")
+ return node
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Fail closed on ambiguous absence; readable out-of-scope roots may be removed."""
+ native_id = record.identity.native_id
+ UUID(native_id)
+ data = await self._query("issue-membership", {"issueId": native_id})
+ payload = self._single_issue(data, native_id)
+ if payload is None:
+ return
+ issue = _Issue.model_validate(payload)
+ if str(issue.team.id) in self.team_ids:
+ raise ValueError("Linear omitted an issue still readable in the selected scope")
async def generate_entities(
self,
@@ -617,48 +414,11 @@ async def generate_entities(
files: FileService | None = None,
node_selections: list[NodeSelectionData] | None = None,
) -> AsyncGenerator[BaseEntity, None]:
- """Generate all entities from Linear.
-
- On first sync (no cursor), all entities are fetched. On subsequent syncs,
- only entities updated since the last sync are fetched.
- """
- cursor_data = cursor.data if cursor else {}
- last_synced_at = cursor_data.get("last_synced_at") or None
- sync_start = datetime.now(tz=timezone.utc).isoformat()
-
- if last_synced_at:
- self.logger.info(f"Incremental sync from {last_synced_at}")
- else:
- self.logger.info("Full sync (first run)")
+ """Legacy crawling is replaced by canonical page capture and offline projection."""
+ raise NotImplementedError("Linear requires canonical page capture")
+ yield # pragma: no cover
- async for entity in self._generate_entities_safe(
- self._generate_team_entities(since=last_synced_at), "team"
- ):
- yield entity
- async for entity in self._generate_entities_safe(
- self._generate_project_entities(since=last_synced_at), "project"
- ):
- yield entity
-
- async for entity in self._generate_entities_safe(
- self._generate_user_entities(since=last_synced_at), "user"
- ):
- yield entity
-
- async for entity in self._generate_entities_safe(
- self._generate_issue_entities(since=last_synced_at, files=files),
- "issue, comment, and attachment",
- ):
- yield entity
-
- if cursor:
- cursor.update(last_synced_at=sync_start)
-
- # ------------------------------------------------------------------
- # Validation
- # ------------------------------------------------------------------
-
- async def validate(self) -> None:
- """Verify Linear OAuth2 token by POSTing a minimal GraphQL query."""
- await self._post("query { viewer { id } }")
+class _ChildConnection(_Model):
+ comments: _Connection | None = None
+ attachments: _Connection | None = None
diff --git a/backend/airweave/platform/sources/linear_queries/attachments.graphql b/backend/airweave/platform/sources/linear_queries/attachments.graphql
new file mode 100644
index 0000000000..034702fa0e
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/attachments.graphql
@@ -0,0 +1,12 @@
+query OwnedLinearIssueAttachments($issueId: ID!, $first: Int!, $after: String) {
+ issues(filter: {id: {eq: $issueId}}, includeArchived: true, first: 2) {
+ nodes {
+ id team { id }
+ attachments(first: $first, after: $after, includeArchived: true, orderBy: createdAt) {
+ nodes { id title url metadata createdAt updatedAt archivedAt issue { id identifier title url team { id name } project { id name } } }
+ pageInfo { hasNextPage endCursor }
+ }
+ }
+ pageInfo { hasNextPage endCursor }
+ }
+}
diff --git a/backend/airweave/platform/sources/linear_queries/comments.graphql b/backend/airweave/platform/sources/linear_queries/comments.graphql
new file mode 100644
index 0000000000..ef6ac1b97c
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/comments.graphql
@@ -0,0 +1,17 @@
+query OwnedLinearIssueComments($issueId: ID!, $first: Int!, $after: String) {
+ issues(filter: {id: {eq: $issueId}}, includeArchived: true, first: 2) {
+ nodes {
+ id team { id }
+ comments(first: $first, after: $after, includeArchived: true, orderBy: createdAt) {
+ nodes {
+ id body url createdAt updatedAt archivedAt resolvedAt
+ user { id name updatedAt }
+ parent { id }
+ issue { id identifier title url team { id name } project { id name } }
+ }
+ pageInfo { hasNextPage endCursor }
+ }
+ }
+ pageInfo { hasNextPage endCursor }
+ }
+}
diff --git a/backend/airweave/platform/sources/linear_queries/identity.graphql b/backend/airweave/platform/sources/linear_queries/identity.graphql
new file mode 100644
index 0000000000..a72d843d35
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/identity.graphql
@@ -0,0 +1,4 @@
+query OwnedLinearIdentity {
+ viewer { id name }
+ organization { id name }
+}
diff --git a/backend/airweave/platform/sources/linear_queries/issue-membership.graphql b/backend/airweave/platform/sources/linear_queries/issue-membership.graphql
new file mode 100644
index 0000000000..69418b8c2f
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/issue-membership.graphql
@@ -0,0 +1,6 @@
+query OwnedLinearIssueMembership($issueId: ID!) {
+ issues(filter: {id: {eq: $issueId}}, includeArchived: true, first: 2) {
+ nodes { id team { id } archivedAt trashed }
+ pageInfo { hasNextPage endCursor }
+ }
+}
diff --git a/backend/airweave/platform/sources/linear_queries/issues.graphql b/backend/airweave/platform/sources/linear_queries/issues.graphql
new file mode 100644
index 0000000000..f4266be157
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/issues.graphql
@@ -0,0 +1,15 @@
+query OwnedLinearIssues($teamIds: [ID!]!, $first: Int!, $after: String) {
+ issues(filter: {team: {id: {in: $teamIds}}}, first: $first, after: $after,
+ includeArchived: true, orderBy: createdAt) {
+ nodes {
+ id identifier title description url createdAt updatedAt archivedAt trashed
+ priority dueDate
+ team { id name key visibility archivedAt updatedAt }
+ project { id name updatedAt archivedAt }
+ assignee { id name updatedAt }
+ creator { id name updatedAt }
+ state { id name type updatedAt }
+ }
+ pageInfo { hasNextPage endCursor }
+ }
+}
diff --git a/backend/airweave/platform/sources/linear_queries/teams.graphql b/backend/airweave/platform/sources/linear_queries/teams.graphql
new file mode 100644
index 0000000000..82d77c9ef3
--- /dev/null
+++ b/backend/airweave/platform/sources/linear_queries/teams.graphql
@@ -0,0 +1,6 @@
+query OwnedLinearTeamChoices($first: Int!, $after: String) {
+ teams(first: $first, after: $after, includeArchived: true, orderBy: createdAt) {
+ nodes { id name key visibility archivedAt updatedAt }
+ pageInfo { hasNextPage endCursor }
+ }
+}
diff --git a/backend/airweave/platform/sources/monday.py b/backend/airweave/platform/sources/monday.py
index a9e5c2e75b..59d791c7a9 100644
--- a/backend/airweave/platform/sources/monday.py
+++ b/backend/airweave/platform/sources/monday.py
@@ -30,12 +30,12 @@
MondayUpdateEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/notion.py b/backend/airweave/platform/sources/notion.py
index 57c2856713..040a4924da 100644
--- a/backend/airweave/platform/sources/notion.py
+++ b/backend/airweave/platform/sources/notion.py
@@ -1,45 +1,201 @@
-"""Notion connector for syncing content from Notion workspaces to Airweave.
+"""Native Notion originals; incomplete discovery never establishes absence.
-This module provides a comprehensive source implementation for extracting databases and pages
-with full aggregated content from Notion, handling API rate limits, and converting API
-responses to entity objects.
+Page/database/data-source topology entries describe enumeration scopes, not grants.
+These originals are independently retrieved roots. Blocks alone inherit access from
+native owners. Properties, comments and file bytes are not complete in this slice.
"""
from __future__ import annotations
-from typing import Any, AsyncGenerator, Dict, List, Optional, Set, Tuple
+import hashlib
+import json
+from collections.abc import AsyncGenerator, Sequence
+from datetime import datetime, timezone
+from functools import partial
+from typing import Annotated, Literal
+from uuid import UUID
-import httpx
-from tenacity import retry, stop_after_attempt
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import SourceAuthError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
-from airweave.domains.storage import FileSkippedException
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ InvalidScanContinuation,
+ ScopeAccessLost,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceAuthError, SourceError
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage.file_service import FileService
+from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.config import NotionConfig
from airweave.platform.decorators import source
-from airweave.platform.entities._base import BaseEntity, Breadcrumb
-from airweave.platform.entities.notion import (
- NotionDatabaseEntity,
- NotionFileEntity,
- NotionPageEntity,
- NotionPropertyEntity,
- _extract_page_title,
- _extract_rich_text_plain,
-)
+from airweave.platform.entities._base import BaseEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
+API_VERSION = "2026-03-11"
+FIELD_SET_VERSION = 2
+_ROOTS = ("page", "database", "data_source")
+_ENDPOINTS = {
+ "page": "pages",
+ "database": "databases",
+ "data_source": "data_sources",
+ "block": "blocks",
+}
+_PAGE_SIZE = 25
+
+
+class _Native(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+
+
+class _Object(_Native):
+ id: UUID
+ object: Literal["page", "database", "data_source", "block"]
+ created_time: AwareDatetime
+ last_edited_time: AwareDatetime
+ in_trash: StrictBool = False
+
+
+class _Parent(_Native):
+ type: str
+ page_id: UUID | None = None
+ block_id: UUID | None = None
+ database_id: UUID | None = None
+ data_source_id: UUID | None = None
+
+
+class _Located(_Object):
+ parent: _Parent
+
+
+class _Block(_Located):
+ type: str
+ has_children: StrictBool
+ synced_block: dict[str, JsonValue] | None = None
+
+
+class _Database(_Object):
+ data_sources: list[dict[str, JsonValue]] = Field(max_length=1000)
+
+
+class _Id(_Native):
+ id: UUID
+
+
+class _BotWorkspace(_Native):
+ workspace_id: UUID
+
+
+class _Principal(_Id):
+ object: Literal["user"]
+ type: Literal["bot"]
+ bot: _BotWorkspace
+
+
+class _Status(_Native):
+ type: Literal["complete", "incomplete"]
+ incomplete_reason: str | None = None
+
+
+class _List(_Native):
+ object: Literal["list"]
+ results: list[dict[str, JsonValue]] = Field(max_length=_PAGE_SIZE)
+ has_more: StrictBool
+ next_cursor: str | None
+ request_status: _Status | None = None
+
+
+class _Property(_Native):
+ # Native IDs already contain percent encoding. Accept one URL path segment,
+ # including valid escapes, and never quote it a second time.
+ id: str = Field(pattern=r"^(?:[A-Za-z0-9._~-]|%[0-9A-Fa-f]{2})+$")
+ type: str
+
+
+class _PageProperties(_Object):
+ properties: dict[str, dict[str, JsonValue]]
+
+
+class _Calculation(_Native):
+ type: str
+
+
+class _PropertyMetadata(_Property):
+ rollup: _Calculation | None = None
+ formula: _Calculation | None = None
+
+
+class _PropertyItem(_PropertyMetadata):
+ object: Literal["property_item"]
+
+
+class _PropertyList(_List):
+ type: Literal["property_item"]
+ property_item: _PropertyMetadata
+
+
+class _PropertyArchiveHeader(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ format_version: Literal[1] = 1
+ page_id: UUID
+ property_id: str
+ notion_version: Literal["2026-03-11"] = API_VERSION
+ page_last_edited_time: AwareDatetime
+
+
+class _PropertyManifest(_PropertyArchiveHeader):
+ name: str
+ property: dict[str, JsonValue]
+ response_count: int = Field(ge=1)
+ value_status: Literal["available", "unsupported"]
+
+
+class _PropertyProgress(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ index: int = Field(default=0, ge=0)
+ inventory: str | None = Field(default=None, pattern=r"^[a-f0-9]{64}$")
+
+
+class _Error(_Native):
+ code: str
+
+
+class _Progress(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ version: Literal[1] = 1
+ cursor: str | None = Field(default=None, min_length=1, max_length=4096)
+ offset: int = Field(default=0, ge=0, le=1000)
+ window_start: AwareDatetime | None = None
+ last_created: AwareDatetime | None = None
+ capped: bool = False
+ cursor_hashes: tuple[Annotated[str, Field(pattern=r"^[a-f0-9]{64}$")], ...] = Field(
+ default=(), max_length=512
+ )
+
@source(
name="Notion",
@@ -57,1223 +213,580 @@
rate_limit_level=RateLimitLevel.CONNECTION,
)
class NotionSource(BaseSource):
- """Notion source connector integrates with the Notion API to extract and synchronize content.
-
- Connects to your Notion workspace.
-
- It provides comprehensive access to databases, pages, and content with advanced content
- aggregation, lazy loading, and file processing capabilities for optimal performance.
- """
-
- TIMEOUT_SECONDS = 60.0
-
- class NotionAccessError(Exception):
- """Non-retryable access/shape error from Notion.
-
- Raised for permission or structural cases that won't succeed on retry
- (e.g., 403/404, or specific 400 validation for linked views).
- """
-
- def __init__(self, status: int, message: str, url: str):
- """Initialize with HTTP status, message, and request URL."""
- super().__init__(f"HTTP {status} - {message} ({url})")
- self.status = status
- self.message = message
- self.url = url
+ """One source registration, native acquisition, and existing SQL-owned page progress."""
+
+ canonical_record_types = (*_ROOTS, "block", "page_property")
+ # Root alternatives remain independent access authorities. Child alternatives
+ # represent native enumeration operations; they emit discovered roots, not edges.
+ canonical_container_parents = {
+ "page": (None, "data_source"),
+ "database": (None, "data_source"),
+ "data_source": (None, "database"),
+ "block": ("page", "block"),
+ "page_property": "page",
+ }
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Versioned acquisition contract; changing it requires explicit cycle restart."""
+ return self._cycle_configuration
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: NotionConfig,
) -> NotionSource:
- """Create a new Notion source."""
- return cls(auth=auth, logger=logger, http_client=http_client)
-
- async def validate(self) -> None:
- """Validate credentials by pinging the Notion current-user endpoint."""
- await self._get("https://api.notion.com/v1/users/me")
-
- # ------------------------------------------------------------------
- # HTTP helpers
- # ------------------------------------------------------------------
+ """Build a native managed/direct-auth source without acquiring additional grants."""
+ instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.config = config
+ fingerprint = (
+ f"notion:{API_VERSION}:{FIELD_SET_VERSION}:partial-native-roots-blocks-properties:"
+ f"{config.expected_workspace_id}:{config.expected_bot_id}"
+ )
+ instance._cycle_configuration = CycleConfiguration.from_source(
+ fingerprint=hashlib.sha256(fingerprint.encode()).hexdigest(),
+ record_types=cls.canonical_record_types,
+ container_parents=cls.canonical_container_parents,
+ completion_policies={kind: "discovery_with_validation" for kind in _ROOTS},
+ )
+ return instance
@retry(
stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
reraise=True,
)
- async def _get(self, url: str) -> dict:
- """Make an authenticated GET request to the Notion API."""
- token = await self.auth.get_token()
- headers = {
- "Authorization": f"Bearer {token}",
- "Notion-Version": "2022-06-28",
- }
-
- response = await self.http_client.get(url, headers=headers, timeout=self.TIMEOUT_SECONDS)
-
+ async def _request(
+ self,
+ path: str,
+ body: dict[str, JsonValue] | None = None,
+ *,
+ params: dict[str, str | int] | None = None,
+ ) -> dict[str, JsonValue] | None:
+ headers = {**await authorization_headers(self.auth), "Notion-Version": API_VERSION}
+ method = self.http_client.get if body is None else self.http_client.post
+ kwargs = {"params": params} if body is None else {"json": body}
+ response = await method(f"https://api.notion.com/v1/{path}", headers=headers, **kwargs)
if response.status_code == 401 and self.auth.supports_refresh:
- self.logger.warning("Received 401 from Notion — attempting token refresh")
- new_token = await self.auth.force_refresh()
- headers["Authorization"] = f"Bearer {new_token}"
- response = await self.http_client.get(
- url, headers=headers, timeout=self.TIMEOUT_SECONDS
- )
-
- if response.status_code != 200:
- msg = ""
- try:
- body_json = response.json()
- if isinstance(body_json, dict):
- msg = str(body_json.get("message", ""))
- except Exception:
- pass
-
- is_pages = "/v1/pages/" in url
- is_databases = "/v1/databases/" in url
-
- if is_pages and response.status_code in (403, 404):
- raise NotionSource.NotionAccessError(
- response.status_code, msg or "access denied", url
- )
- if is_databases and (
- response.status_code in (403, 404)
- or (
- response.status_code == 400
- and "does not contain any data sources" in (msg or "").lower()
- )
- ):
- raise NotionSource.NotionAccessError(
- response.status_code, msg or "validation_error", url
- )
-
+ headers = {
+ **await authorization_headers(self.auth, refresh=True),
+ "Notion-Version": API_VERSION,
+ }
+ response = await method(f"https://api.notion.com/v1/{path}", headers=headers, **kwargs)
+ if (
+ response.status_code == 404
+ and _Error.model_validate(response.json()).code == "object_not_found"
+ ):
+ return None
raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ response, source_short_name="notion", token_provider_kind=self.auth.provider_kind
)
- return response.json()
+ result = response.json()
+ if not isinstance(result, dict):
+ raise ValueError("Notion returned a non-object response")
+ return result
- @retry(
- stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
- reraise=True,
- )
- async def _post(self, url: str, json_data: dict) -> dict:
- """Make an authenticated POST request to the Notion API."""
- token = await self.auth.get_token()
- headers = {
- "Authorization": f"Bearer {token}",
- "Notion-Version": "2022-06-28",
- }
-
- response = await self.http_client.post(
- url, headers=headers, json=json_data, timeout=self.TIMEOUT_SECONDS
- )
-
- if response.status_code == 401 and self.auth.supports_refresh:
- self.logger.warning("Received 401 from Notion POST — attempting token refresh")
- new_token = await self.auth.force_refresh()
- headers["Authorization"] = f"Bearer {new_token}"
- response = await self.http_client.post(
- url, headers=headers, json=json_data, timeout=self.TIMEOUT_SECONDS
+ async def validate(self) -> None:
+ """Verify identity independently of content discovery capabilities."""
+ payload = await self._request("users/me")
+ principal = _Principal.model_validate(payload)
+ if (principal.bot.workspace_id, principal.id) != (
+ self.config.expected_workspace_id,
+ self.config.expected_bot_id,
+ ):
+ raise SourceAuthError(
+ "Notion workspace or connection bot does not match the configured account",
+ source_short_name="notion",
+ status_code=200,
+ token_provider_kind=self.auth.provider_kind,
)
- raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
+ async def _retrieve(self, kind: str, native_id: str) -> dict[str, JsonValue] | None:
+ native_id = str(UUID(native_id))
+ payload = await self._request(f"{_ENDPOINTS[kind]}/{native_id}")
+ if payload is not None:
+ parsed = _Object.model_validate(payload)
+ if parsed.object != kind or str(parsed.id) != native_id:
+ raise ValueError("Notion returned the wrong exact object identity")
+ return payload
+
+ @staticmethod
+ def _record(
+ kind: str, payload: dict[str, JsonValue], *, parent: RecordIdentity | None = None
+ ) -> CaptureRecord:
+ parsed = _Object.model_validate(payload)
+ if parsed.object != kind:
+ raise ValueError("Notion returned an unexpected object kind")
+ return CaptureRecord(
+ identity=RecordIdentity(record_type=kind, native_id=str(parsed.id)),
+ parent=parent,
+ allow_reparent=parent is not None,
+ payload=payload,
+ payload_schema_version=FIELD_SET_VERSION,
+ completeness="partial",
+ source_created_at=parsed.created_time,
+ source_updated_at=parsed.last_edited_time,
+ observed_at=datetime.now(timezone.utc),
+ kind="delete" if parsed.in_trash else "upsert",
+ removal_reason="provider_deleted" if parsed.in_trash else None,
)
- return response.json()
-
- # ------------------------------------------------------------------
- # Search / query helpers
- # ------------------------------------------------------------------
-
- async def _search_objects( # noqa: C901
- self, object_type: str
- ) -> AsyncGenerator[dict, None]:
- """Search for objects of a specific type, excluding archived objects."""
- url = "https://api.notion.com/v1/search"
- has_more = True
- start_cursor = None
-
- total_found = 0
- total_filtered = 0
-
- while has_more:
- json_data: Dict[str, Any] = {
- "filter": {"property": "object", "value": object_type},
- "page_size": 100,
- }
- if start_cursor:
- json_data["start_cursor"] = start_cursor
-
- try:
- response = await self._post(url, json_data)
- results = response.get("results", [])
- total_found += len(results)
-
- self.logger.debug(f"Search returned {len(results)} {object_type}(s) in this batch")
-
- filtered_count = 0
- for obj in results:
- is_archived = obj.get("archived", False)
- is_trashed = obj.get("in_trash", False)
-
- if is_archived or is_trashed:
- filtered_count += 1
- total_filtered += 1
- else:
- yield obj
-
- if filtered_count > 0:
- self.logger.warning(
- f"Filtered {filtered_count} {object_type}(s) in this batch "
- f"(archived or trashed)"
- )
- has_more = response.get("has_more", False)
- start_cursor = response.get("next_cursor")
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Search access issue for {object_type}: {e}")
- raise
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error searching for {object_type}: {e}")
- raise
-
- yielded_count = total_found - total_filtered
- self.logger.debug(
- f"Search complete for {object_type}: "
- f"found={total_found}, filtered={total_filtered}, yielded={yielded_count}"
+ @staticmethod
+ def _unavailable(record: SourceRecord) -> CaptureRecord:
+ return CaptureRecord(
+ identity=record.identity,
+ parent=record.parent,
+ payload={},
+ kind="delete",
+ removal_reason="scope_removed",
+ completeness="partial",
+ observed_at=datetime.now(timezone.utc),
)
- async def _query_database_pages( # noqa: C901
- self, database_id: str
- ) -> AsyncGenerator[dict, None]:
- """Query all pages in a database, excluding archived pages."""
- url = f"https://api.notion.com/v1/databases/{database_id}/query"
- has_more = True
- start_cursor = None
-
- total_found = 0
- total_filtered = 0
-
- while has_more:
- json_data: Dict[str, Any] = {"page_size": 100}
- if start_cursor:
- json_data["start_cursor"] = start_cursor
-
- try:
- response = await self._post(url, json_data)
- results = response.get("results", [])
- total_found += len(results)
-
- filtered_count = 0
- for page in results:
- is_archived = page.get("archived", False)
- is_trashed = page.get("in_trash", False)
-
- if is_archived or is_trashed:
- filtered_count += 1
- total_filtered += 1
- else:
- yield page
-
- if filtered_count > 0:
- self.logger.warning(
- f"Filtered {filtered_count} page(s) in this batch (archived or trashed)"
- )
-
- has_more = response.get("has_more", False)
- start_cursor = response.get("next_cursor")
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Database query access issue for {database_id}: {e}")
- raise
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error querying database {database_id}: {e}")
- raise
-
- yielded_count = total_found - total_filtered
- self.logger.debug(
- f"Database query complete for {database_id}: "
- f"found={total_found}, filtered={total_filtered}, yielded={yielded_count}"
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Search omission is repaired by an exact native read, never inferred deletion."""
+ if record.identity.record_type not in _ROOTS or record.parent is not None:
+ raise ValueError("Notion known-object refresh requires an independent root")
+ payload = await self._retrieve(record.identity.record_type, record.identity.native_id)
+ return (
+ self._unavailable(record)
+ if payload is None
+ else self._record(record.identity.record_type, payload)
)
- # ------------------------------------------------------------------
- # Child database processing
- # ------------------------------------------------------------------
-
- async def _process_child_databases( # noqa: C901
- self, files: FileService | None = None
- ) -> AsyncGenerator[BaseEntity, None]:
- """Process child databases discovered during page content extraction."""
- self.logger.debug("Processing child databases")
-
- while self._child_databases_to_process:
- unprocessed_databases = [
- db_id
- for db_id in self._child_databases_to_process
- if db_id not in self._processed_databases
- ]
-
- if not unprocessed_databases:
- break
-
- self.logger.debug(f"Processing {len(unprocessed_databases)} child databases")
-
- for database_id in unprocessed_databases:
- try:
- self.logger.debug(f"Processing child database: {database_id}")
-
- schema = await self._get(f"https://api.notion.com/v1/databases/{database_id}")
-
- breadcrumbs = self._child_database_breadcrumbs.get(database_id, [])
-
- database_entity = NotionDatabaseEntity.from_api(schema, breadcrumbs=breadcrumbs)
- yield database_entity
- self._processed_databases.add(database_id)
-
- async for page in self._query_database_pages(database_id):
- page_id = page["id"]
- if page_id in self._processed_pages:
- continue
-
- try:
- db_breadcrumb = Breadcrumb(
- entity_id=database_id,
- name=database_entity.title,
- entity_type=NotionDatabaseEntity.__name__,
- )
- page_breadcrumbs = breadcrumbs + [db_breadcrumb]
-
- (
- page_entity,
- file_entities,
- ) = await self._create_comprehensive_page_entity(
- page, page_breadcrumbs, database_id, schema
- )
- yield page_entity
-
- for file_entity in file_entities:
- processed = await self._process_and_yield_file(file_entity, files)
- if processed:
- yield processed
-
- self._processed_pages.add(page_id)
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(
- f"Access issue processing child database page {page_id}: {e}"
- )
- continue
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(
- f"Error processing child database page {page_id}: {e}"
- )
- continue
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(
- f"Child database {database_id} not accessible ({e.status}). Skipping."
- )
- self._processed_databases.add(database_id)
- continue
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing child database {database_id}: {e}")
- self._processed_databases.add(database_id)
- continue
-
- # ------------------------------------------------------------------
- # Breadcrumbs
- # ------------------------------------------------------------------
-
- async def _build_page_breadcrumbs(self, page: dict) -> List[Breadcrumb]:
- """Build breadcrumbs for a page by traversing up the parent hierarchy."""
- breadcrumbs: List[Breadcrumb] = []
- current_page = page
-
- while True:
- parent = current_page.get("parent", {})
- parent_type = parent.get("type", "")
-
- if parent_type == "page_id":
- parent_id = parent.get("page_id")
- try:
- parent_page = await self._get(f"https://api.notion.com/v1/pages/{parent_id}")
- parent_title = _extract_page_title(parent_page) or "Untitled Page"
- breadcrumbs.insert(
- 0,
- Breadcrumb(
- entity_id=parent_id,
- name=parent_title,
- entity_type=NotionPageEntity.__name__,
- ),
- )
- current_page = parent_page
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Could not fetch parent page {parent_id}: {e}")
- break
- elif parent_type == "database_id":
- break
- else:
- break
-
- return breadcrumbs
-
- # ------------------------------------------------------------------
- # Page entity creation with content aggregation
- # ------------------------------------------------------------------
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Only native block-child listings establish their exact-owner absence."""
+ if record.identity.record_type != "block" or record.parent is None:
+ raise ValueError("Notion discovery requires exact known-object refresh")
+ payload = await self._retrieve("block", record.identity.native_id)
+ if payload is None:
+ return
+ block = _Block.model_validate(payload)
+ if block.in_trash or self._block_parent(block) != record.parent:
+ return
+ raise ValueError("Notion omitted a block still readable under its captured owner")
+
+ @staticmethod
+ def _block_parent(block: _Block) -> RecordIdentity:
+ if block.parent.type == "page_id" and block.parent.page_id is not None:
+ return RecordIdentity(record_type="page", native_id=str(block.parent.page_id))
+ if block.parent.type == "block_id" and block.parent.block_id is not None:
+ return RecordIdentity(record_type="block", native_id=str(block.parent.block_id))
+ raise ValueError("Notion block has an unsupported or missing native owner")
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Notion UUIDs are global; do not rewrite native object identity when it moves."""
+ if record_type not in self.capture_cycle_configuration.children_of(
+ parent.identity.record_type
+ ):
+ raise ValueError("Unsupported Notion enumeration scope")
+ return CompletedScope(
+ record_type=record_type,
+ parent=parent.identity,
+ container_id=parent.identity.native_id if record_type == "page_property" else None,
+ )
- async def _create_comprehensive_page_entity(
+ async def capture_page(
self,
- page: dict,
- breadcrumbs: List[Breadcrumb],
- database_id: Optional[str] = None,
- database_schema: Optional[dict] = None,
- ) -> Tuple[NotionPageEntity, List[NotionFileEntity]]:
- """Create a comprehensive page entity with full aggregated content."""
- page_id = page["id"]
- title = _extract_page_title(page)
-
- self.logger.debug(f"Processing page: {title} ({page_id})")
-
- self._child_databases_to_process.clear()
-
- content_result = await self._aggregate_page_content(page_id, breadcrumbs)
-
- property_entities: List[NotionPropertyEntity] = []
- formatted_properties: Dict[str, Any] = {}
- if database_id and database_schema:
- property_entities = await self._extract_page_properties(
- page, database_id, database_schema
- )
- formatted_properties = self._create_formatted_properties_dict(
- page.get("properties", {}), database_schema.get("properties", {})
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Return bounded observations; SQL commits them together with this continuation."""
+ if scope.record_type == "page_property":
+ if (
+ parent is None
+ or parent.identity != scope.parent
+ or scope != self.child_scope(parent, "page_property")
+ ):
+ raise ValueError("Notion property scope has a different native page owner")
+ return await self._property_page(
+ parent, _PropertyProgress.model_validate(continuation.value), files
)
-
- properties_text = self._generate_properties_text_for_page(formatted_properties, title)
-
- page_entity = NotionPageEntity.from_api(
- page,
- breadcrumbs=breadcrumbs,
- content=content_result["content"],
- formatted_properties=formatted_properties,
- properties_text=properties_text,
- property_entities=property_entities,
- content_blocks_count=content_result["blocks_count"],
- max_depth=content_result["max_depth"],
+ progress = _Progress.model_validate(continuation.value)
+ if scope.container_id is not None:
+ raise ValueError("Notion native identities do not have a synthetic container")
+ if scope.parent is None:
+ if parent is not None or scope.record_type not in _ROOTS:
+ raise ValueError("Invalid Notion discovery scope")
+ return await self._search(scope.record_type, progress)
+ if parent is None or parent.identity != scope.parent:
+ raise ValueError("Notion enumeration owner is missing or mismatched")
+ self.child_scope(parent, scope.record_type)
+ if scope.record_type == "block":
+ return await self._blocks(parent, progress)
+ if parent.identity.record_type == "database":
+ return await self._database_sources(parent, progress)
+ if scope.record_type == "database":
+ located = _Located.model_validate(parent.payload)
+ native = located.parent.database_id
+ roots = await self._roots("database", [str(native)] if native else [])
+ return self._page((), progress, True, roots)
+ return await self._query_rows(parent, progress)
+
+ @staticmethod
+ def _page(
+ records: Sequence[CaptureRecord],
+ progress: _Progress | _PropertyProgress,
+ final: bool,
+ discovered: Sequence[CaptureRecord] = (),
+ ) -> CapturePage:
+ return CapturePage(
+ records=tuple(records),
+ discovered_records=tuple(discovered),
+ continuation=ScanContinuation(value=progress.model_dump(mode="json")),
+ final=final,
)
- files_with_breadcrumbs = []
- for file_entity in content_result["files"]:
- file_entity.breadcrumbs = breadcrumbs.copy()
- files_with_breadcrumbs.append(file_entity)
+ @staticmethod
+ def _following(page: _List, progress: _Progress) -> tuple[str | None, tuple[str, ...]]:
+ if page.has_more:
+ if not page.next_cursor or page.next_cursor == progress.cursor:
+ raise ValueError("Notion returned a missing or repeated cursor")
+ digest = hashlib.sha256(page.next_cursor.encode()).hexdigest()
+ if digest in progress.cursor_hashes:
+ raise ValueError("Notion pagination entered a cursor cycle")
+ # Detect recent cycles without imposing a total page-count ceiling.
+ # Longer-period loops remain subject to the existing runtime limits.
+ return page.next_cursor, (*progress.cursor_hashes[-511:], digest)
+ if page.next_cursor is not None:
+ raise ValueError("Notion returned an inconsistent terminal cursor")
+ return None, progress.cursor_hashes
+
+ async def _roots(self, kind: str, identities: list[str]) -> list[CaptureRecord]:
+ result = []
+ for native in dict.fromkeys(identities):
+ payload = await self._retrieve(kind, native)
+ if payload is not None:
+ record = self._record(kind, payload)
+ # Side admissions cannot claim deletion of a different root inventory.
+ if record.kind == "upsert":
+ result.append(record)
+ return result
+
+ async def _search(self, kind: str, progress: _Progress) -> CapturePage:
+ if kind == "database":
+ return self._page((), progress, True)
+ body: dict[str, JsonValue] = {
+ "page_size": _PAGE_SIZE,
+ "filter": {"property": "object", "value": kind},
+ }
+ if progress.cursor:
+ body["start_cursor"] = progress.cursor
+ page = _List.model_validate(await self._request("search", body))
+ if page.request_status and page.request_status.type == "incomplete":
+ raise SourceError(
+ "Notion search was capped; discovery has not completed", source_short_name="notion"
+ )
+ records = []
+ for item in page.results:
+ parsed = _Object.model_validate(item)
+ if parsed.object != kind:
+ raise ValueError("Notion search returned a different filtered object type")
+ payload = await self._retrieve(kind, str(parsed.id))
+ if payload is not None:
+ records.append(self._record(kind, payload))
+ cursor, cursor_hashes = self._following(page, progress)
+ return self._page(
+ records,
+ progress.model_copy(update={"cursor": cursor, "cursor_hashes": cursor_hashes}),
+ cursor is None,
+ )
- self.logger.debug(
- f"Page entity created: {len(content_result['content'])} chars, "
- f"{content_result['blocks_count']} blocks, "
- f"{len(files_with_breadcrumbs)} files, "
- f"max depth {content_result['max_depth']}"
+ async def _database_sources(self, parent: SourceRecord, progress: _Progress) -> CapturePage:
+ # Revalidate exact current database, not a stale embedded data-source inventory.
+ payload = await self._retrieve("database", parent.identity.native_id)
+ if payload is None:
+ raise ScopeAccessLost("Notion database is unavailable", removal_reason="scope_removed")
+ database = _Database.model_validate(payload)
+ if database.in_trash:
+ raise ScopeAccessLost("Notion database is in trash", removal_reason="scope_removed")
+ identities = database.data_sources
+ ids = [str(_Id.model_validate(item).id) for item in identities]
+ # One bounded response is re-read on resume. A changed list must restart rather
+ # than skip objects by an offset into a different inventory.
+ digest = hashlib.sha256("\n".join(ids).encode()).hexdigest()
+ if progress.cursor and progress.cursor != digest:
+ raise InvalidScanContinuation("Notion database inventory changed; restart this scope")
+ roots = await self._roots(
+ "data_source", ids[progress.offset : progress.offset + _PAGE_SIZE]
+ )
+ offset = progress.offset + _PAGE_SIZE
+ return self._page(
+ (),
+ progress.model_copy(update={"offset": min(offset, len(ids)), "cursor": digest}),
+ offset >= len(ids),
+ roots,
)
- return page_entity, files_with_breadcrumbs
-
- async def _aggregate_page_content(
- self, page_id: str, page_breadcrumbs: List[Breadcrumb]
- ) -> Dict[str, Any]:
- """Aggregate all content from a page into a single markdown string."""
- content_parts: List[str] = []
- files: List[NotionFileEntity] = []
- blocks_count = 0
- max_depth = 0
-
- async for block_content in self._extract_blocks_recursive(page_id, 0, page_breadcrumbs):
- content_parts.append(block_content["content"])
- files.extend(block_content["files"])
- blocks_count += 1
- max_depth = max(max_depth, block_content["depth"])
- self._stats["total_blocks_processed"] += 1
-
- if max_depth > self._stats["max_page_depth"]:
- self._stats["max_page_depth"] = max_depth
-
- self._stats["total_files_found"] += len(files)
-
- full_content = "\n\n".join(filter(None, content_parts))
- cleaned_content = self.clean_content_for_embedding(full_content)
-
- return {
- "content": cleaned_content,
- "files": files,
- "blocks_count": blocks_count,
- "max_depth": max_depth,
+ async def _query_rows(self, parent: SourceRecord, progress: _Progress) -> CapturePage:
+ body: dict[str, JsonValue] = {
+ "page_size": _PAGE_SIZE,
+ "sorts": [{"timestamp": "created_time", "direction": "ascending"}],
}
-
- async def _extract_blocks_recursive(
- self,
- block_id: str,
- depth: int,
- page_breadcrumbs: List[Breadcrumb],
- ) -> AsyncGenerator[Dict[str, Any], None]:
- """Recursively extract and format blocks into markdown content."""
- url = f"https://api.notion.com/v1/blocks/{block_id}/children"
- has_more = True
- start_cursor = None
-
- while has_more:
- url_with_params = url if not start_cursor else f"{url}?start_cursor={start_cursor}"
-
- try:
- response = await self._get(url_with_params)
- if not self._is_valid_response(response, block_id):
- break
-
- blocks = response.get("results", [])
- async for result in self._process_blocks(blocks, block_id, depth, page_breadcrumbs):
- yield result
-
- has_more = response.get("has_more", False)
- start_cursor = response.get("next_cursor")
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Access issue extracting blocks from {block_id}: {e}")
- break
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(
- f"Error extracting blocks from {block_id}: {type(e).__name__}: {e}"
+ if progress.cursor:
+ body["start_cursor"] = progress.cursor
+ if progress.window_start:
+ body["filter"] = {
+ "timestamp": "created_time",
+ "created_time": {"on_or_after": progress.window_start.isoformat()},
+ }
+ payload = await self._request(f"data_sources/{parent.identity.native_id}/query", body)
+ if payload is None:
+ raise ScopeAccessLost(
+ "Notion data source is unavailable", removal_reason="scope_removed"
+ )
+ page = _List.model_validate(payload)
+ roots = []
+ last = progress.last_created
+ for item in page.results:
+ parsed = _Object.model_validate(item)
+ if parsed.object not in {"page", "data_source"}:
+ raise ValueError("Notion data source returned an unexpected row type")
+ if last is not None and parsed.created_time < last:
+ raise ValueError("Notion data source violated created-time ordering")
+ last = parsed.created_time
+ roots.extend(await self._roots(parsed.object, [str(parsed.id)]))
+ cursor, cursor_hashes = self._following(page, progress)
+ if (
+ page.request_status
+ and page.request_status.type == "incomplete"
+ and page.request_status.incomplete_reason != "query_result_limit_reached"
+ ):
+ raise SourceError(
+ "Notion query was incomplete for an unsupported reason", source_short_name="notion"
+ )
+ capped = progress.capped or bool(
+ page.request_status and page.request_status.type == "incomplete"
+ )
+ next_progress = progress.model_copy(
+ update={
+ "cursor": cursor,
+ "cursor_hashes": cursor_hashes,
+ "last_created": last,
+ "capped": capped,
+ }
+ )
+ if cursor is None and capped:
+ if last is None or (
+ progress.window_start is not None and last <= progress.window_start
+ ):
+ raise SourceError(
+ "Notion query limit cannot be split by creation time; capture is incomplete",
+ source_short_name="notion",
)
- break
-
- def _is_valid_response(self, response: Any, block_id: str) -> bool:
- """Validate API response format."""
- if not response or not isinstance(response, dict):
- self.logger.warning(
- f"Invalid response format for blocks from {block_id}: {type(response)}"
+ next_progress = _Progress(window_start=last)
+ return self._page((), next_progress, False, roots)
+ return self._page((), next_progress, cursor is None, roots)
+
+ async def _blocks(self, parent: SourceRecord, progress: _Progress) -> CapturePage:
+ if parent.identity.record_type == "block":
+ block = _Block.model_validate(parent.payload)
+ if (
+ not block.has_children
+ or block.type in {"child_page", "child_database"}
+ or (block.synced_block and block.synced_block.get("synced_from") is not None)
+ ):
+ return self._page((), progress, True)
+ params: dict[str, str | int] = {"page_size": _PAGE_SIZE}
+ if progress.cursor:
+ params["start_cursor"] = progress.cursor
+ payload = await self._request(f"blocks/{parent.identity.native_id}/children", params=params)
+ if payload is None:
+ raise ScopeAccessLost(
+ "Notion block owner is unavailable", removal_reason="scope_removed"
)
- return False
- return True
-
- async def _process_blocks(
- self,
- blocks: List[dict],
- block_id: str,
- depth: int,
- page_breadcrumbs: List[Breadcrumb],
- ) -> AsyncGenerator[Dict[str, Any], None]:
- """Process a list of blocks and their children."""
- for block in blocks:
- if not block or not isinstance(block, dict):
- self.logger.warning(f"Skipping invalid block in {block_id}: {type(block)}")
- continue
-
- block_result = await self._format_block_content(block, depth, page_breadcrumbs)
- if block_result and block_result.get("content"):
- yield block_result
-
- if block.get("has_children", False) and block.get("id"):
- async for child_result in self._extract_blocks_recursive(
- block["id"], depth + 1, page_breadcrumbs
- ):
- yield child_result
-
- async def _format_block_content(
- self, block: dict, depth: int, page_breadcrumbs: List[Breadcrumb]
- ) -> Dict[str, Any]:
- """Format a single block into markdown content."""
- if not self._validate_block(block):
- return {"content": "", "files": [], "depth": depth}
-
- block_type = block.get("type", "")
- if not block_type:
- self.logger.warning(f"Block missing type field: {block.get('id', 'unknown')}")
- return {"content": "", "files": [], "depth": depth}
-
- block_content = block.get(block_type, {}) or {}
- content, files = self._dispatch_block_formatter(
- block_type, block_content, block, depth, page_breadcrumbs
+ page = _List.model_validate(payload)
+ if page.request_status and page.request_status.type == "incomplete":
+ raise ValueError("Notion block enumeration was incomplete")
+ records, roots = [], []
+ for item in page.results:
+ block = _Block.model_validate(item)
+ if self._block_parent(block) != parent.identity:
+ raise ValueError("Notion block has a different native owner")
+ records.append(self._record("block", item, parent=parent.identity))
+ target = {"child_page": "page", "child_database": "database"}.get(block.type)
+ if target and not block.in_trash:
+ roots.extend(await self._roots(target, [str(block.id)]))
+ cursor, cursor_hashes = self._following(page, progress)
+ return self._page(
+ records,
+ progress.model_copy(update={"cursor": cursor, "cursor_hashes": cursor_hashes}),
+ cursor is None,
+ roots,
)
- return {"content": content, "files": files, "depth": depth}
-
- def _validate_block(self, block: dict) -> bool:
- """Validate block structure."""
- if not block or not isinstance(block, dict):
- self.logger.warning(f"Invalid block format: {type(block)}")
- return False
- return True
-
- def _dispatch_block_formatter(
+ async def _property_inventory(
self,
- block_type: str,
- block_content: dict,
- block: dict,
- depth: int,
- page_breadcrumbs: List[Breadcrumb],
- ) -> Tuple[str, List[NotionFileEntity]]:
- """Dispatch to appropriate block formatter based on type."""
- if block_type == "paragraph":
- return self._extract_rich_text_markdown(block_content.get("rich_text", [])), []
- if block_type in ["heading_1", "heading_2", "heading_3"]:
- return self._format_heading_block(block_content, block_type), []
- if block_type in ["bulleted_list_item", "numbered_list_item", "to_do"]:
- return self._format_list_blocks(block_content, block_type, depth), []
- if block_type in ["quote", "callout", "code"]:
- return self._format_text_blocks(block_content, block_type), []
- if block_type in ["image", "video", "file", "pdf"]:
- return self._format_file_block(block_content, block, block_type)
- if block_type in ["embed", "bookmark", "equation", "divider"]:
- return self._format_simple_blocks(block_content, block_type), []
- if block_type in ["child_page", "child_database"]:
- return (
- self._format_child_blocks(block_content, block, block_type, page_breadcrumbs),
- [],
+ page_id: str,
+ ) -> tuple[str, list[tuple[str, dict[str, JsonValue]]], datetime]:
+ payload = await self._retrieve("page", page_id)
+ if payload is None or _Object.model_validate(payload).in_trash:
+ raise ScopeAccessLost(
+ "Notion property page is unavailable", removal_reason="scope_removed"
)
- return self._format_other_blocks(block_content, block_type), []
-
- def _format_heading_block(self, block_content: dict, block_type: str) -> str:
- """Format heading blocks."""
- level = int(block_type.split("_")[1])
- text = self._extract_rich_text_markdown(block_content.get("rich_text", []))
- prefix = "▸ " if block_content.get("is_toggleable", False) else ""
- return f"{'#' * level} {prefix}{text}"
-
- def _format_list_blocks(self, block_content: dict, block_type: str, depth: int) -> str:
- """Format list and todo blocks."""
- text = self._extract_rich_text_markdown(block_content.get("rich_text", []))
- indent = " " * depth
-
- if block_type == "bulleted_list_item":
- return f"{indent}- {text}"
- elif block_type == "numbered_list_item":
- return f"{indent}1. {text}"
- else: # to_do
- checkbox = "- [x]" if block_content.get("checked", False) else "- [ ]"
- return f"{indent}{checkbox} {text}"
-
- def _format_text_blocks(self, block_content: dict, block_type: str) -> str:
- """Format quote, callout, and code blocks."""
- if block_type == "quote":
- text = self._extract_rich_text_markdown(block_content.get("rich_text", []))
- return f"> {text}"
- elif block_type == "callout":
- return self._format_callout_block(block_content)
- else: # code
- return self._format_code_block(block_content)
-
- def _format_simple_blocks(self, block_content: dict, block_type: str) -> str:
- """Format simple blocks like embed, bookmark, equation, divider."""
- if block_type in ["embed", "bookmark"]:
- url = block_content.get("url", "")
- caption = _extract_rich_text_plain(block_content.get("caption", []))
- if len(url) > 200 or "?" in url and len(url.split("?", 1)[1]) > 100:
- content = f"[{block_type.title()}]"
- else:
- content = f"[{block_type.title()}]({url})"
- if caption:
- content += f" - {caption}"
- return content
- elif block_type == "equation":
- expression = block_content.get("expression", "")
- return f"$$\n{expression}\n$$"
- elif block_type == "divider":
- return "---"
- else:
- return "**Table of Contents**"
-
- def _format_child_blocks(
- self, block_content: dict, block: dict, block_type: str, page_breadcrumbs: List[Breadcrumb]
- ) -> str:
- """Format child page and database blocks."""
- if block_type == "child_page":
- title = block_content.get("title", "Untitled Page")
- return f"📄 **[{title}]** (Child Page)"
- else:
- return self._format_child_database_block(block_content, block, page_breadcrumbs)
-
- def _format_other_blocks(self, block_content: dict, block_type: str) -> str:
- """Format other block types including table, column, etc."""
- if block_type in ["table", "column_list"]:
- return f"**[{block_type.replace('_', ' ').title()}]**"
- elif block_type in ["table_row", "column"]:
- return ""
- elif block_type == "unsupported":
- return "*[Unsupported block type]*"
- else:
- return self._format_unknown_block(block_content, block_type)
-
- def _format_callout_block(self, block_content: dict) -> str:
- """Format callout blocks."""
- icon = block_content.get("icon", {})
- icon_text = icon.get("emoji", "💡") if icon.get("type") == "emoji" else "💡"
- text = self._extract_rich_text_markdown(block_content.get("rich_text", []))
- return f"**{icon_text} {text}**"
-
- def _format_code_block(self, block_content: dict) -> str:
- """Format code blocks."""
- language = block_content.get("language", "")
- code_text = _extract_rich_text_plain(block_content.get("rich_text", []))
- caption = _extract_rich_text_plain(block_content.get("caption", []))
- content = f"```{language}\n{code_text}\n```"
- if caption:
- content += f"\n*{caption}*"
- return content
-
- def _format_file_block(
- self, block_content: dict, block: dict, block_type: str
- ) -> Tuple[str, List[NotionFileEntity]]:
- """Format file blocks and return content and file entities."""
- file_entity = NotionFileEntity.from_api(
- block_content, parent_id=block["id"], breadcrumbs=[]
+ page = _PageProperties.model_validate(payload)
+ rows = sorted(
+ page.properties.items(), key=lambda item: _Property.model_validate(item[1]).id
)
- files = [file_entity]
- caption = _extract_rich_text_plain(block_content.get("caption", []))
-
- if block_type == "image":
- display_name = file_entity.name if file_entity.name != "Untitled File" else "Image"
- content = f"[Image: {display_name}]"
- else:
- content = f"[File: {file_entity.name}]"
-
- if caption:
- content += f" - {caption}"
-
- return content, files
-
- def _format_child_database_block(
- self, block_content: dict, block: dict, page_breadcrumbs: List[Breadcrumb]
- ) -> str:
- """Format child database blocks."""
- title = block_content.get("title", "Untitled Database")
- database_id = block["id"]
-
- self._child_databases_to_process.add(database_id)
-
- child_db_breadcrumbs = page_breadcrumbs.copy()
- self._child_database_breadcrumbs[database_id] = child_db_breadcrumbs
-
- self._stats["child_databases_found"] += 1
- return f"🗃️ **[{title}]** (Child Database)"
-
- def _format_unknown_block(self, block_content: dict, block_type: str) -> str:
- """Format unknown block types."""
- if "rich_text" in block_content:
- return self._extract_rich_text_markdown(block_content.get("rich_text", []))
+ identities = [
+ (_Property.model_validate(value).id, name, _Property.model_validate(value).type)
+ for name, value in rows
+ ]
+ if len({item[0] for item in identities}) != len(identities):
+ raise ValueError("Notion page contains duplicate property IDs")
+ material = json.dumps(
+ [page.last_edited_time.isoformat(), identities], ensure_ascii=False
+ ).encode()
+ return hashlib.sha256(material).hexdigest(), rows, page.last_edited_time
+
+ @staticmethod
+ def _property_response(
+ payload: dict[str, JsonValue], prop: _Property, progress: _Progress
+ ) -> tuple[_Progress, bool, Literal["available", "unsupported"]]:
+ if payload.get("object") == "property_item":
+ if progress.cursor is not None:
+ raise InvalidScanContinuation(
+ "Notion property changed response shape while reading"
+ )
+ item = _PropertyItem.model_validate(payload)
+ metadata = item
+ final = True
+ following = progress
else:
- return f"*[{block_type.replace('_', ' ').title()}]*"
-
- # ------------------------------------------------------------------
- # Rich text extraction
- # ------------------------------------------------------------------
-
- def _extract_rich_text_markdown(self, rich_text: List[dict]) -> str:
- """Extract rich text and convert to markdown formatting."""
- if not rich_text or not isinstance(rich_text, list):
- return ""
-
- result_parts = []
- for text_obj in rich_text:
- if not text_obj or not isinstance(text_obj, dict):
- continue
-
- text = text_obj.get("plain_text", "")
- if not text:
- continue
-
- annotations = text_obj.get("annotations") or {}
- href = text_obj.get("href")
- formatted_text = self._apply_markdown_annotations(text, annotations, href)
- result_parts.append(formatted_text)
-
- return "".join(result_parts)
-
- def _apply_markdown_annotations(self, text: str, annotations: dict, href: str = None) -> str:
- """Apply markdown formatting based on annotations."""
- if annotations.get("bold"):
- text = f"**{text}**"
- if annotations.get("italic"):
- text = f"*{text}*"
- if annotations.get("strikethrough"):
- text = f"~~{text}~~"
- if annotations.get("underline"):
- text = f"{text}"
- if annotations.get("code"):
- text = f"`{text}`"
- if href:
- text = f"[{text}]({href})"
- return text
-
- # ------------------------------------------------------------------
- # Property extraction
- # ------------------------------------------------------------------
-
- async def _extract_page_properties(
- self, page: dict, database_id: str, database_schema: dict
- ) -> List[NotionPropertyEntity]:
- """Extract database page properties as structured entities."""
- page_id = page["id"]
- page_properties = page.get("properties", {})
- schema_properties = database_schema.get("properties", {})
-
- property_entities = []
-
- for prop_name, prop_value in page_properties.items():
- if prop_name in schema_properties:
- schema_prop = schema_properties[prop_name]
-
- try:
- property_entity = self._create_property_entity(
- prop_name, prop_value, schema_prop, page_id, database_id
- )
- property_entities.append(property_entity)
-
- except Exception as e:
- self.logger.warning(
- f"Error processing property {prop_name} for page {page_id}: {e}"
- )
- continue
-
- return property_entities
-
- def _create_property_entity(
- self, prop_name: str, prop_value: dict, schema_prop: dict, page_id: str, database_id: str
- ) -> NotionPropertyEntity:
- """Create a property entity from page property data."""
- prop_type = prop_value.get("type", "")
- formatted_value = self._format_property_value(prop_value, prop_type)
- schema_prop_id = schema_prop.get("id", prop_name)
- property_key = f"{page_id}_{schema_prop_id}"
-
- return NotionPropertyEntity(
- entity_id=property_key,
- breadcrumbs=[],
- name=prop_name,
- created_at=None,
- updated_at=None,
- property_key=property_key,
- property_id=schema_prop_id,
- property_name=prop_name,
- property_type=prop_type,
- page_id=page_id,
- database_id=database_id,
- value=prop_value.get(prop_type),
- formatted_value=formatted_value,
+ page = _PropertyList.model_validate(payload)
+ if page.request_status and page.request_status.type == "incomplete":
+ raise SourceError(
+ "Notion property response is incomplete", source_short_name="notion"
+ )
+ metadata = page.property_item
+ for value in page.results:
+ item = _PropertyItem.model_validate(value)
+ # Rollup show_original may return the target property's items.
+ if prop.type != "rollup" and (item.id != prop.id or item.type != prop.type):
+ raise ValueError("Notion property list contains another property's values")
+ cursor, hashes = NotionSource._following(page, progress)
+ following = progress.model_copy(update={"cursor": cursor, "cursor_hashes": hashes})
+ final = cursor is None
+ if metadata.id != prop.id or metadata.type != prop.type:
+ raise ValueError("Notion returned a different property identity or type")
+ calculation = metadata.rollup if prop.type == "rollup" else metadata.formula
+ if prop.type in {"rollup", "formula"} and calculation is None:
+ raise ValueError("Notion calculation has no native result metadata")
+ if final and calculation and calculation.type == "incomplete":
+ raise SourceError(
+ "Notion property calculation did not complete", source_short_name="notion"
+ )
+ status = "unsupported" if calculation and calculation.type == "unsupported" else "available"
+ return following, final, status
+
+ async def _collect_property(
+ self, page_id: str, prop: _Property, fingerprint: str, page_last_edited_time: datetime
+ ) -> tuple[bytes, int, Literal["available", "unsupported"]]:
+ """Retain native pages once each; retry starts this property without partial publication."""
+ header = (
+ _PropertyArchiveHeader(
+ page_id=UUID(page_id),
+ property_id=prop.id,
+ page_last_edited_time=page_last_edited_time,
+ )
+ .model_dump_json()
+ .encode()
)
-
- # ------------------------------------------------------------------
- # Utility methods
- # ------------------------------------------------------------------
-
- def _format_property_value(self, prop_value: dict, prop_type: str) -> str:
- """Format property value for human readability."""
- if not prop_value or prop_type not in prop_value:
- return ""
-
- value = prop_value[prop_type]
-
- formatters = {
- "title": lambda v: _extract_rich_text_plain(v),
- "rich_text": lambda v: _extract_rich_text_plain(v),
- "number": lambda v: str(v) if v is not None else "",
- "url": lambda v: str(v) if v is not None else "",
- "email": lambda v: str(v) if v is not None else "",
- "phone_number": lambda v: str(v) if v is not None else "",
- "checkbox": lambda v: "Yes" if v else "No",
- "select": lambda v: self._format_select_properties(v),
- "status": lambda v: self._format_select_properties(v),
- "multi_select": lambda v: ", ".join([opt.get("name", "") for opt in v]) if v else "",
- "date": lambda v: self._format_date_property(v),
- "people": lambda v: self._format_people_property(v),
- "files": lambda v: f"{len(v)} file(s)" if v else "0 files",
- "created_time": lambda v: v or "",
- "last_edited_time": lambda v: v or "",
- "created_by": lambda v: v.get("name", "Unknown User") if v else "",
- "last_edited_by": lambda v: v.get("name", "Unknown User") if v else "",
- }
-
- if prop_type in formatters:
- return formatters[prop_type](value)
-
- return self._format_complex_property_types(prop_type, value)
-
- def _format_complex_property_types(self, prop_type: str, value: Any) -> str:
- """Format complex property types that need special handling."""
- if prop_type == "relation":
- return f"{len(value)} relation(s)" if value else "0 relations"
- elif prop_type == "rollup":
- return self._format_rollup_property(value)
- elif prop_type == "formula":
- return self._format_formula_property(value)
- elif prop_type == "unique_id":
- return self._format_unique_id_property(value)
- elif prop_type == "verification":
- return self._format_verification_property(value)
- else:
- return str(value) if value else ""
-
- def _format_unique_id_property(self, value: dict) -> str:
- prefix = value.get("prefix", "")
- number = value.get("number", "")
- return f"{prefix}{number}" if prefix else str(number)
-
- def _format_verification_property(self, value: dict) -> str:
- state = value.get("state", "")
- return state.title() if state else ""
-
- def _format_select_properties(self, value: dict) -> str:
- return value.get("name", "") if value else ""
-
- def _format_date_property(self, value: dict) -> str:
- if value and value.get("start"):
- start = value["start"]
- end = value.get("end")
- return f"{start} - {end}" if end else start
- return ""
-
- def _format_people_property(self, value: List[dict]) -> str:
- names = []
- for person in value:
- if person.get("type") == "person":
- names.append(person.get("name", "Unknown"))
- elif person.get("type") == "bot":
- names.append(person.get("name", "Bot"))
- return ", ".join(names)
-
- def _format_formula_property(self, value: dict) -> str:
- formula_type = value.get("type", "")
- if formula_type in ["string", "number", "boolean", "date"]:
- return str(value.get(formula_type, ""))
- return ""
-
- def _format_rollup_property(self, value: dict) -> str:
- rollup_type = value.get("type", "")
- if rollup_type in ["string", "number", "boolean", "date"]:
- return str(value.get(rollup_type, ""))
- elif rollup_type == "array":
- return f"{len(value.get('array', []))} item(s)"
- return ""
-
- def _create_formatted_properties_dict(
- self, page_properties: dict, database_schema: dict
- ) -> Dict[str, Any]:
- """Create a clean, formatted properties dictionary for better searchability."""
- formatted: Dict[str, Any] = {}
-
- for prop_name, prop_value in page_properties.items():
- if not prop_value:
- continue
-
- prop_type = prop_value.get("type", "")
- formatted_value = self._format_property_value(prop_value, prop_type)
-
- if not formatted_value:
- continue
-
- formatted[prop_name] = formatted_value
-
- if prop_type in ["select", "status", "multi_select"] and prop_name in database_schema:
- schema_prop = database_schema[prop_name]
- if prop_type == "multi_select":
- options = schema_prop.get("multi_select", {}).get("options", [])
- else:
- options = schema_prop.get(prop_type, {}).get("options", [])
-
- if options:
- formatted[f"{prop_name}_options"] = [
- opt.get("name", "") for opt in options if opt.get("name")
- ]
-
- return formatted
-
- def _generate_properties_text_for_page(
- self, properties: Dict[str, Any], page_title: str
- ) -> str:
- """Generate human-readable text from properties for embedding."""
- if not properties:
- return ""
-
- text_parts = []
-
- priority_keys = [
- "Product Name",
- "Name",
- "Title",
- "Status",
- "Priority",
- "Launch Status",
- "Owner",
- "Team",
- "Description",
- ]
-
- for key in priority_keys:
- if key in properties:
- value = properties[key]
- if value and str(value).strip():
- if key in ["Product Name", "Name", "Title"] and value == page_title:
- continue
- text_parts.append(f"{key}: {value}")
-
- for key, value in properties.items():
- if key not in priority_keys and not key.endswith("_options"):
- if value and str(value).strip():
- formatted_key = key.replace("_", " ").title()
- text_parts.append(f"{formatted_key}: {value}")
-
- return " | ".join(text_parts) if text_parts else ""
-
- # ------------------------------------------------------------------
- # File processing
- # ------------------------------------------------------------------
-
- async def _process_and_yield_file( # noqa: C901
- self, file_entity: NotionFileEntity, files: FileService | None = None
- ) -> Optional[NotionFileEntity]:
- """Process a file entity by downloading it and setting local_path."""
- try:
- if file_entity.needs_refresh():
- self.logger.warning(f"Skipping file {file_entity.name} - URL expired")
- return None
-
- if not file_entity.url or file_entity.url.startswith("http") is False:
- self.logger.debug(f"Skipping file {file_entity.name} - no valid download URL")
- return None
-
- if not files:
- return None
-
- try:
- await files.download_from_url(
- entity=file_entity,
- client=self.http_client,
- auth=self.auth,
- logger=self.logger,
+ prefix = header[:-1] + b',"responses":['
+ responses: list[bytes] = []
+ total = len(prefix) + len(b"]}")
+ following = _Progress()
+ while True:
+ params: dict[str, str | int] = {"page_size": _PAGE_SIZE}
+ if following.cursor:
+ params["start_cursor"] = following.cursor
+ payload = await self._request(f"pages/{page_id}/properties/{prop.id}", params=params)
+ if payload is None:
+ # The page may have been removed, or its schema changed. Do not
+ # withdraw the whole page based solely on a property endpoint404.
+ current, _, _ = await self._property_inventory(page_id)
+ if current != fingerprint:
+ raise InvalidScanContinuation(
+ "Notion property inventory changed during retrieval"
+ )
+ raise SourceError(
+ "Notion property is unavailable while its page remains readable",
+ source_short_name="notion",
)
-
- if not file_entity.local_path:
- raise ValueError(f"Download failed - no local path set for {file_entity.name}")
-
- self.logger.debug(
- f"Successfully downloaded file {file_entity.name} "
- f"to local_path: {file_entity.local_path}"
+ following, final, status = self._property_response(payload, prop, following)
+ encoded = json.dumps(
+ payload, ensure_ascii=False, allow_nan=False, separators=(",", ":"), sort_keys=True
+ ).encode()
+ total += len(encoded) + bool(responses)
+ if total > MAX_FILE_SIZE_BYTES:
+ raise SourceError(
+ "Notion property exceeds the retained-file size limit; capture is incomplete",
+ source_short_name="notion",
)
- return file_entity
-
- except FileSkippedException as e:
- self.logger.debug(f"Skipping file {file_entity.name}: {e.reason}")
- return None
-
- except httpx.HTTPStatusError as e:
- if e.response.status_code == 401:
- raise
- self.logger.warning(f"Failed to download file {file_entity.name}: {e}")
- return None
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Access issue processing file {file_entity.name}: {e}")
- return None
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing file {file_entity.name}: {e}")
- return None
-
- # ------------------------------------------------------------------
- # Main entry point
- # ------------------------------------------------------------------
+ responses.append(encoded)
+ if final:
+ break
+ return prefix + b",".join(responses) + b"]}", len(responses), status
+
+ async def _property_page(
+ self, parent: SourceRecord, progress: _PropertyProgress, files: FileService
+ ) -> CapturePage:
+ page_id = parent.identity.native_id
+ fingerprint, properties, page_last_edited_time = await self._property_inventory(page_id)
+ if progress.inventory is not None and progress.inventory != fingerprint:
+ raise InvalidScanContinuation(
+ "Notion page properties changed; restart their enumeration"
+ )
+ if progress.index > len(properties):
+ raise ValueError("Notion property continuation is outside its inventory")
+ if progress.index == len(properties):
+ return self._page((), progress.model_copy(update={"inventory": fingerprint}), True)
+ name, native_property = properties[progress.index]
+ prop = _Property.model_validate(native_property)
+ archive, count, status = await self._collect_property(
+ page_id, prop, fingerprint, page_last_edited_time
+ )
+ current, _, _ = await self._property_inventory(page_id)
+ if current != fingerprint:
+ raise InvalidScanContinuation("Notion page changed before property capture completed")
+ blob = await files.store_canonical_blob(archive, media_type="application/json")
+ record = CaptureRecord(
+ identity=RecordIdentity(
+ record_type="page_property", native_id=prop.id, container_id=page_id
+ ),
+ parent=parent.identity,
+ payload=_PropertyManifest(
+ page_id=UUID(page_id),
+ property_id=prop.id,
+ name=name,
+ property=native_property,
+ response_count=count,
+ value_status=status,
+ page_last_edited_time=page_last_edited_time,
+ ).model_dump(mode="json"),
+ payload_schema_version=FIELD_SET_VERSION,
+ completeness="partial",
+ blobs=(blob,),
+ observed_at=datetime.now(timezone.utc),
+ )
+ next_progress = _PropertyProgress(index=progress.index + 1, inventory=fingerprint)
+ return self._page((record,), next_progress, next_progress.index == len(properties))
- async def generate_entities( # noqa: C901
+ async def generate_entities(
self,
*,
cursor: SyncCursor | None = None,
files: FileService | None = None,
node_selections: list[NodeSelectionData] | None = None,
) -> AsyncGenerator[BaseEntity, None]:
- """Generate all entities from Notion using streaming discovery."""
- self.logger.debug("Starting Notion entity generation with content aggregation")
-
- self._processed_pages: Set[str] = set()
- self._processed_databases: Set[str] = set()
- self._child_databases_to_process: Set[str] = set()
- self._child_database_breadcrumbs: Dict[str, List[Breadcrumb]] = {}
- self._stats: Dict[str, int] = {
- "api_calls": 0,
- "rate_limit_waits": 0,
- "databases_found": 0,
- "child_databases_found": 0,
- "pages_found": 0,
- "total_blocks_processed": 0,
- "total_files_found": 0,
- "max_page_depth": 0,
- }
-
- try:
- self.logger.debug("PHASE 1 & 2: Streaming database discovery and schema analysis")
- async for entity in self._stream_database_discovery(files):
- yield entity
- self.logger.debug(
- f"Phase 1 & 2 complete: {self._stats['databases_found']} databases found"
- )
-
- self.logger.debug("PHASE 3: Streaming standalone page discovery")
- async for entity in self._stream_page_discovery(files):
- yield entity
- self.logger.debug(
- f"Phase 3 complete: {self._stats['pages_found']} standalone pages found"
- )
-
- self.logger.debug("PHASE 4: Processing child databases")
- async for entity in self._process_child_databases(files):
- yield entity
- self.logger.debug(
- f"Phase 4 complete: {self._stats['child_databases_found']} child databases found"
- )
-
- self.logger.debug("Notion sync complete. Final stats:")
- for key, value in self._stats.items():
- self.logger.debug(f" {key}: {value}")
-
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(
- f"Error during streaming Notion entity generation: {e}", exc_info=True
- )
- raise
-
- async def _stream_database_discovery(
- self, files: FileService | None = None
- ) -> AsyncGenerator[BaseEntity, None]:
- """Discover databases and delegate per-database processing."""
- self.logger.debug("Streaming database discovery...")
-
- async for database in self._search_objects("database"):
- database_id = database["id"]
- if database_id in self._processed_databases:
- continue
- self._stats["databases_found"] += 1
- async for entity in self._process_single_database(database_id, files):
- yield entity
-
- async def _process_single_database( # noqa: C901
- self, database_id: str, files: FileService | None = None
- ) -> AsyncGenerator[BaseEntity, None]:
- """Process one database: fetch schema, emit DB entity, pages, files, child DBs."""
- try:
- self.logger.debug(f"Fetching schema for database: {database_id}")
- schema = await self._get(f"https://api.notion.com/v1/databases/{database_id}")
- self._processed_databases.add(database_id)
-
- database_entity = NotionDatabaseEntity.from_api(schema, breadcrumbs=[])
- yield database_entity
-
- database_title = _extract_rich_text_plain(schema.get("title", []))
- self.logger.debug(f"Processing pages in database: {database_title}")
-
- async for page in self._query_database_pages(database_id):
- page_id = page["id"]
- if page_id in self._processed_pages:
- continue
- try:
- breadcrumbs = [
- Breadcrumb(
- entity_id=database_id,
- name=database_title or "Untitled Database",
- entity_type=NotionDatabaseEntity.__name__,
- )
- ]
- page_entity, file_entities = await self._create_comprehensive_page_entity(
- page, breadcrumbs, database_id, schema
- )
- yield page_entity
- for file_entity in file_entities:
- processed = await self._process_and_yield_file(file_entity, files)
- if processed:
- yield processed
- self._processed_pages.add(page_id)
- async for child_entity in self._process_child_databases(files):
- yield child_entity
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Access issue processing database page {page_id}: {e}")
- continue
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing database page {page_id}: {e}")
- continue
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Access issue processing database {database_id}: {e}")
- return
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing database {database_id}: {e}")
- return
-
- async def _stream_page_discovery(
- self, files: FileService | None = None
- ) -> AsyncGenerator[BaseEntity, None]:
- """Stream page discovery and immediately yield page entities."""
- self.logger.debug("Streaming page discovery...")
-
- async for page in self._search_objects("page"):
- page_id = page["id"]
- if page_id in self._processed_pages:
- continue
-
- parent = page.get("parent", {})
- parent_type = parent.get("type", "")
-
- if parent_type == "database_id":
- continue
-
- self._stats["pages_found"] += 1
-
- try:
- full_page = await self._get(f"https://api.notion.com/v1/pages/{page_id}")
-
- breadcrumbs = await self._build_page_breadcrumbs(full_page)
-
- page_entity, file_entities = await self._create_comprehensive_page_entity(
- full_page, breadcrumbs
- )
- yield page_entity
-
- for file_entity in file_entities:
- processed = await self._process_and_yield_file(file_entity, files)
- if processed:
- yield processed
-
- self._processed_pages.add(page_id)
-
- async for child_entity in self._process_child_databases(files):
- yield child_entity
-
- except NotionSource.NotionAccessError as e:
- self.logger.warning(f"Access issue processing standalone page {page_id}: {e}")
- continue
- except SourceAuthError:
- raise
- except Exception as e:
- self.logger.warning(f"Error processing standalone page {page_id}: {e}")
- continue
+ """The legacy formatter/crawler is replaced, not maintained as a second authority."""
+ raise NotImplementedError("Notion requires canonical page capture")
+ yield # pragma: no cover
diff --git a/backend/airweave/platform/sources/onedrive.py b/backend/airweave/platform/sources/onedrive.py
index 6f61e9151d..f7bb68244d 100644
--- a/backend/airweave/platform/sources/onedrive.py
+++ b/backend/airweave/platform/sources/onedrive.py
@@ -33,13 +33,13 @@
from airweave.platform.entities._base import BaseEntity
from airweave.platform.entities.onedrive import OneDriveDriveEntity, OneDriveDriveItemEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/onenote.py b/backend/airweave/platform/sources/onenote.py
index 57c266581c..9e30f6bf96 100644
--- a/backend/airweave/platform/sources/onenote.py
+++ b/backend/airweave/platform/sources/onenote.py
@@ -36,12 +36,12 @@
OneNoteSectionEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/outlook_calendar.py b/backend/airweave/platform/sources/outlook_calendar.py
index 9d294b35f8..f3d7fab243 100644
--- a/backend/airweave/platform/sources/outlook_calendar.py
+++ b/backend/airweave/platform/sources/outlook_calendar.py
@@ -15,9 +15,10 @@
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderError
from airweave.domains.browse_tree.types import NodeSelectionData
from airweave.domains.sources.exceptions import SourceAuthError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.sources.token_providers.protocol import SourceAuthProvider
from airweave.domains.storage import FileSkippedException
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
@@ -30,12 +31,13 @@
OutlookCalendarEventEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.composio_transport import ComposioProxyError
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
@@ -65,38 +67,38 @@ class OutlookCalendarSource(BaseSource):
GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0"
+ # Query/projection capability; only explicit fresh-source configuration opts
+ # an instance into the composed canonical page adapter below.
+ canonical_record_types = ("calendar", "event")
+
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: OutlookCalendarConfig,
) -> "OutlookCalendarSource":
"""Create a new Outlook Calendar source instance."""
- return cls(auth=auth, logger=logger, http_client=http_client)
+ instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.graph = OutlookGraphClient(
+ auth, http_client, cls.short_name, config.expected_principal_id
+ )
+ if config.capture_originals:
+ from airweave.platform.sources.outlook_calendar_capture import OutlookCalendarCapture
+
+ instance._capture_page_source = await OutlookCalendarCapture.create(
+ graph=instance.graph, config=config
+ )
+ else:
+ await instance.graph.verify_principal()
+ return instance
# ------------------------------------------------------------------
# HTTP helpers
# ------------------------------------------------------------------
- async def _authed_headers(self) -> Dict[str, str]:
- """Build Authorization + Accept headers with a fresh token."""
- token = await self.auth.get_token()
- return {
- "Authorization": f"Bearer {token}",
- "Accept": "application/json",
- }
-
- async def _refresh_and_get_headers(self) -> Dict[str, str]:
- """Force-refresh the token and return updated headers."""
- new_token = await self.auth.force_refresh()
- return {
- "Authorization": f"Bearer {new_token}",
- "Accept": "application/json",
- }
-
@retry(
stop=stop_after_attempt(5),
retry=retry_if_rate_limit_or_timeout,
@@ -104,25 +106,8 @@ async def _refresh_and_get_headers(self) -> Dict[str, str]:
reraise=True,
)
async def _get(self, url: str, params: Optional[dict] = None) -> Any:
- """Make an authenticated GET request to Microsoft Graph API."""
- self.logger.debug(f"Making authenticated GET request to: {url} with params: {params}")
-
- headers = await self._authed_headers()
- response = await self.http_client.get(url, headers=headers, params=params)
-
- if response.status_code == 401 and self.auth.supports_refresh:
- self.logger.warning(
- f"Got 401 Unauthorized from Microsoft Graph API at {url}, refreshing token..."
- )
- headers = await self._refresh_and_get_headers()
- response = await self.http_client.get(url, headers=headers, params=params)
-
- raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
- )
- return response.json()
+ """Read only within the attested Microsoft Graph mailbox."""
+ return await self.graph.get(url, params=params)
# ------------------------------------------------------------------
# Entity generators
@@ -141,7 +126,7 @@ async def _generate_calendar_entities(
try:
while url:
- self.logger.debug(f"Fetching calendars from: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
calendars = data.get("value", [])
self.logger.info(f"Retrieved {len(calendars)} calendars")
@@ -184,7 +169,7 @@ async def _generate_calendar_entities(
self.logger.info(f"Completed calendar generation. Total calendars: {calendar_count}")
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error generating calendar entities: {str(e)}")
@@ -215,7 +200,7 @@ async def _generate_event_entities(
try:
while url:
- self.logger.debug(f"Fetching events from: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url, params=params)
events = data.get("value", [])
self.logger.info(f"Retrieved {len(events)} events from calendar {calendar_name}")
@@ -236,7 +221,12 @@ async def _generate_event_entities(
event_data, cal_breadcrumb, files=files
):
yield entity
- except SourceAuthError:
+ except (
+ SourceAuthError,
+ AuthProviderError,
+ ComposioProxyError,
+ OutlookBoundaryError,
+ ):
raise
except Exception as e:
self.logger.warning(f"Error processing event {event_id}: {str(e)}")
@@ -251,7 +241,7 @@ async def _generate_event_entities(
f"Total events: {event_count}"
)
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error generating events for calendar {calendar_name}: {str(e)}")
@@ -299,7 +289,7 @@ async def _process_event(
self.logger.debug(
f"Processed {attachment_count} attachments for event {event_subject}"
)
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error processing attachments for event {event_id}: {str(e)}")
@@ -318,7 +308,7 @@ async def _process_event_attachments(
try:
while url:
- self.logger.debug(f"Making request to: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
attachments = data.get("value", [])
self.logger.debug(f"Retrieved {len(attachments)} attachments for event {event_id}")
@@ -340,7 +330,7 @@ async def _process_event_attachments(
if url:
self.logger.debug("Following pagination link")
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error processing attachments for event {event_id}: {str(e)}")
@@ -413,7 +403,12 @@ async def _process_single_attachment( # noqa: C901
self.logger.debug(f"Skipping attachment {attachment_name}: {e.reason}")
return None
- except SourceAuthError:
+ except (
+ SourceAuthError,
+ AuthProviderError,
+ ComposioProxyError,
+ OutlookBoundaryError,
+ ):
raise
except Exception as e:
@@ -422,7 +417,7 @@ async def _process_single_attachment( # noqa: C901
else:
return file_entity
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error processing attachment {attachment_id}: {str(e)}")
@@ -440,6 +435,7 @@ async def generate_entities(
node_selections: list[NodeSelectionData] | None = None,
) -> AsyncGenerator[BaseEntity, None]:
"""Generate all Outlook Calendar entities: Calendars, Events and Attachments."""
+ await self.graph.verify_principal()
self.logger.info("Starting Outlook Calendar entity generation")
entity_count = 0
@@ -462,7 +458,7 @@ async def generate_entities(
)
yield event_entity
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error in entity generation: {str(e)}", exc_info=True)
@@ -474,6 +470,7 @@ async def generate_entities(
async def validate(self) -> None:
"""Validate credentials by pinging the calendars endpoint."""
+ await self.graph.verify_principal()
await self._get(
f"{self.GRAPH_BASE_URL}/me/calendars",
params={"$top": "1"},
diff --git a/backend/airweave/platform/sources/outlook_calendar_capture.py b/backend/airweave/platform/sources/outlook_calendar_capture.py
new file mode 100644
index 0000000000..4c226e6571
--- /dev/null
+++ b/backend/airweave/platform/sources/outlook_calendar_capture.py
@@ -0,0 +1,364 @@
+"""Calendar-owned Graph originals using the shared resumable full-scan engine.
+
+Unbounded singles and masters precede bounded native expansion in the SAME event
+scope. Previously captured occurrences remain retained while exact reads succeed;
+the completed window describes expansion coverage, not a deletion boundary.
+"""
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from urllib.parse import quote, unquote, urlsplit
+
+from pydantic import AwareDatetime, TypeAdapter, ValidationError
+
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import CapturePage, CapturePlan
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation, ScanState
+from airweave.domains.entities.canonical.scope_execution import ScopePlan
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.configs.config import CalendarOccurrenceWindow, OutlookCalendarConfig
+from airweave.platform.sources.outlook_calendar_models import (
+ OutlookCalendar,
+ OutlookCalendarContinuation,
+ OutlookCalendarEvent,
+ OutlookCalendarPage,
+ OutlookCalendarScopeContext,
+)
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient
+
+BASE = "https://graph.microsoft.com/v1.0/me"
+
+
+class OutlookCalendarCapture:
+ """Provider interpretation only; SQL owns progress, parent visibility and publication."""
+
+ canonical_record_types = ("calendar", "event")
+ canonical_container_parents = {"event": "calendar"}
+
+ def __init__(self, graph: OutlookGraphClient, config: OutlookCalendarConfig) -> None:
+ """Use create to attest the principal before obtaining this composition."""
+ self.graph = graph
+ self.config = config.model_copy(deep=True)
+ self.fingerprint = hashlib.sha256(
+ json.dumps(
+ {"version": 1, "id_type": "ImmutableId", "config": config.model_dump(mode="json")},
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+
+ @classmethod
+ async def create(
+ cls, *, graph: OutlookGraphClient, config: OutlookCalendarConfig
+ ) -> "OutlookCalendarCapture":
+ """Require trusted mailbox identity; credential lifecycle remains in GraphClient."""
+ instance = cls(graph, config)
+ if not config.expected_principal_id or (
+ graph.expected_principal_id != config.expected_principal_id
+ ):
+ raise instance._error("Original Outlook calendars require a bound principal")
+ await graph.verify_principal()
+ instance._require_principal()
+ return instance
+
+ @staticmethod
+ def _error(message: str) -> OutlookBoundaryError:
+ return OutlookBoundaryError(message, source_short_name="outlook_calendar")
+
+ def _require_principal(self) -> None:
+ if not self.config.expected_principal_id or (
+ self.graph.expected_principal_id != self.config.expected_principal_id
+ or self.graph.verified_principal_id != self.config.expected_principal_id
+ ):
+ raise self._error("Original Outlook calendar principal is not attested")
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Neither listing omissions nor ambiguous404 imply provider deletion."""
+ self._require_principal()
+ return CycleConfiguration.from_source(
+ fingerprint=self.fingerprint,
+ record_types=self.canonical_record_types,
+ container_parents=self.canonical_container_parents,
+ completion_policies={
+ "calendar": "discovery_with_validation",
+ "event": "discovery_with_validation",
+ },
+ )
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Resolve rolling bounds once; this first version does not claim incremental capture."""
+ await self.graph.verify_principal()
+ self._require_principal()
+ return CapturePlan(
+ mode="full",
+ source_plan={"window": self.config.resolved_window().model_dump(mode="json")},
+ )
+
+ def _scope(self, scope: CompletedScope, parent: SourceRecord | None) -> None:
+ root = scope.record_type == "calendar"
+ if root and scope.container_id is None and scope.parent is None and parent is None:
+ return
+ if (
+ scope.record_type == "event"
+ and scope.parent is not None
+ and scope.parent.record_type == "calendar"
+ and scope.parent.container_id is None
+ and scope.container_id == scope.parent.native_id
+ and parent is not None
+ and parent.identity == scope.parent
+ ):
+ return
+ raise self._error("Outlook calendar scope does not match its native owner")
+
+ async def prepare_scope(
+ self,
+ scope: CompletedScope,
+ cycle: CaptureCycle,
+ previous: ScanState | None,
+ *,
+ parent: SourceRecord | None,
+ force_full: bool,
+ ) -> ScopePlan:
+ """Both phases share one full sweep and one immutable expansion window."""
+ self._scope(scope, parent)
+ if cycle.configuration != self.capture_cycle_configuration or cycle.mode != "full":
+ raise self._error("Outlook calendar cycle belongs to another configuration")
+ try:
+ window = CalendarOccurrenceWindow.model_validate(cycle.source_plan["window"])
+ context = OutlookCalendarScopeContext(
+ fingerprint=self.fingerprint, calendar_id=scope.container_id, window=window
+ )
+ return ScopePlan(request_context=context.model_dump(mode="json"))
+ except (ValidationError, KeyError):
+ raise self._error("Invalid Outlook calendar scope plan") from None
+
+ def initial_scope_continuation(
+ self, scope: CompletedScope, cycle: CaptureCycle, plan: ScopePlan
+ ) -> ScanContinuation:
+ """Resume only the persisted scope request, never newly resolved wall-clock bounds."""
+ self._require_principal()
+ try:
+ context = OutlookCalendarScopeContext.model_validate(plan.request_context)
+ except ValidationError:
+ raise self._error("Invalid Outlook calendar scope context") from None
+ if (
+ plan.mode != "full"
+ or cycle.mode != "full"
+ or plan.starting_checkpoint is not None
+ or cycle.configuration != self.capture_cycle_configuration
+ or context.fingerprint != self.fingerprint
+ or context.calendar_id != scope.container_id
+ or context.window.model_dump(mode="json") != cycle.source_plan.get("window")
+ ):
+ raise self._error("Outlook calendar scope context changed")
+ return self._continuation(
+ OutlookCalendarContinuation(
+ context=context, phase="calendars" if scope.record_type == "calendar" else "events"
+ )
+ )
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """One native calendar owns all event kinds and their eventual attachment children."""
+ if (
+ parent.identity.record_type != "calendar"
+ or parent.parent is not None
+ or record_type != "event"
+ ):
+ raise self._error("Invalid Outlook calendar child scope")
+ return CompletedScope(
+ record_type="event", container_id=parent.identity.native_id, parent=parent.identity
+ )
+
+ def _continuation(self, state: OutlookCalendarContinuation) -> ScanContinuation:
+ try:
+ return ScanContinuation(value=state.model_dump(mode="json"))
+ except ValidationError:
+ raise self._error(
+ "Outlook calendar continuation exceeds its bounded capacity"
+ ) from None
+
+ def _collection_url(self, url: str, collection: str) -> None:
+ try:
+ actual, expected = urlsplit(url), urlsplit(collection)
+ if (
+ actual.scheme == "https"
+ and actual.netloc == "graph.microsoft.com"
+ and unquote(actual.path) == unquote(expected.path)
+ and not actual.fragment
+ and not any(ord(char) < 32 for char in url)
+ ):
+ return
+ except ValueError:
+ pass
+ raise self._error("Outlook calendar continuation changed its native collection")
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Commit a complete bounded page with its next phase; late failures do not advance it."""
+ self._require_principal()
+ self._scope(scope, parent)
+ try:
+ state = OutlookCalendarContinuation.model_validate(continuation.value)
+ except ValidationError:
+ raise self._error("Invalid Outlook calendar continuation") from None
+ root = scope.record_type == "calendar"
+ if (
+ state.context.fingerprint != self.fingerprint
+ or state.context.calendar_id != scope.container_id
+ or (state.phase not in (("calendars",) if root else ("events", "calendarView")))
+ ):
+ raise self._error("Outlook calendar continuation belongs to another scope")
+ collection = f"{BASE}/calendars"
+ params: dict[str, str | int] | None = {"$top": 100, "$select": "id"}
+ if not root:
+ collection += f"/{quote(scope.container_id or '', safe='')}/{state.phase}"
+ if state.phase == "calendarView":
+ # Graph's default response representation is UTC. Preserve that native JSON.
+ params = {
+ "$top": 100,
+ "startDateTime": state.context.window.start.isoformat(),
+ "endDateTime": state.context.window.end.isoformat(),
+ }
+ url = state.next_link or collection
+ self._collection_url(url, collection)
+ raw = await self.graph.get(
+ url, params=None if state.next_link else params, immutable_ids=True
+ )
+ try:
+ page = OutlookCalendarPage.model_validate(raw)
+ except ValidationError:
+ raise self._error("Invalid Outlook calendar collection response") from None
+ ids = tuple(item.id for item in page.value)
+ if len(set(ids)) != len(ids):
+ raise self._error("Outlook calendar page repeated a native identity")
+ progress = self._advance(state, page, collection, url)
+ observations = []
+ for native_id in ids:
+ if root:
+ observations.append(await self._calendar(native_id))
+ else:
+ observations.append(await self._event(scope.parent, native_id))
+ return CapturePage(
+ records=tuple(observations),
+ continuation=progress,
+ final=progress.value["phase"] == "done",
+ )
+
+ def _advance(
+ self,
+ state: OutlookCalendarContinuation,
+ page: OutlookCalendarPage,
+ collection: str,
+ url: str,
+ ) -> ScanContinuation:
+ """Validate bounded progress before fetching exact originals."""
+ links = state.recent_links
+ if page.next_link:
+ self._collection_url(page.next_link, collection)
+ digest = hashlib.sha256(page.next_link.encode()).hexdigest()
+ if page.next_link == url or digest in links:
+ raise self._error("Outlook calendar pagination repeated a recent link")
+ links = (*links, digest)[-128:]
+ phase = (
+ state.phase
+ if page.next_link
+ else ("calendarView" if state.phase == "events" else "done")
+ )
+ return self._continuation(
+ state.model_copy(
+ update={
+ "next_link": page.next_link,
+ "recent_links": links if page.next_link else (),
+ "phase": phase,
+ }
+ )
+ )
+
+ async def _calendar(self, native_id: str) -> CaptureRecord:
+ raw = await self.graph.get(
+ f"{BASE}/calendars/{quote(native_id, safe='')}", immutable_ids=True
+ )
+ try:
+ calendar = OutlookCalendar.model_validate(raw)
+ except ValidationError:
+ raise self._error("Invalid Outlook calendar metadata") from None
+ if calendar.id != native_id:
+ raise self._error("Outlook calendar identity changed during capture")
+ return CaptureRecord(
+ identity=RecordIdentity(record_type="calendar", native_id=native_id),
+ payload=raw,
+ observed_at=datetime.now(timezone.utc),
+ descendant_visibility_fields=("canViewPrivateItems",),
+ )
+
+ async def _event(self, parent: RecordIdentity | None, native_id: str) -> CaptureRecord:
+ if parent is None:
+ raise self._error("Outlook event original requires its calendar owner")
+ calendar_url = f"{BASE}/calendars/{quote(parent.native_id, safe='')}"
+ url = f"{calendar_url}/events/{quote(native_id, safe='')}"
+ raw = await self.graph.get(url, immutable_ids=True)
+ try:
+ event = OutlookCalendarEvent.model_validate(raw)
+ created = (
+ TypeAdapter(AwareDatetime).validate_python(event.createdDateTime)
+ if event.createdDateTime
+ else None
+ )
+ updated = (
+ TypeAdapter(AwareDatetime).validate_python(event.lastModifiedDateTime)
+ if event.lastModifiedDateTime
+ else None
+ )
+ except ValidationError:
+ raise self._error("Invalid Outlook event metadata") from None
+ if event.id != native_id:
+ raise self._error("Outlook event identity changed during capture")
+ return CaptureRecord(
+ identity=RecordIdentity(
+ record_type="event", native_id=native_id, container_id=parent.native_id
+ ),
+ parent=parent,
+ payload=raw,
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=created,
+ source_updated_at=updated,
+ # Native attachments and selected-only master cancellation metadata are not
+ # retained yet. hasAttachments=false must not certify full content capture.
+ completeness="partial",
+ )
+
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Historical instances remain originals; an ambiguous404 fails before reconciliation."""
+ self._require_principal()
+ if (
+ record.identity.record_type == "calendar"
+ and record.parent is None
+ and record.identity.container_id is None
+ ):
+ return await self._calendar(record.identity.native_id)
+ if (
+ record.identity.record_type == "event"
+ and record.parent is not None
+ and record.parent.record_type == "calendar"
+ and record.parent.container_id is None
+ and record.identity.container_id == record.parent.native_id
+ ):
+ return await self._event(record.parent, record.identity.native_id)
+ raise self._error("Invalid Outlook exact-read identity")
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """No arbitrary Graph failure establishes authoritative native absence."""
+ raise self._error("Outlook calendar omissions require exact current observation")
diff --git a/backend/airweave/platform/sources/outlook_calendar_models.py b/backend/airweave/platform/sources/outlook_calendar_models.py
new file mode 100644
index 0000000000..dfb481edb2
--- /dev/null
+++ b/backend/airweave/platform/sources/outlook_calendar_models.py
@@ -0,0 +1,100 @@
+"""Native Outlook calendar boundaries; originals are retained without model dumping."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator
+
+from airweave.platform.configs.config import CalendarOccurrenceWindow
+from airweave.platform.sources.outlook_mail_models import OutlookBody, OutlookRecipient
+
+
+class OutlookCalendar(BaseModel):
+ """Calendar identity and permission metadata, not a Google calendar representation."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(min_length=1)
+ name: str | None = None
+ owner: dict[str, JsonValue] | None = None
+ canEdit: bool | None = None
+ canViewPrivateItems: bool | None = None
+ isDefaultCalendar: bool | None = None
+
+
+class OutlookDateTimeZone(BaseModel):
+ """Graph wall time plus native zone, including Windows or custom zone names."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ dateTime: str = Field(min_length=1)
+ timeZone: str = Field(min_length=1)
+
+
+class OutlookCalendarEvent(BaseModel):
+ """The native event kind distinguishes masters, singles and expanded instances."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(min_length=1)
+ changeKey: str = Field(min_length=1)
+ type: Literal["singleInstance", "seriesMaster", "occurrence", "exception"]
+ start: OutlookDateTimeZone
+ end: OutlookDateTimeZone
+ isAllDay: bool
+ isCancelled: bool
+ hasAttachments: bool
+ subject: str | None = None
+ body: OutlookBody | None = None
+ bodyPreview: str | None = None
+ organizer: OutlookRecipient | None = None
+ attendees: list[dict[str, JsonValue]] = Field(default_factory=list)
+ seriesMasterId: str | None = None
+ originalStart: str | None = None
+ originalStartTimeZone: str | None = None
+ originalEndTimeZone: str | None = None
+ recurrence: dict[str, JsonValue] | None = None
+ cancelledOccurrences: list[str] | None = None
+ iCalUId: str | None = None
+ webLink: str | None = None
+ createdDateTime: str | None = None
+ lastModifiedDateTime: str | None = None
+
+
+class OutlookCalendarID(BaseModel):
+ """A collection item is discovery evidence, not the retained exact original."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(min_length=1)
+ removed: JsonValue = Field(default=None, alias="@removed")
+
+ @model_validator(mode="after")
+ def full_item(self):
+ """These full-list endpoints have no delta tombstone contract."""
+ if "removed" in self.model_fields_set:
+ raise ValueError("Unexpected removed item in a full calendar listing")
+ return self
+
+
+class OutlookCalendarPage(BaseModel):
+ """An absent array or unexpected delta removal is never an empty successful page."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ value: list[OutlookCalendarID] = Field(max_length=500)
+ next_link: str | None = Field(default=None, alias="@odata.nextLink", min_length=1)
+
+
+class OutlookCalendarScopeContext(BaseModel):
+ """Immutable per-scope request authority, persisted by the shared page engine."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ version: Literal[1] = 1
+ fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
+ calendar_id: str | None = Field(default=None, min_length=1)
+ window: CalendarOccurrenceWindow
+
+
+class OutlookCalendarContinuation(BaseModel):
+ """Only one collection's opaque continuation, never the expanding original archive."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ context: OutlookCalendarScopeContext
+ phase: Literal["calendars", "events", "calendarView", "done"]
+ next_link: str | None = Field(default=None, min_length=1)
+ recent_links: tuple[str, ...] = Field(default=(), max_length=128)
diff --git a/backend/airweave/platform/sources/outlook_graph.py b/backend/airweave/platform/sources/outlook_graph.py
new file mode 100644
index 0000000000..a24091ca51
--- /dev/null
+++ b/backend/airweave/platform/sources/outlook_graph.py
@@ -0,0 +1,232 @@
+"""Outlook's native mailbox boundary over the existing managed/direct transport."""
+
+from contextlib import aclosing
+from urllib.parse import unquote, urlsplit
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
+
+from airweave.domains.sources.exceptions import SourceError
+from airweave.domains.sources.token_providers.protocol import (
+ ManagedAuthProvider,
+ SourceAuthProvider,
+ authorization_headers,
+)
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.http_client.composio_transport import ComposioProxyError
+from airweave.platform.sources.http_helpers import raise_for_status
+
+
+class OutlookBoundaryError(SourceError):
+ """Unsafe mailbox scope or invalid principal evidence; never a skipped object."""
+
+
+class OutlookDeltaExpiredError(SourceError):
+ """Documented native delta expiration, without retaining private error bodies."""
+
+
+class GraphError(BaseModel):
+ """Only the native error code is interpreted; messages are never exposed."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ code: str
+
+
+class GraphErrorEnvelope(BaseModel):
+ """Structured expiry evidence, distinct from transport and authentication errors."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ error: GraphError
+
+
+class GraphPrincipal(BaseModel):
+ """Only the native, case-sensitive principal ID establishes mailbox identity."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ id: str = Field(min_length=1, pattern=r"^\S+$")
+
+
+class OutlookGraphClient:
+ """Validate native scope and identity; credential custody stays in the transport."""
+
+ def __init__(
+ self,
+ auth: SourceAuthProvider,
+ http_client: AirweaveHttpClient,
+ source_name: str,
+ expected_principal_id: str | None,
+ ) -> None:
+ """Bind the expected principal supplied by the trusted connection owner."""
+ self.auth = auth
+ self.http_client = http_client
+ self.source_name = source_name
+ self.expected_principal_id = expected_principal_id
+ self.verified_principal_id: str | None = None
+ if isinstance(auth, ManagedAuthProvider) and expected_principal_id is None:
+ raise self._boundary_error("Managed Outlook requires a trusted expected principal ID")
+
+ def _boundary_error(self, message: str) -> OutlookBoundaryError:
+ return OutlookBoundaryError(
+ message,
+ source_short_name=self.source_name,
+ )
+
+ def _validate_url(self, url: str) -> None:
+ try:
+ parsed = urlsplit(url)
+ path = unquote(parsed.path)
+ valid = (
+ parsed.scheme == "https"
+ and parsed.netloc == "graph.microsoft.com"
+ and not parsed.fragment
+ and (path == "/v1.0/me" or path.startswith("/v1.0/me/"))
+ and "\\" not in path
+ and not any(segment in (".", "..") for segment in path.split("/"))
+ and "%" not in path
+ and not any(ord(char) < 32 for char in url)
+ )
+ except ValueError:
+ valid = False
+ if not valid:
+ raise self._boundary_error("Microsoft Graph URL is outside the bound mailbox")
+
+ async def _request(
+ self, url: str, params: dict | None = None, *, immutable_ids: bool = False
+ ) -> httpx.Response:
+ self._validate_url(url)
+ headers = await authorization_headers(self.auth)
+ # Derived folder types remain explicit even when the request selects only id.
+ headers["Accept"] = "application/json;odata.metadata=minimal"
+ if immutable_ids:
+ headers["Prefer"] = 'IdType="ImmutableId"'
+ try:
+ return await self.http_client.get(
+ url, headers=headers, params=params, follow_redirects=False
+ )
+ except ComposioProxyError:
+ raise
+ except httpx.TimeoutException:
+ raise httpx.TimeoutException("Microsoft Graph request timed out") from None
+ except httpx.RequestError:
+ raise httpx.RequestError("Microsoft Graph request failed") from None
+
+ def _check_response(self, response: httpx.Response) -> None:
+ # Native errors may echo tokens or private data. Keep status/retry timing only.
+ safe = httpx.Response(
+ response.status_code,
+ headers={"Retry-After": response.headers["Retry-After"]}
+ if "Retry-After" in response.headers
+ else {},
+ json={},
+ )
+ raise_for_status(
+ safe, source_short_name=self.source_name, token_provider_kind=self.auth.provider_kind
+ )
+
+ async def verify_principal(self) -> None:
+ """Re-attest on validation/enumeration; a failed check invalidates old evidence."""
+ self.verified_principal_id = None
+ if self.expected_principal_id is None:
+ return # Legacy direct accounts need not acquire a new User.Read permission.
+ response = await self._request("https://graph.microsoft.com/v1.0/me", {"$select": "id"})
+ if response.status_code == 401 and self.auth.supports_refresh:
+ await authorization_headers(self.auth, refresh=True)
+ response = await self._request("https://graph.microsoft.com/v1.0/me", {"$select": "id"})
+ self._check_response(response)
+ try:
+ principal = GraphPrincipal.model_validate(response.json())
+ except (ValidationError, ValueError):
+ raise self._boundary_error(
+ "Microsoft Graph returned an invalid principal identity"
+ ) from None
+ if principal.id != self.expected_principal_id:
+ raise self._boundary_error(
+ "Microsoft Graph principal does not match the expected identity"
+ )
+ self.verified_principal_id = principal.id
+
+ async def get(
+ self, url: str, params: dict | None = None, *, immutable_ids: bool = False
+ ) -> dict:
+ """Preserve legacy IDs while rejecting requests without current attestation."""
+ if self.expected_principal_id is not None and (
+ self.verified_principal_id != self.expected_principal_id
+ ):
+ raise self._boundary_error("Microsoft Graph principal must be attested before reading")
+ response = await self._request(url, params, immutable_ids=immutable_ids)
+ if response.status_code == 401:
+ self.verified_principal_id = None
+ if self.auth.supports_refresh:
+ await authorization_headers(self.auth, refresh=True)
+ await self.verify_principal()
+ response = await self._request(url, params, immutable_ids=immutable_ids)
+ if response.status_code == 401:
+ self.verified_principal_id = None
+ if urlsplit(url).path.lower().endswith("/messages/delta") and response.status_code in (
+ 400,
+ 404,
+ 410,
+ ):
+ try:
+ error = GraphErrorEnvelope.model_validate(response.json())
+ except ValueError:
+ error = None
+ if error is not None and error.error.code == "syncStateNotFound":
+ raise OutlookDeltaExpiredError(
+ "Microsoft Graph delta state expired", source_short_name=self.source_name
+ )
+ self._check_response(response)
+ try:
+ data = TypeAdapter(dict[str, JsonValue]).validate_python(response.json())
+ if "error" in data:
+ raise self._boundary_error("Microsoft Graph returned an error envelope")
+ return data
+ except ValueError:
+ raise self._boundary_error(
+ "Microsoft Graph returned an invalid JSON response"
+ ) from None
+
+ async def mime_bytes(self, url: str, *, max_bytes: int) -> bytes:
+ """Bound original MIME reads; canonical requests alone opt into immutable IDs."""
+ if self.expected_principal_id is None or (
+ self.verified_principal_id != self.expected_principal_id
+ ):
+ raise self._boundary_error("Canonical MIME requires an attested principal")
+ self._validate_url(url)
+ for attempt in range(2):
+ headers = await authorization_headers(self.auth)
+ headers["Prefer"] = 'IdType="ImmutableId"'
+ try:
+ async with self.http_client.stream(
+ "GET", url, headers=headers, follow_redirects=False
+ ) as response:
+ if response.status_code == 401:
+ self.verified_principal_id = None
+ if not attempt and self.auth.supports_refresh:
+ await authorization_headers(self.auth, refresh=True)
+ await self.verify_principal()
+ continue
+ self._check_response(response)
+ return await self._read_mime_response(response, max_bytes)
+ except ComposioProxyError:
+ raise
+ except httpx.TimeoutException:
+ raise httpx.TimeoutException("Microsoft Graph MIME request timed out") from None
+ except httpx.RequestError:
+ raise httpx.RequestError("Microsoft Graph MIME request failed") from None
+ raise self._boundary_error("Microsoft Graph MIME authentication failed")
+
+ async def _read_mime_response(self, response: httpx.Response, max_bytes: int) -> bytes:
+ media_type = response.headers.get("content-type", "").split(";", 1)[0].strip().lower()
+ if media_type not in ("message/rfc822", "text/plain", "application/octet-stream"):
+ raise self._boundary_error("Microsoft Graph returned a non-MIME response")
+ content = bytearray()
+ async with aclosing(response.aiter_bytes()) as chunks:
+ async for chunk in chunks:
+ if len(content) + len(chunk) > max_bytes:
+ raise FileSkippedException("Original MIME exceeds byte limit", "message")
+ content.extend(chunk)
+ if not content:
+ raise self._boundary_error("Microsoft Graph returned empty MIME")
+ return bytes(content)
diff --git a/backend/airweave/platform/sources/outlook_mail.py b/backend/airweave/platform/sources/outlook_mail.py
index cdc3ce6ccc..e50b565d73 100644
--- a/backend/airweave/platform/sources/outlook_mail.py
+++ b/backend/airweave/platform/sources/outlook_mail.py
@@ -16,9 +16,10 @@
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderError
from airweave.domains.browse_tree.types import NodeSelectionData
from airweave.domains.sources.exceptions import SourceAuthError, SourceEntityNotFoundError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.sources.token_providers.protocol import SourceAuthProvider
from airweave.domains.storage import FileSkippedException
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
@@ -34,12 +35,13 @@
OutlookMessageEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.composio_transport import ComposioProxyError
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient
from airweave.platform.utils.filename_utils import safe_filename
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
@@ -71,17 +73,32 @@ class OutlookMailSource(BaseSource):
GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0"
+ # Offline projection/read capability. This does not opt legacy source instances
+ # into CanonicalPageSource; immutable mailbox capture is composed separately.
+ canonical_record_types = ("message",)
+
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: OutlookMailConfig,
) -> "OutlookMailSource":
"""Create a new Outlook Mail source instance."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.graph = OutlookGraphClient(
+ auth, http_client, cls.short_name, config.expected_principal_id
+ )
+ if config.capture_originals:
+ from airweave.platform.sources.outlook_mail_capture import OutlookMailCapture
+
+ instance._capture_page_source = await OutlookMailCapture.create(
+ graph=instance.graph, config=config
+ )
+ else:
+ await instance.graph.verify_principal()
instance.after_date = config.after_date
instance.included_folders = config.included_folders
instance.excluded_folders = config.excluded_folders
@@ -91,16 +108,6 @@ async def create(
# HTTP helpers
# ------------------------------------------------------------------
- async def _authed_headers(self) -> Dict[str, str]:
- """Build Authorization header with a fresh token."""
- token = await self.auth.get_token()
- return {"Authorization": f"Bearer {token}"}
-
- async def _refresh_and_get_headers(self) -> Dict[str, str]:
- """Force-refresh the token and return updated headers."""
- new_token = await self.auth.force_refresh()
- return {"Authorization": f"Bearer {new_token}"}
-
@retry(
stop=stop_after_attempt(5),
retry=retry_if_rate_limit_or_timeout,
@@ -108,27 +115,8 @@ async def _refresh_and_get_headers(self) -> Dict[str, str]:
reraise=True,
)
async def _get(self, url: str, params: Optional[dict] = None) -> dict:
- """Make an authenticated GET request to Microsoft Graph API."""
- self.logger.debug(f"Making authenticated GET request to: {url} with params: {params}")
-
- headers = await self._authed_headers()
- response = await self.http_client.get(url, headers=headers, params=params)
-
- if response.status_code == 401 and self.auth.supports_refresh:
- self.logger.warning(
- f"Got 401 Unauthorized from Microsoft Graph API at {url}, refreshing token..."
- )
- headers = await self._refresh_and_get_headers()
- response = await self.http_client.get(url, headers=headers, params=params)
-
- raise_for_status(
- response,
- source_short_name=self.short_name,
- token_provider_kind=self.auth.provider_kind,
- )
- data = response.json()
- self.logger.debug(f"Received response from {url} - Status: {response.status_code}")
- return data
+ """Read only within the attested Microsoft Graph mailbox."""
+ return await self.graph.get(url, params=params)
# ------------------------------------------------------------------
# Cursor helpers
@@ -234,7 +222,7 @@ async def _process_folder_messages(
folder_entity, folder_breadcrumb, files=files
):
yield entity
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -263,7 +251,7 @@ async def _process_child_folders(
files=files,
):
yield child_entity
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -278,7 +266,7 @@ async def _init_and_store_message_delta_for_folder(
"""Initialize the per-folder message delta link and store it in the cursor."""
try:
delta_url = f"{self.GRAPH_BASE_URL}/me/mailFolders/{folder_entity.id}/messages/delta"
- self.logger.debug(f"Calling delta endpoint: {delta_url}")
+ self.logger.debug("Reading Microsoft Graph page")
delta_data = await self._get(delta_url)
attempts = 0
@@ -310,7 +298,7 @@ async def _init_and_store_message_delta_for_folder(
f"{folder_entity.display_name} while initializing delta."
)
break
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -387,7 +375,7 @@ async def _generate_folder_entities(
try:
while url:
- self.logger.debug(f"Making request to: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
folders = data.get("value", [])
self.logger.debug(f"Retrieved {len(folders)} folders")
@@ -400,10 +388,10 @@ async def _generate_folder_entities(
next_link = data.get("@odata.nextLink")
if next_link:
- self.logger.debug(f"Following pagination link: {next_link}")
+ self.logger.debug("Reading Microsoft Graph page")
url = next_link if next_link else None
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error fetching folders: {str(e)}")
@@ -477,7 +465,12 @@ async def _generate_message_entities( # noqa: C901
message_data, folder_entity.display_name, folder_breadcrumb, files=files
):
yield entity
- except SourceAuthError:
+ except (
+ SourceAuthError,
+ AuthProviderError,
+ ComposioProxyError,
+ OutlookBoundaryError,
+ ):
raise
except Exception as e:
self.logger.warning(f"Error processing message {message_id}: {str(e)}")
@@ -493,7 +486,7 @@ async def _generate_message_entities( # noqa: C901
)
break
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -561,7 +554,7 @@ async def _process_message(
self.logger.debug(
f"Processed {attachment_count} attachments for message {message_id}"
)
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -667,7 +660,7 @@ async def _process_single_attachment( # noqa: C901
self.logger.debug(f"Skipping attachment {attachment_name}: {e.reason}")
return None
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
@@ -688,7 +681,7 @@ async def _process_attachments(
try:
while url:
- self.logger.debug(f"Making request to: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
attachments = data.get("value", [])
self.logger.debug(
@@ -712,7 +705,7 @@ async def _process_attachments(
if url:
self.logger.debug("Following pagination link")
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error processing attachments for message {message_id}: {str(e)}")
@@ -736,7 +729,7 @@ async def _process_delta_changes(
url = f"{self.GRAPH_BASE_URL}/me/mailFolders/{folder_id}/messages/delta"
params = {"$deltatoken": delta_token}
while url:
- self.logger.debug(f"Fetching delta changes from: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url, params=params)
params = None
@@ -769,7 +762,7 @@ async def _process_delta_changes(
if url:
self.logger.debug("Following delta pagination")
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -814,7 +807,7 @@ async def _initialize_folders_delta_link(self, cursor: SyncCursor | None = None)
"""Initialize and store the delta link for the mailFolders collection."""
try:
init_url = f"{self.GRAPH_BASE_URL}/me/mailFolders/delta"
- self.logger.debug(f"Initializing folders delta link via: {init_url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(init_url)
safety_counter = 0
@@ -834,7 +827,7 @@ async def _initialize_folders_delta_link(self, cursor: SyncCursor | None = None)
else:
self.logger.warning("No deltaLink or nextLink while initializing folders delta")
break
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Failed to initialize folders delta link: {e}")
@@ -864,7 +857,7 @@ async def _process_folders_delta_changes(
cursor.update(folders_delta_link=new_delta_link)
self.logger.debug("Updated folders_delta_link for next incremental run")
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error processing folders delta changes: {e}")
@@ -873,7 +866,7 @@ async def _iterate_delta_pages(self, start_url: str) -> AsyncGenerator[Dict[str,
"""Iterate delta/next pages starting from a delta or nextLink URL."""
url = start_url
while url:
- self.logger.debug(f"Fetching folders delta changes from: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
yield data
url = data.get("@odata.nextLink")
@@ -1006,7 +999,7 @@ async def _emit_folder_add_or_update( # noqa: C901
msg_delta_data = await self._get(next_link)
else:
break
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Failed to initialize message delta for folder {folder_id}: {e}")
@@ -1025,7 +1018,7 @@ async def _process_delta_changes_url( # noqa: C901
try:
url = delta_url
while url:
- self.logger.debug(f"Fetching delta changes from: {url}")
+ self.logger.debug("Reading Microsoft Graph page")
data = await self._get(url)
changes = data.get("value", [])
@@ -1079,7 +1072,7 @@ async def _process_delta_changes_url( # noqa: C901
)
break
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(
@@ -1135,6 +1128,7 @@ async def generate_entities( # noqa: C901
Supports both full sync (first run) and incremental sync (subsequent runs)
using Microsoft Graph delta API.
"""
+ await self.graph.verify_principal()
self.logger.debug("===== STARTING OUTLOOK MAIL ENTITY GENERATION =====")
entity_count = 0
@@ -1178,7 +1172,7 @@ async def generate_entities( # noqa: C901
)
yield entity
- except SourceAuthError:
+ except (SourceAuthError, AuthProviderError, ComposioProxyError, OutlookBoundaryError):
raise
except Exception as e:
self.logger.warning(f"Error in entity generation: {str(e)}", exc_info=True)
@@ -1190,6 +1184,7 @@ async def generate_entities( # noqa: C901
async def validate(self) -> None:
"""Validate credentials by pinging the mailFolders endpoint."""
+ await self.graph.verify_principal()
await self._get(
f"{self.GRAPH_BASE_URL}/me/mailFolders",
params={"$top": "1"},
diff --git a/backend/airweave/platform/sources/outlook_mail_capture.py b/backend/airweave/platform/sources/outlook_mail_capture.py
new file mode 100644
index 0000000000..f6b0b27cd9
--- /dev/null
+++ b/backend/airweave/platform/sources/outlook_mail_capture.py
@@ -0,0 +1,479 @@
+"""Mailbox-owned Outlook originals with bounded per-folder delta progress."""
+
+import hashlib
+import json
+from datetime import datetime, timezone
+from urllib.parse import quote, unquote, urlsplit
+
+from pydantic import AwareDatetime, TypeAdapter, ValidationError
+
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CycleConfiguration,
+ ProviderCheckpoint,
+)
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidCaptureCheckpoint,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceGoneError
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.configs.config import OutlookMailConfig
+from airweave.platform.sources.outlook_graph import (
+ OutlookBoundaryError,
+ OutlookDeltaExpiredError,
+ OutlookGraphClient,
+)
+from airweave.platform.sources.outlook_mail_models import (
+ OutlookDeltaPage,
+ OutlookFolder,
+ OutlookFolderPage,
+ OutlookMailCheckpoint,
+ OutlookMailContinuation,
+ OutlookMessage,
+ OutlookMessageID,
+)
+
+BASE = "https://graph.microsoft.com/v1.0/me"
+
+
+class OutlookMailCapture:
+ """One root scan uses shared page commits, blobs and exact omission validation.
+
+ MIME-backed schema1 originals are partial until native attachment inventory is
+ qualified. Folder removals are signals, never proof of mailbox deletion.
+ """
+
+ canonical_record_types = ("message",)
+ canonical_container_parents: dict[str, str | tuple[str | None, ...]] = {}
+
+ def __init__(
+ self,
+ graph: OutlookGraphClient,
+ config: OutlookMailConfig,
+ included_ids: frozenset[str],
+ excluded_ids: frozenset[str],
+ ) -> None:
+ """Construct through create so the native principal and folder aliases are resolved."""
+ self.graph = graph
+ self.config = config.model_copy(deep=True)
+ self.included_ids = included_ids
+ self.excluded_ids = excluded_ids
+ self.after = (
+ datetime.strptime(config.after_date, "%Y/%m/%d").replace(tzinfo=timezone.utc)
+ if config.after_date
+ else None
+ )
+ self.fingerprint = hashlib.sha256(
+ json.dumps(
+ {
+ "version": 2,
+ "id_type": "ImmutableId",
+ "config": self.config.model_dump(mode="json"),
+ "included_ids": sorted(included_ids),
+ "excluded_ids": sorted(excluded_ids),
+ },
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+
+ @classmethod
+ async def create(
+ cls, *, graph: OutlookGraphClient, config: OutlookMailConfig
+ ) -> "OutlookMailCapture":
+ """Require a trusted identity and resolve native folder aliases, never display names."""
+ if not config.expected_principal_id or (
+ graph.expected_principal_id != config.expected_principal_id
+ ):
+ raise OutlookBoundaryError(
+ "Canonical Outlook requires a bound principal", source_short_name="outlook_mail"
+ )
+ await graph.verify_principal()
+ folders = {}
+ for alias in dict.fromkeys((*config.included_folders, *config.excluded_folders)):
+ raw = await graph.get(
+ f"{BASE}/mailFolders/{quote(alias, safe='')}", params={"$select": "id"}
+ )
+ try:
+ folders[alias] = OutlookMessageID.model_validate(raw).id
+ except ValidationError:
+ raise OutlookBoundaryError(
+ "Invalid Outlook folder identity", source_short_name="outlook_mail"
+ ) from None
+ return cls(
+ graph,
+ config,
+ frozenset(folders[x] for x in config.included_folders),
+ frozenset(folders[x] for x in config.excluded_folders),
+ )
+
+ def _require_principal(self) -> None:
+ if not self.config.expected_principal_id or (
+ self.graph.verified_principal_id != self.config.expected_principal_id
+ or self.graph.expected_principal_id != self.config.expected_principal_id
+ ):
+ raise OutlookBoundaryError(
+ "Canonical Outlook principal is not attested", source_short_name="outlook_mail"
+ )
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Observed enumeration is not a snapshot; omissions require exact validation."""
+ self._require_principal()
+ return CycleConfiguration.from_source(
+ fingerprint=self.fingerprint,
+ record_types=self.canonical_record_types,
+ container_parents={},
+ completion_policies={"message": "discovery_with_validation"},
+ )
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Only a compatible promoted mailbox checkpoint authorizes changes mode."""
+ await self.graph.verify_principal()
+ self._require_principal()
+ if (
+ previous is None
+ or previous.configuration != self.capture_cycle_configuration
+ or previous.promoted_checkpoint is None
+ ):
+ # An explicit empty starting map opts into terminal checkpoint publication.
+ # It is not a completed native baseline and full mode never reuses its tokens.
+ return CapturePlan(
+ starting_checkpoint=ProviderCheckpoint(
+ value=OutlookMailCheckpoint(fingerprint=self.fingerprint).model_dump(
+ mode="json"
+ )
+ )
+ )
+ checkpoint = previous.promoted_checkpoint.checkpoint
+ self._checkpoint(checkpoint)
+ return CapturePlan(mode="changes", starting_checkpoint=checkpoint)
+
+ def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation:
+ """Bind progress to the persisted principal, selection and immutable ID format."""
+ self._require_principal()
+ if cycle.configuration != self.capture_cycle_configuration:
+ raise self._error("Outlook capture configuration changed")
+ links = {}
+ if cycle.mode == "changes":
+ if cycle.starting_checkpoint is None:
+ raise self._error("Outlook changes require a complete mailbox checkpoint")
+ links = self._checkpoint(cycle.starting_checkpoint).folder_links
+ elif cycle.mode != "full":
+ raise self._error("Unsupported Outlook capture mode")
+ return self._continuation(
+ OutlookMailContinuation(
+ fingerprint=self.fingerprint, mode=cycle.mode, folder_links=links
+ )
+ )
+
+ @staticmethod
+ def _error(message: str) -> OutlookBoundaryError:
+ return OutlookBoundaryError(message, source_short_name="outlook_mail")
+
+ def _checkpoint(self, checkpoint: ProviderCheckpoint) -> OutlookMailCheckpoint:
+ try:
+ state = OutlookMailCheckpoint.model_validate(checkpoint.value)
+ except ValidationError:
+ raise self._error("Invalid Outlook mailbox checkpoint") from None
+ if state.fingerprint != self.fingerprint:
+ raise self._error("Outlook checkpoint belongs to another capture scope")
+ for folder, link in state.folder_links.items():
+ self._collection_url(link, self._delta_url(folder))
+ return state
+
+ def _continuation(self, state: OutlookMailContinuation) -> ScanContinuation:
+ try:
+ return ScanContinuation(value=state.model_dump(mode="json"))
+ except ValidationError:
+ raise self._error("Outlook continuation exceeds the 64 KiB capacity") from None
+
+ def _state(self, continuation: ScanContinuation) -> OutlookMailContinuation:
+ try:
+ state = OutlookMailContinuation.model_validate(continuation.value)
+ except ValidationError:
+ raise self._error("Invalid Outlook capture continuation") from None
+ if state.fingerprint != self.fingerprint:
+ raise self._error("Outlook continuation belongs to another capture scope")
+ return state
+
+ @staticmethod
+ def _delta_url(folder: str) -> str:
+ return f"{BASE}/mailFolders/{quote(folder, safe='')}/messages/delta"
+
+ def _collection_url(self, url: str, expected: str) -> None:
+ """Opaque queries remain unchanged, but cannot redirect to another collection."""
+ self.graph._validate_url(url)
+ path = unquote(urlsplit(url).path).replace("/mailfolders", "/mailFolders")
+ expected_path = unquote(urlsplit(expected).path)
+ allowed = {expected_path}
+ prefix = "/v1.0/me/mailFolders/"
+ if expected_path.startswith(prefix):
+ folder, suffix = expected_path[len(prefix) :].split("/", 1)
+ allowed.add(f"/v1.0/me/mailFolders('{folder}')/{suffix}")
+ if path.rstrip("/") not in allowed:
+ raise self._error("Outlook continuation changed the native collection")
+
+ def _physical_folder(self, folder: OutlookFolder) -> bool:
+ """Virtual views duplicate physical messages but may contain real child folders."""
+ if folder.odata_type == "#microsoft.graph.mailSearchFolder":
+ if folder.id in self.included_ids or folder.id in self.excluded_ids:
+ raise self._error("Outlook search-folder selections are not supported")
+ return False
+ if folder.odata_type != "#microsoft.graph.mailFolder":
+ raise self._error("Outlook unknown folder type is not qualified for delta capture")
+ return True
+
+ async def _topology_page(self, state: OutlookMailContinuation) -> OutlookMailContinuation:
+ """Enumerate each folder's hidden children; topology never owns message visibility."""
+ folder = state.folders_to_visit[0]
+ collection = (
+ f"{BASE}/mailFolders/{quote(folder, safe='')}/childFolders"
+ if folder
+ else f"{BASE}/mailFolders"
+ )
+ url = state.next_link or collection
+ self._collection_url(url, collection)
+ raw = await self.graph.get(
+ url,
+ params=None
+ if state.next_link
+ else {"includeHiddenFolders": "true", "$select": "id", "$top": 25},
+ immutable_ids=True,
+ )
+ try:
+ page = OutlookFolderPage.model_validate(raw)
+ except ValidationError:
+ raise self._error("Invalid Outlook folder inventory") from None
+ discovered = list(state.discovered_folders)
+ physical = list(state.remaining_folders)
+ queue = list(state.folders_to_visit)
+ for item in page.value:
+ if self._physical_folder(item):
+ physical.append(item.id)
+ if item.id in discovered:
+ raise InvalidCaptureCheckpoint("Outlook folder topology repeated an identity")
+ discovered.append(item.id)
+ queue.append(item.id)
+ if page.next_link:
+ self._collection_url(page.next_link, collection)
+ if page.next_link == url:
+ raise self._error("Outlook topology continuation did not advance")
+ else:
+ queue.pop(0)
+ if queue:
+ return state.model_copy(
+ update={
+ "folders_to_visit": tuple(queue),
+ "discovered_folders": tuple(discovered),
+ "remaining_folders": tuple(physical),
+ "next_link": page.next_link,
+ }
+ )
+ if set(state.folder_links) - set(physical):
+ raise InvalidCaptureCheckpoint("Outlook folder disappeared; mailbox baseline required")
+ return state.model_copy(
+ update={
+ "phase": "messages",
+ "folders_to_visit": (),
+ "discovered_folders": (),
+ "remaining_folders": tuple(sorted(physical)),
+ "next_link": None,
+ }
+ )
+
+ def _page(
+ self, state: OutlookMailContinuation, records: tuple[CaptureRecord, ...] = ()
+ ) -> CapturePage:
+ continuation = self._continuation(state)
+ final = state.phase == "messages" and not state.remaining_folders
+ # Check the final encoding as well: checkpoints escape Unicode while cursors do not.
+ try:
+ checkpoint = ProviderCheckpoint(
+ value=OutlookMailCheckpoint(
+ fingerprint=self.fingerprint, folder_links=state.folder_links
+ ).model_dump(mode="json")
+ )
+ except ValidationError:
+ raise self._error("Outlook checkpoint exceeds the 64 KiB capacity") from None
+ return CapturePage(
+ records=records,
+ continuation=continuation,
+ final=final,
+ provider_checkpoint=checkpoint if final else None,
+ )
+
+ async def _delta_page(self, state: OutlookMailContinuation) -> OutlookMailContinuation:
+ """Interpret one opaque folder delta page before exact hydration."""
+ folder = state.remaining_folders[0]
+ collection = self._delta_url(folder)
+ url = state.next_link or state.folder_links.get(folder) or collection
+ self._collection_url(url, collection)
+ try:
+ raw = await self.graph.get(
+ url,
+ params={"$select": "id"} if url == collection else None,
+ immutable_ids=True,
+ )
+ except (OutlookDeltaExpiredError, SourceGoneError):
+ raise InvalidCaptureCheckpoint(
+ "Outlook folder delta requires a fresh baseline"
+ ) from None
+ try:
+ page = OutlookDeltaPage.model_validate(raw)
+ except ValidationError:
+ raise self._error("Invalid Outlook message delta") from None
+ link = page.next_link or page.delta_link
+ self._collection_url(link, collection)
+ if page.next_link == url:
+ raise self._error("Outlook delta continuation did not advance")
+ links = dict(state.folder_links)
+ if page.delta_link:
+ links[folder] = page.delta_link
+ state = state.model_copy(
+ update={
+ "folder_links": links,
+ "next_link": page.next_link,
+ "pending_ids": tuple(dict.fromkeys(item.id for item in page.value)),
+ }
+ )
+ self._page(state) # Enforce working-state capacity before MIME/blob work.
+ return state
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """One topology page or exact message observation shares its atomic cursor commit."""
+ self._require_principal()
+ if scope != CompletedScope(record_type="message") or parent is not None:
+ raise self._error("Outlook requires the mailbox message root")
+ state = self._state(continuation)
+ if state.phase == "topology":
+ return self._page(await self._topology_page(state))
+ if not state.remaining_folders:
+ return self._page(state)
+ if not state.pending_ids:
+ state = await self._delta_page(state)
+ records = ()
+ if state.pending_ids:
+ observation = await self.message(state.pending_ids[0], files=files)
+ # Full omissions exact-refresh known roots. Changes must carry scope withdrawals.
+ if observation.kind == "upsert" or state.mode == "changes":
+ records = (observation,)
+ state = state.model_copy(update={"pending_ids": state.pending_ids[1:]})
+ if not state.pending_ids and state.next_link is None:
+ state = state.model_copy(update={"remaining_folders": state.remaining_folders[1:]})
+ return self._page(state, records)
+
+ @staticmethod
+ def _date(raw: str | None) -> datetime | None:
+ if raw is None:
+ return None
+ try:
+ return TypeAdapter(AwareDatetime).validate_python(raw)
+ except ValidationError:
+ raise OutlookBoundaryError(
+ "Invalid Outlook message timestamp", source_short_name="outlook_mail"
+ ) from None
+
+ async def message(self, native_id: str, *, files: FileService) -> CaptureRecord:
+ """Exact current mailbox state decides membership; ambiguous404 never means deletion."""
+ self._require_principal()
+ url = f"{BASE}/messages/{quote(native_id, safe='')}"
+ raw = await self.graph.get(url, immutable_ids=True)
+ try:
+ message = OutlookMessage.model_validate(raw)
+ except ValidationError:
+ raise OutlookBoundaryError(
+ "Invalid Outlook message metadata", source_short_name="outlook_mail"
+ ) from None
+ if message.id != native_id:
+ raise OutlookBoundaryError(
+ "Outlook message identity mismatch", source_short_name="outlook_mail"
+ )
+ identity = RecordIdentity(record_type="message", native_id=native_id)
+ outside = message.parentFolderId in self.excluded_ids or (
+ bool(self.included_ids) and message.parentFolderId not in self.included_ids
+ )
+ received = self._date(message.receivedDateTime)
+ if self.after is not None and not outside:
+ if received is None:
+ raise OutlookBoundaryError(
+ "Outlook date selection requires a received timestamp",
+ source_short_name="outlook_mail",
+ )
+ if received < self.after:
+ outside = True
+ if outside:
+ return CaptureRecord(
+ identity=identity,
+ payload={"id": native_id},
+ kind="delete",
+ removal_reason="scope_removed",
+ observed_at=datetime.now(timezone.utc),
+ )
+ blobs = ()
+ try:
+ content = await self.graph.mime_bytes(
+ f"{url}/$value", max_bytes=files.MAX_FILE_SIZE_BYTES
+ )
+ blob = await files.store_canonical_blob(content, media_type="message/rfc822")
+ blobs = (blob,)
+ except FileSkippedException:
+ pass # Explicit metadata-only original; transport/storage failures are not skips.
+ current = await self.graph.get(
+ url, params={"$select": "id,changeKey,parentFolderId"}, immutable_ids=True
+ )
+ if any(current.get(key) != raw.get(key) for key in ("id", "changeKey", "parentFolderId")):
+ raise OutlookBoundaryError(
+ "Outlook message changed during original capture", source_short_name="outlook_mail"
+ )
+ return CaptureRecord(
+ identity=identity,
+ payload=raw,
+ payload_schema_version=1,
+ observed_at=datetime.now(timezone.utc),
+ blobs=blobs,
+ completeness="partial" if blobs else "metadata_only",
+ source_created_at=self._date(message.createdDateTime),
+ source_updated_at=self._date(message.lastModifiedDateTime),
+ )
+
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Exact omitted-root reads must succeed before completion."""
+ if (
+ record.identity.record_type != "message"
+ or record.identity.container_id is not None
+ or record.parent is not None
+ ):
+ raise OutlookBoundaryError(
+ "Outlook omission identity is outside mailbox scope",
+ source_short_name="outlook_mail",
+ )
+ return await self.message(record.identity.native_id, files=files)
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Messages own their MIME originals; there are no independently scanned children."""
+ raise OutlookBoundaryError(
+ "Outlook message capture has no child scopes", source_short_name="outlook_mail"
+ )
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """The engine must use exact omission observations, never an absence assertion."""
+ raise OutlookBoundaryError(
+ "Outlook requires exact omission validation", source_short_name="outlook_mail"
+ )
diff --git a/backend/airweave/platform/sources/outlook_mail_models.py b/backend/airweave/platform/sources/outlook_mail_models.py
new file mode 100644
index 0000000000..05d2a88572
--- /dev/null
+++ b/backend/airweave/platform/sources/outlook_mail_models.py
@@ -0,0 +1,147 @@
+"""Typed Outlook native metadata and bounded full-capture continuation."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+class OutlookAddress(BaseModel):
+ """Native email address; sender and from retain separate meanings."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ name: str | None = None
+ address: str | None = None
+
+
+class OutlookRecipient(BaseModel):
+ """Graph recipient envelope."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ emailAddress: OutlookAddress
+
+
+class OutlookBody(BaseModel):
+ """The full native body, distinct from bodyPreview."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ contentType: Literal["text", "html"]
+ content: str
+
+
+class OutlookMessage(BaseModel):
+ """Validated native fields; capture retains the original JSON rather than a model dump."""
+
+ model_config = ConfigDict(extra="ignore", strict=True, populate_by_name=True)
+ id: str = Field(min_length=1)
+ changeKey: str = Field(min_length=1)
+ parentFolderId: str = Field(min_length=1)
+ subject: str | None = None
+ conversationId: str | None = None
+ internetMessageId: str | None = None
+ sender: OutlookRecipient | None = None
+ from_: OutlookRecipient | None = Field(default=None, alias="from")
+ replyTo: list[OutlookRecipient] = Field(default_factory=list)
+ toRecipients: list[OutlookRecipient] = Field(default_factory=list)
+ ccRecipients: list[OutlookRecipient] = Field(default_factory=list)
+ bccRecipients: list[OutlookRecipient] = Field(default_factory=list)
+ body: OutlookBody | None = None
+ bodyPreview: str | None = None
+ hasAttachments: bool = False
+ isDraft: bool = False
+ receivedDateTime: str | None = None
+ sentDateTime: str | None = None
+ createdDateTime: str | None = None
+ lastModifiedDateTime: str | None = None
+ webLink: str | None = None
+
+
+class OutlookMessageID(BaseModel):
+ """A list result is discovery, never an authoritative current message payload."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(min_length=1)
+
+
+class OutlookMessagePage(BaseModel):
+ """Explicit array required: missing metadata cannot become an empty successful page."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ value: list[OutlookMessageID] = Field(max_length=500)
+ next_link: str | None = Field(default=None, alias="@odata.nextLink", min_length=1)
+
+
+class OutlookFolder(OutlookMessageID):
+ """Folder topology includes hidden folders and explicitly identifies search views."""
+
+ # OData minimal metadata omits the declared base type, but requires derived types.
+ odata_type: str = Field(default="#microsoft.graph.mailFolder", alias="@odata.type")
+
+
+class OutlookFolderPage(BaseModel):
+ """One native topology page; missing arrays never mean an empty mailbox."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ value: list[OutlookFolder] = Field(max_length=500)
+ next_link: str | None = Field(default=None, alias="@odata.nextLink", min_length=1)
+
+
+class OutlookDeltaPage(OutlookMessagePage):
+ """Both ordinary and removed IDs require exact mailbox hydration."""
+
+ delta_link: str | None = Field(default=None, alias="@odata.deltaLink", min_length=1)
+
+ @model_validator(mode="after")
+ def terminal_or_next(self):
+ """Graph supplies exactly one continuation or terminal link."""
+ if (self.next_link is None) == (self.delta_link is None):
+ raise ValueError("Delta page requires exactly one next or delta link")
+ return self
+
+
+class OutlookMailCheckpoint(BaseModel):
+ """One complete mailbox round; byte capacity is enforced by ProviderCheckpoint."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ version: Literal[2] = 2
+ fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
+ folder_links: dict[str, str] = Field(default_factory=dict)
+
+ @model_validator(mode="after")
+ def nonempty_links(self):
+ """An absent folder token cannot masquerade as a completed baseline."""
+ if any(not folder or not link for folder, link in self.folder_links.items()):
+ raise ValueError("Folder checkpoints require nonempty identities and links")
+ return self
+
+
+class OutlookMailContinuation(OutlookMailCheckpoint):
+ """Working topology and one pending delta page share the existing bounded cursor."""
+
+ mode: Literal["full", "changes"] = "full"
+ phase: Literal["topology", "messages"] = "topology"
+ # Empty identity represents only the mailbox topology root.
+ folders_to_visit: tuple[str, ...] = ("",)
+ discovered_folders: tuple[str, ...] = ()
+ # Physical folders accumulated during topology, then consumed during message capture.
+ remaining_folders: tuple[str, ...] = ()
+ pending_ids: tuple[str, ...] = Field(default=(), max_length=500)
+ next_link: str | None = Field(default=None, min_length=1)
+
+ @model_validator(mode="after")
+ def coherent_progress(self):
+ """Malformed progress must not silently skip discovery or pending originals."""
+ for values in (self.discovered_folders, self.remaining_folders, self.pending_ids):
+ if len(values) != len(set(values)) or any(not value for value in values):
+ raise ValueError("Outlook progress identities must be nonempty and unique")
+ if len(self.folders_to_visit) != len(set(self.folders_to_visit)):
+ raise ValueError("Outlook topology queue must be unique")
+ if self.phase == "topology":
+ if not self.folders_to_visit or self.pending_ids:
+ raise ValueError("Invalid topology phase progress")
+ if not set(self.remaining_folders).issubset(self.discovered_folders):
+ raise ValueError("Physical folders must have been discovered")
+ elif self.folders_to_visit or self.discovered_folders:
+ raise ValueError("Message phase cannot contain topology work")
+ elif not self.remaining_folders and (self.pending_ids or self.next_link):
+ raise ValueError("Completed message progress cannot contain pending work")
+ return self
diff --git a/backend/airweave/platform/sources/pipedrive.py b/backend/airweave/platform/sources/pipedrive.py
index a8903616d5..69ce0a9e0b 100644
--- a/backend/airweave/platform/sources/pipedrive.py
+++ b/backend/airweave/platform/sources/pipedrive.py
@@ -27,12 +27,12 @@
PipedriveProductEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
diff --git a/backend/airweave/platform/sources/powerpoint.py b/backend/airweave/platform/sources/powerpoint.py
index de1074e8cf..f05acc8a76 100644
--- a/backend/airweave/platform/sources/powerpoint.py
+++ b/backend/airweave/platform/sources/powerpoint.py
@@ -35,12 +35,12 @@
from airweave.platform.entities._base import BaseEntity
from airweave.platform.entities.powerpoint import PowerPointPresentationEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/records/calendar_pages.py b/backend/airweave/platform/sources/records/calendar_pages.py
new file mode 100644
index 0000000000..18ddefb587
--- /dev/null
+++ b/backend/airweave/platform/sources/records/calendar_pages.py
@@ -0,0 +1,312 @@
+"""Native Calendar pages under the existing fenced mixed-scope capture engine."""
+
+import hashlib
+import json
+from typing import Literal
+from urllib.parse import quote
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+from airweave.domains.entities.canonical.calendar import CalendarScopeContext
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CycleConfiguration,
+ ProviderCheckpoint,
+)
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidScopeCheckpoint,
+ RequiredScopeAccessLost,
+ ScopeAccessLost,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation, ScanState
+from airweave.domains.entities.canonical.scope_execution import ScopePlan
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError, SourceGoneError
+from airweave.platform.configs.config import CalendarOccurrenceWindow, GoogleCalendarConfig
+from airweave.platform.sources.records.google_calendar import BASE, CalendarPage, GetJSON, record
+
+
+class CalendarProgress(BaseModel):
+ """Bounded per-scope continuation; original payloads remain only in the store."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ version: Literal[1] = 1
+ mode: Literal["full", "changes"]
+ context: CalendarScopeContext
+ sync_token: str | None = None
+ page_token: str | None = None
+ recent_tokens: tuple[str, ...] = Field(default=(), max_length=128)
+ selected_seen: tuple[str, ...] = Field(default=(), max_length=250)
+ selection_unavailable: bool = False
+ pages: int = Field(default=0, ge=0)
+ records: int = Field(default=0, ge=0)
+ bytes: int = Field(default=0, ge=0)
+
+ def continuation(self) -> ScanContinuation:
+ """Validate the complete continuation against the engine byte budget."""
+ return ScanContinuation(value=self.model_dump(mode="json"))
+
+
+class CalendarPages:
+ """Acquisition only; SQL owns every page, scope publication and cycle transition."""
+
+ def __init__(self, get: GetJSON, config: GoogleCalendarConfig):
+ """Reuse the source authenticated native HTTP boundary."""
+ self.get, self.config = get, config
+
+ @property
+ def configuration(self) -> CycleConfiguration:
+ """Fingerprint native selection and capture semantics, not moving wall-clock bounds."""
+ raw = {"version": 3, "config": self.config.model_dump(mode="json")}
+ return CycleConfiguration(
+ fingerprint=hashlib.sha256(json.dumps(raw, sort_keys=True).encode()).hexdigest(),
+ parents={
+ "calendar": (None,),
+ "event": ("calendar",),
+ "event_occurrence": ("calendar",),
+ },
+ scope_changes=("event",),
+ )
+
+ def prepare(self) -> CapturePlan:
+ """Resolve a rolling window once, before the durable cycle is created."""
+ return CapturePlan(
+ mode="mixed",
+ source_plan={"window": self.config.resolved_window().model_dump(mode="json")},
+ )
+
+ def scope_plan(
+ self,
+ scope: CompletedScope,
+ cycle: CaptureCycle,
+ previous: ScanState | None,
+ *,
+ parent: SourceRecord | None,
+ force_full: bool,
+ ) -> ScopePlan:
+ """Reuse native tokens only with compatible completed full and current publication."""
+ params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"}
+ timezone = None
+ if scope.record_type != "calendar":
+ if (
+ parent is None
+ or parent.identity != scope.parent
+ or scope.container_id != parent.identity.native_id
+ ):
+ raise ValueError("Calendar child requires its exact native owner")
+ params = {"maxResults": 250, "showDeleted": "true", "singleEvents": "false"}
+ if scope.record_type == "event_occurrence":
+ window = CalendarOccurrenceWindow.model_validate(cycle.source_plan["window"])
+ timezone = parent.payload.get("timeZone")
+ if not isinstance(timezone, str) or not timezone:
+ raise ValueError("Calendar has no timezone for occurrence capture")
+ params.update(
+ singleEvents="true",
+ timeMin=window.start.isoformat(),
+ timeMax=window.end.isoformat(),
+ timeZone=timezone,
+ )
+ elif scope.record_type != "event":
+ raise ValueError("Unknown Calendar scope")
+ context = CalendarScopeContext(parameters=params, timezone=timezone).model_dump(mode="json")
+ proposed = ScopePlan(request_context=context)
+ execution = previous.execution if previous else None
+ if (
+ scope.record_type == "event"
+ and not force_full
+ and previous
+ and previous.fingerprint == cycle.configuration.fingerprint
+ and execution
+ and execution.last_full
+ and execution.published
+ and execution.last_full.policy == "exhaustive"
+ and execution.published.policy == "exhaustive"
+ and execution.last_full.matches(proposed, previous.parent_visibility_epoch)
+ and execution.published.matches(proposed, previous.parent_visibility_epoch)
+ and execution.published.checkpoint
+ ):
+ return proposed.model_copy(
+ update={"mode": "changes", "starting_checkpoint": execution.published.checkpoint}
+ )
+ return proposed
+
+ @staticmethod
+ def initial(plan: ScopePlan) -> ScanContinuation:
+ """Only the persisted scope plan determines native page parameters."""
+ token = None
+ if plan.starting_checkpoint:
+ token = plan.starting_checkpoint.value.get("sync_token")
+ if not isinstance(token, str) or not token:
+ raise ValueError("Calendar scope checkpoint has no native sync token")
+ return CalendarProgress(
+ mode=plan.mode,
+ context=CalendarScopeContext.model_validate(plan.request_context),
+ sync_token=token,
+ ).continuation()
+
+ async def page(self, scope: CompletedScope, continuation: ScanContinuation) -> CapturePage:
+ """Fetch bounded native data; no cursor or coverage writes happen here."""
+ state = CalendarProgress.model_validate(continuation.value)
+ if state.selection_unavailable:
+ raise ValueError("A requested calendar is not accessible in current membership")
+ native = await self._fetch(scope, state)
+ root = scope.record_type == "calendar"
+ if len(native.items) > 250 or (native.nextPageToken and native.nextSyncToken):
+ raise ValueError("Calendar returned an invalid bounded page")
+ observations = tuple(
+ record(scope.record_type, item, scope.container_id) for item in native.items
+ )
+ if len({item.identity for item in observations}) != len(observations):
+ raise ValueError("Calendar page repeated a native identity")
+ seen, unavailable = set(state.selected_seen), False
+ if root:
+ observations, seen, unavailable = await self._selected(native, state, observations)
+ token_history = state.recent_tokens
+ if native.nextPageToken:
+ digest = hashlib.sha256(native.nextPageToken.encode()).hexdigest()
+ if digest in token_history:
+ raise ValueError("Calendar returned a recent repeated page token")
+ token_history = (*token_history, digest)[-128:]
+ count = state.records + len(observations)
+ size = state.bytes + sum(len(json.dumps(item.payload).encode()) for item in observations)
+ pages = state.pages + 1
+ if scope.record_type == "event_occurrence" and (
+ count > 10000 or size > 20 * 1024 * 1024 or pages > 100
+ ):
+ raise ValueError("Expanded Calendar capture exceeds bounded horizon budget")
+ checkpoint = None
+ if scope.record_type == "event" and not native.nextPageToken:
+ checkpoint = ProviderCheckpoint(value={"sync_token": native.completed_token()})
+ progress = state.model_copy(
+ update={
+ "page_token": native.nextPageToken,
+ "recent_tokens": token_history,
+ "selected_seen": tuple(sorted(seen)),
+ "selection_unavailable": unavailable,
+ "records": count,
+ "bytes": size,
+ "pages": pages,
+ }
+ )
+ return CapturePage(
+ records=observations,
+ continuation=progress.continuation(),
+ final=not native.nextPageToken and not unavailable,
+ provider_checkpoint=checkpoint,
+ )
+
+ async def _fetch(self, scope: CompletedScope, state: CalendarProgress) -> CalendarPage:
+ params = dict(state.context.parameters)
+ if state.page_token:
+ params["pageToken"] = state.page_token
+ if state.sync_token:
+ params["syncToken"] = state.sync_token
+ root = scope.record_type == "calendar"
+ url = (
+ f"{BASE}/users/me/calendarList"
+ if root
+ else f"{BASE}/calendars/{quote(scope.container_id or '', safe='')}/events"
+ )
+ try:
+ native = (
+ CalendarPage()
+ if root and self.config.calendar_ids == ()
+ else CalendarPage.model_validate(await self.get(url, params=params))
+ )
+ except SourceGoneError as error:
+ if scope.record_type == "event" and state.mode == "changes":
+ raise InvalidScopeCheckpoint("Calendar incremental token expired") from error
+ raise
+ except SourceEntityNotFoundError as error:
+ if root:
+ raise
+ error_type = (
+ RequiredScopeAccessLost if self.config.calendar_ids is not None else ScopeAccessLost
+ )
+ raise error_type(
+ "Calendar is no longer accessible", removal_reason="scope_removed"
+ ) from error
+ return native
+
+ async def _selected(
+ self, native: CalendarPage, state: CalendarProgress, observations: tuple[CaptureRecord, ...]
+ ):
+ seen = set(state.selected_seen)
+ if self.config.calendar_ids is not None:
+ selected = set(self.config.calendar_ids)
+ observations = tuple(
+ item for item in observations if item.identity.native_id in selected
+ )
+ seen.update(item.identity.native_id for item in observations)
+ unavailable = (
+ any(item.kind == "delete" for item in observations)
+ and self.config.calendar_ids is not None
+ )
+ if not native.nextPageToken and self.config.calendar_ids is not None:
+ missing = set(self.config.calendar_ids) - seen
+ # Commit verified withdrawals before surfacing the configured-selection failure.
+ extra = []
+ for native_id in sorted(missing):
+ await self.confirm_member_absent(native_id)
+ extra.append(
+ record("calendar", {"id": native_id, "deleted": True}).model_copy(
+ update={"removal_reason": "scope_removed"}
+ )
+ )
+ present = {item.identity.native_id for item in observations}
+ observations += tuple(item for item in extra if item.identity.native_id not in present)
+ unavailable = unavailable or bool(missing)
+ return observations, seen, unavailable
+
+ async def confirm_member_absent(self, native_id: str) -> None:
+ """Only exact missing/deleted membership confirms a list omission."""
+ if self.config.calendar_ids is not None and native_id not in self.config.calendar_ids:
+ return
+ try:
+ payload = await self.get(f"{BASE}/users/me/calendarList/{quote(native_id, safe='')}")
+ except SourceEntityNotFoundError:
+ return
+ if payload.get("id") != native_id:
+ raise ValueError("Calendar returned a different membership identity")
+ if payload.get("deleted") is not True and payload.get("accessRole") not in {
+ "none",
+ "freeBusyReader",
+ }:
+ raise ValueError("Calendar listing omitted accessible membership; retry capture")
+
+
+class _NativeError(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ reason: str
+ location: str | None = None
+ locationType: str | None = None
+
+
+class _NativeErrorBody(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ errors: tuple[_NativeError, ...] = ()
+
+
+class _NativeErrorEnvelope(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ error: _NativeErrorBody
+
+
+def invalid_page_token(raw: dict) -> bool:
+ """Only explicit native pageToken parameter evidence permits resetting pagination."""
+ try:
+ error = _NativeErrorEnvelope.model_validate(raw)
+ except ValidationError:
+ return False
+ return any(
+ item.location == "pageToken"
+ and item.locationType == "parameter"
+ and item.reason in {"badRequest", "invalidArgument", "invalid"}
+ for item in error.error.errors
+ )
diff --git a/backend/airweave/platform/sources/records/google_calendar.py b/backend/airweave/platform/sources/records/google_calendar.py
new file mode 100644
index 0000000000..7651f6f31a
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_calendar.py
@@ -0,0 +1,77 @@
+"""Native Calendar resource validation; page progress belongs to the canonical engine."""
+
+from collections.abc import Awaitable, Callable
+from datetime import datetime, timezone
+
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
+
+from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event
+from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity
+
+GetJSON = Callable[..., Awaitable[dict]]
+BASE = "https://www.googleapis.com/calendar/v3"
+
+
+class CalendarPage(BaseModel):
+ """Validate pagination shape without stripping fields from original resources."""
+
+ model_config = ConfigDict(extra="ignore")
+ items: list[dict[str, JsonValue]] = Field(default_factory=list)
+ nextPageToken: str | None = None
+ nextSyncToken: str | None = None
+
+ def completed_token(self) -> str:
+ """An exhausted event listing must supply its durable continuation boundary."""
+ if not self.nextSyncToken:
+ raise ValueError("Calendar event enumeration ended without nextSyncToken")
+ return self.nextSyncToken
+
+
+def record(
+ kind: str, payload: dict[str, JsonValue], calendar_id: str | None = None
+) -> CaptureRecord:
+ """Preserve sparse cancellation identity as well as full live provider JSON."""
+ native_id = payload.get("id")
+ if not isinstance(native_id, str) or not native_id:
+ raise ValueError("Calendar returned a resource without an ID")
+ recurring_cancellation = kind == "event" and is_cancelled_recurring_event(payload)
+ if (
+ kind == "event"
+ and payload.get("status") == "cancelled"
+ and payload.get("recurringEventId")
+ and not recurring_cancellation
+ ):
+ raise ValueError("Cancelled recurring event lacks its original occurrence identity")
+ # Google requires clients to retain this exclusion for the lifetime of its series.
+ # Keep its native ID/container so cancellation and reinstatement update the same row.
+ unavailable_role = kind == "calendar" and payload.get("accessRole") in {
+ "none",
+ "freeBusyReader",
+ }
+ deleted = (
+ unavailable_role
+ or (payload.get("status") == "cancelled" and not recurring_cancellation)
+ or payload.get("deleted") is True
+ )
+ return CaptureRecord(
+ identity=RecordIdentity(record_type=kind, native_id=native_id, container_id=calendar_id),
+ payload=payload,
+ descendant_visibility_fields=("accessRole",) if kind == "calendar" else (),
+ kind="delete" if deleted else "upsert",
+ removal_reason=(
+ "access_revoked"
+ if unavailable_role
+ else "scope_removed"
+ if kind == "calendar"
+ else "provider_deleted"
+ )
+ if deleted
+ else None,
+ source_created_at=TypeAdapter(AwareDatetime).validate_python(payload["created"])
+ if payload.get("created")
+ else None,
+ source_updated_at=TypeAdapter(AwareDatetime).validate_python(payload["updated"])
+ if payload.get("updated")
+ else None,
+ observed_at=datetime.now(timezone.utc),
+ )
diff --git a/backend/airweave/platform/sources/records/google_docs_content.py b/backend/airweave/platform/sources/records/google_docs_content.py
new file mode 100644
index 0000000000..c627dcd399
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_docs_content.py
@@ -0,0 +1,136 @@
+"""Bounded native Docs acquisition under the owning Drive file transaction."""
+
+import json
+from urllib.parse import quote
+
+import httpx
+from pydantic import JsonValue, TypeAdapter
+
+from airweave.domains.entities.canonical.requests import BlobReference, CaptureRecord
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.http_client.bounded_response import bounded_response_bytes
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.records.workspace_manifest import (
+ MANIFEST_MIME,
+ DocsGap,
+ DocsState,
+ ExportState,
+ WorkspaceManifestV1,
+ canonical_json,
+ document_media_gaps,
+ parse_manifest,
+ validate_document,
+)
+
+MAX_DOCS_BYTES = 32 * 1024 * 1024
+_DOCUMENT = TypeAdapter(dict[str, JsonValue])
+
+
+async def read_native_document(
+ file_id: str, *, client: AirweaveHttpClient, auth: SourceAuthProvider
+) -> dict[str, JsonValue]:
+ """No field mask: preserve native tab structure, including unknown returned fields."""
+ url = "https://docs.googleapis.com/v1/documents/" + quote(file_id, safe="")
+ parameters = {
+ "includeTabsContent": "true",
+ "suggestionsViewMode": "DEFAULT_FOR_CURRENT_ACCESS",
+ "commentsViewMode": "COMMENTS_VIEW_MODE_OMITTED",
+ }
+ for attempt in range(2):
+ headers = await authorization_headers(auth, refresh=attempt == 1)
+ headers["Accept-Encoding"] = "identity"
+ async with client.stream(
+ "GET", url, headers=headers, params=parameters, timeout=60.0
+ ) as response:
+ if response.status_code == 401 and auth.supports_refresh and attempt == 0:
+ continue
+ # Classify errors without reading an unbounded/private error body. A large
+ # 401/429/5xx can never become a successful missing-representation gap.
+ if not response.is_success:
+ raise_for_status(
+ httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ content=b"",
+ request=response.request,
+ ),
+ source_short_name="google_drive",
+ token_provider_kind=auth.provider_kind,
+ )
+ if response.headers.get("content-encoding", "identity").lower() != "identity":
+ raise ValueError("Docs capture requires identity response encoding")
+ body = await bounded_response_bytes(response, MAX_DOCS_BYTES, label="Docs JSON")
+ document = _DOCUMENT.validate_python(json.loads(body))
+ validate_document(document, file_id=file_id)
+ return document
+ raise ValueError("Docs authentication refresh did not complete")
+
+
+async def capture_document_parts(
+ record: CaptureRecord,
+ *,
+ export: ExportState,
+ export_blob: BlobReference | None,
+ files: FileService,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+) -> CaptureRecord:
+ """Stage every part then manifest; the caller still must check final Drive version."""
+ parts = [export_blob] if export_blob is not None else []
+ try:
+ document = await read_native_document(record.identity.native_id, client=client, auth=auth)
+ except FileSkippedException:
+ native = DocsState(
+ status="unavailable",
+ embedded_media="not_retained",
+ missing=(DocsGap(reason="read_size_limit"),),
+ )
+ else:
+ blob = await files.store_canonical_blob(
+ canonical_json(document), media_type="application/json"
+ )
+ parts.append(blob)
+ gaps = document_media_gaps(document)
+ native = DocsState(
+ status="complete",
+ document_blob=blob.sha256,
+ embedded_media="not_retained" if gaps else "retained",
+ missing=gaps,
+ )
+ version = record.payload["version"]
+ manifest = WorkspaceManifestV1(
+ file_id=record.identity.native_id, drive_version=version, export=export, native=native
+ )
+ content = canonical_json(manifest.model_dump(mode="json"))
+ manifest_blob = (
+ await files.store_canonical_blob(content, media_type=MANIFEST_MIME)
+ ).model_copy(update={"role": "representation_manifest"})
+ # Workspace-local identity is bytes, never descriptor order. Other providers keep MIME paths.
+ unique: dict[str, BlobReference] = {}
+ for part in parts:
+ previous = unique.get(part.sha256)
+ if previous is not None and (previous.key, previous.size_bytes) != (
+ part.key,
+ part.size_bytes,
+ ):
+ raise ValueError("Workspace identical digest has conflicting storage descriptors")
+ if previous is not None and previous.media_type != part.media_type:
+ part = part.model_copy(update={"media_type": "application/octet-stream"})
+ unique[part.sha256] = part
+ blobs = tuple(unique[key] for key in sorted(unique)) + (manifest_blob,)
+ parse_manifest(content, file_id=record.identity.native_id, drive_version=version, blobs=blobs)
+ complete = native.status == "complete" and not native.missing
+ return CaptureRecord.model_validate(
+ record.model_dump()
+ | {
+ "blobs": blobs,
+ "content_hash": manifest_blob.sha256,
+ "completeness": "complete" if complete else "partial",
+ }
+ )
diff --git a/backend/airweave/platform/sources/records/google_drive.py b/backend/airweave/platform/sources/records/google_drive.py
new file mode 100644
index 0000000000..36bb28b45b
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_drive.py
@@ -0,0 +1,33 @@
+"""Native Drive identity and original record mapping shared with body capture."""
+
+from collections.abc import Awaitable, Callable
+from datetime import datetime, timezone
+
+from pydantic import JsonValue
+
+from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity
+from airweave.platform.entities.google_drive import _parse_drive_dt
+
+BASE = "https://www.googleapis.com/drive/v3"
+GetJSON = Callable[..., Awaitable[dict]]
+
+
+def file_record(payload: dict[str, JsonValue], *, removed: bool = False) -> CaptureRecord:
+ """A Drive removal may mean lost access; do not invent permanent provider deletion."""
+ native_id = payload.get("id")
+ if not isinstance(native_id, str) or not native_id:
+ raise ValueError("Drive resource has no file identity")
+ folder = payload.get("mimeType") in {
+ "application/vnd.google-apps.folder",
+ "application/vnd.google-apps.shortcut",
+ }
+ return CaptureRecord(
+ identity=RecordIdentity(record_type="file", native_id=native_id),
+ payload=payload,
+ kind="delete" if removed else "upsert",
+ removal_reason="scope_removed" if removed else None,
+ completeness="complete" if folder or removed else "metadata_only",
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=_parse_drive_dt(payload.get("createdTime")),
+ source_updated_at=_parse_drive_dt(payload.get("modifiedTime")),
+ )
diff --git a/backend/airweave/platform/sources/records/google_drive_content.py b/backend/airweave/platform/sources/records/google_drive_content.py
new file mode 100644
index 0000000000..58b7e2fdff
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_drive_content.py
@@ -0,0 +1,143 @@
+"""Immutable Drive bodies captured before a record may claim content completeness."""
+
+from urllib.parse import quote, urlencode
+
+import httpx
+
+from airweave.core.logging import ContextualLogger
+from airweave.domains.entities.canonical.requests import BlobReference, CaptureRecord
+from airweave.domains.sources.token_providers.protocol import SourceAuthProvider
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.entities.google_drive import GOOGLE_EXPORT_FORMATS
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.sources.records.google_docs_content import capture_document_parts
+from airweave.platform.sources.records.google_drive import BASE, GetJSON
+from airweave.platform.sources.records.google_sheets_content import capture_spreadsheet_parts
+from airweave.platform.sources.records.sheets_manifest import SHEETS_MIME
+from airweave.platform.sources.records.workspace_manifest import DOCS_MIME, ExportState
+
+
+async def capture_file_content(
+ record: CaptureRecord,
+ *,
+ files: FileService,
+ get: GetJSON,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ logger: ContextualLogger,
+ capture_native_sheets: bool = False,
+) -> CaptureRecord:
+ """Reuse downloads/storage; never pair bytes with metadata from another file version."""
+ if record.kind == "delete" or record.completeness == "complete":
+ return record
+ payload = record.payload
+ mime = payload.get("mimeType")
+ version = payload.get("version")
+ if not isinstance(mime, str) or not isinstance(version, str):
+ return record
+ capabilities = payload.get("capabilities")
+ if isinstance(capabilities, dict) and capabilities.get("canDownload") is False:
+ return record
+ url = f"{BASE}/files/{quote(record.identity.native_id, safe='')}"
+ target = _download_target(url, mime)
+ if target is None:
+ return record
+ download_url, media_type = target
+ blob, export_state = await _download_representation(
+ download_url, media_type, mime, files=files, client=client, auth=auth, logger=logger
+ )
+ blob = _label_representation(blob, payload.get("name"), mime)
+ has_native_parts = False
+ if mime == DOCS_MIME:
+ record = await capture_document_parts(
+ record,
+ export=export_state,
+ export_blob=blob,
+ files=files,
+ client=client,
+ auth=auth,
+ )
+ has_native_parts = True
+ elif mime == SHEETS_MIME and capture_native_sheets:
+ record = await capture_spreadsheet_parts(
+ record, export=export_state, export_blob=blob, files=files, client=client, auth=auth
+ )
+ has_native_parts = True
+ elif blob is None:
+ return record
+ latest = await get(url, params={"fields": "version", "supportsAllDrives": "true"})
+ if latest.get("version") != version:
+ raise ValueError(
+ "Drive file changed during content capture; retry before advancing checkpoint"
+ )
+ if has_native_parts:
+ return record
+ return record.model_copy(
+ update={"blobs": (blob,), "content_hash": blob.sha256, "completeness": "complete"}
+ )
+
+
+def _label_representation(
+ blob: BlobReference | None, name: object, mime: str
+) -> BlobReference | None:
+ """Native document titles need the selected export extension; binary names are exact."""
+ if blob is None or not isinstance(name, str) or not name:
+ return blob
+ export = GOOGLE_EXPORT_FORMATS.get(mime)
+ filename = name + export[1] if export and not name.lower().endswith(export[1]) else name
+ return blob.model_copy(update={"filename": filename})
+
+
+def _download_target(url: str, mime: str) -> tuple[str, str] | None:
+ if not mime.startswith("application/vnd.google-apps."):
+ return url + "?alt=media&supportsAllDrives=true", mime
+ export = GOOGLE_EXPORT_FORMATS.get(mime)
+ if export is None:
+ return None
+ return url + "/export?" + urlencode({"mimeType": export[0]}), export[0]
+
+
+async def _download_representation(
+ url: str,
+ media_type: str,
+ mime: str,
+ *,
+ files: FileService,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ logger: ContextualLogger,
+) -> tuple[BlobReference | None, ExportState]:
+ """Export-size limits are supported gaps; unrelated provider failures propagate."""
+ try:
+ blob = await files.capture_canonical_url(
+ url=url, client=client, auth=auth, logger=logger, media_type=media_type
+ )
+ except httpx.HTTPStatusError as error:
+ if not _export_limit_reached(error, mime):
+ raise
+ return None, ExportState(status="unavailable", reason="export_size_limit")
+ except FileSkippedException:
+ return None, ExportState(status="unavailable", reason="read_size_limit")
+ return blob, ExportState(status="retained", blob=blob.sha256)
+
+
+def _export_limit_reached(error: httpx.HTTPStatusError, mime: str) -> bool:
+ """Only the explicit native-export size failure permits retaining metadata."""
+ if error.response.status_code != 403 or not mime.startswith("application/vnd.google-apps."):
+ return False
+ try:
+ payload = error.response.json()
+ except ValueError:
+ return False
+ if not isinstance(payload, dict) or not isinstance(payload.get("error"), dict):
+ return False
+ errors = payload["error"].get("errors")
+ return (
+ isinstance(errors, list)
+ and bool(errors)
+ and all(
+ isinstance(item, dict) and item.get("reason") == "exportSizeLimitExceeded"
+ for item in errors
+ )
+ )
diff --git a/backend/airweave/platform/sources/records/google_drive_pages.py b/backend/airweave/platform/sources/records/google_drive_pages.py
new file mode 100644
index 0000000000..e00505e688
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_drive_pages.py
@@ -0,0 +1,261 @@
+"""Drive pages retain native identities; SQL owns acknowledgements and checkpoint promotion."""
+
+import hashlib
+from collections.abc import Awaitable, Callable
+from typing import Annotated, Literal
+from urllib.parse import quote
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
+
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CycleConfiguration,
+ ProviderCheckpoint,
+)
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ CapturePlan,
+ InvalidCaptureCheckpoint,
+)
+from airweave.domains.entities.canonical.requests import CaptureRecord
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError
+from airweave.platform.sources.records.google_drive import BASE, GetJSON, file_record
+
+NativeID = Annotated[str, Field(min_length=1, max_length=256)]
+Token = Annotated[str, Field(min_length=1, max_length=8192)]
+
+
+class _ParameterError(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ location: str | None = None
+ locationType: str | None = None
+ reason: str | None = None
+
+
+class _ErrorBody(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ errors: list[_ParameterError] = Field(default_factory=list)
+
+
+class _ErrorEnvelope(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ error: _ErrorBody
+
+
+def rejected_listing_token(raw: dict) -> bool:
+ """Only parameter-specific native evidence permits discarding an inventory token."""
+ try:
+ body = _ErrorEnvelope.model_validate(raw)
+ except ValidationError:
+ return False
+ return any(
+ item.location == "pageToken"
+ and item.locationType == "parameter"
+ and item.reason in {"invalid", "badRequest", "invalidArgument"}
+ for item in body.error.errors
+ )
+
+
+class _File(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: NativeID
+
+
+class _Change(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ changeType: Literal["file", "drive"]
+ fileId: NativeID | None = None
+
+ @model_validator(mode="after")
+ def identity(self):
+ if self.changeType == "file" and self.fileId is None:
+ raise ValueError("Drive file change omitted its identity")
+ return self
+
+
+class _Page(BaseModel):
+ model_config = ConfigDict(extra="ignore", strict=True)
+ files: list[_File] = Field(default_factory=list, max_length=100)
+ changes: list[_Change] = Field(default_factory=list, max_length=100)
+ nextPageToken: Token | None = None
+ newStartPageToken: Token | None = None
+ incompleteSearch: bool = False
+
+
+class _Boundary(BaseModel):
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ page_token: Token
+
+
+class _Progress(BaseModel):
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ version: Literal[1] = 1
+ phase: Literal["list", "changes", "done"]
+ boundary: Token
+ token: Token | None = None
+ pending: tuple[NativeID, ...] = Field(default=(), max_length=100)
+ loaded: bool = False
+ terminal: Token | None = None
+ recent: tuple[str, ...] = Field(default=(), max_length=128)
+
+ @model_validator(mode="after")
+ def valid(self):
+ if len(set(self.pending)) != len(self.pending):
+ raise ValueError("Drive pending identities must be unique")
+ if self.terminal and (self.phase != "changes" or self.token is not None):
+ raise ValueError("Drive terminal boundary conflicts with page continuation")
+ return self
+
+ def continuation(self):
+ return ScanContinuation(
+ value=type(self).model_validate(self.model_dump()).model_dump(mode="json")
+ )
+
+
+class DrivePages:
+ """Acquire bounded native pages through the existing authenticated transport."""
+
+ def __init__(self, get: GetJSON):
+ """Retain only the authenticated read function, never durable state."""
+ self.get = get
+
+ async def prepare(
+ self, previous: CaptureCycle | None, configuration: CycleConfiguration
+ ) -> CapturePlan:
+ """Choose changes only from matching completed capture evidence."""
+ if (
+ previous
+ and previous.configuration == configuration
+ and previous.last_full_capture
+ and previous.promoted_checkpoint
+ ):
+ return CapturePlan(
+ mode="changes", starting_checkpoint=previous.promoted_checkpoint.checkpoint
+ )
+ raw = await self.get(f"{BASE}/changes/startPageToken", params={"supportsAllDrives": "true"})
+ return CapturePlan(
+ starting_checkpoint=ProviderCheckpoint(
+ value=_Boundary(page_token=raw["startPageToken"]).model_dump()
+ )
+ )
+
+ @staticmethod
+ def initial(cycle: CaptureCycle) -> ScanContinuation:
+ """Resume the original persisted boundary, including after a crash."""
+ if cycle.starting_checkpoint is None:
+ raise ValueError("Drive requires a persisted native boundary")
+ boundary = _Boundary.model_validate(cycle.starting_checkpoint.value).page_token
+ return _Progress(
+ phase="changes" if cycle.mode == "changes" else "list",
+ boundary=boundary,
+ token=boundary if cycle.mode == "changes" else None,
+ ).continuation()
+
+ async def current(self, native_id: str) -> CaptureRecord:
+ """Only exact current404 proves the file unavailable to this connection."""
+ try:
+ raw = await self.get(
+ f"{BASE}/files/{quote(native_id, safe='')}",
+ params={"fields": "*", "supportsAllDrives": "true"},
+ )
+ except SourceEntityNotFoundError:
+ return file_record({"id": native_id}, removed=True)
+ if raw.get("id") != native_id:
+ raise ValueError("Drive exact read returned a different file identity")
+ return file_record(raw)
+
+ async def page(
+ self,
+ continuation: ScanContinuation,
+ hydrate: Callable[[CaptureRecord], Awaitable[CaptureRecord]],
+ ) -> CapturePage:
+ """Hydrate one file while retaining at most100 remaining native identities."""
+ state = _Progress.model_validate(continuation.value)
+ if state.phase == "done":
+ raise ValueError("Completed Drive scan cannot fetch another page")
+ if not state.pending and not state.loaded:
+ state = await self._load(state)
+ records = ()
+ if state.pending:
+ record = await hydrate(await self.current(state.pending[0]))
+ records = (record,)
+ state = state.model_copy(update={"pending": state.pending[1:]})
+ final = False
+ checkpoint = None
+ if not state.pending:
+ if state.token is not None:
+ state = state.model_copy(update={"loaded": False})
+ elif state.phase == "list":
+ state = state.model_copy(
+ update={
+ "phase": "changes",
+ "token": state.boundary,
+ "loaded": False,
+ "recent": (),
+ }
+ )
+ else:
+ if state.terminal is None:
+ raise ValueError("Drive final changes page omitted newStartPageToken")
+ checkpoint = ProviderCheckpoint(value={"page_token": state.terminal})
+ state = state.model_copy(update={"phase": "done", "terminal": None})
+ final = True
+ return CapturePage(
+ records=records,
+ continuation=state.continuation(),
+ final=final,
+ provider_checkpoint=checkpoint,
+ )
+
+ async def _load(self, state: _Progress):
+ params = {
+ "pageSize": 100,
+ "fields": "files(id),nextPageToken,incompleteSearch"
+ if state.phase == "list"
+ else "changes(fileId,changeType),nextPageToken,newStartPageToken",
+ "spaces": "drive",
+ "includeItemsFromAllDrives": "true",
+ "supportsAllDrives": "true",
+ }
+ if state.token:
+ params["pageToken"] = state.token
+ if state.phase == "list":
+ params["corpora"] = "allDrives"
+ else:
+ params.update(includeRemoved="true", includeCorpusRemovals="true")
+ page = _Page.model_validate(
+ await self.get(
+ f"{BASE}/{'files' if state.phase == 'list' else 'changes'}", params=params
+ )
+ )
+ if page.incompleteSearch:
+ raise ValueError("Drive reported incomplete enumeration; checkpoint was not advanced")
+ if any(item.changeType == "drive" for item in page.changes):
+ raise InvalidCaptureCheckpoint("Drive membership changed; fresh full capture required")
+ if (
+ state.phase == "changes"
+ and page.nextPageToken is None
+ and page.newStartPageToken is None
+ ):
+ raise ValueError("Drive changes omitted terminal boundary")
+ recent = state.recent
+ if page.nextPageToken:
+ digest = hashlib.sha256(page.nextPageToken.encode()).hexdigest()
+ if digest in recent:
+ raise ValueError("Drive pagination repeated a recent cursor")
+ recent = (*recent[-127:], digest)
+ identities = (
+ [item.id for item in page.files]
+ if state.phase == "list"
+ else [item.fileId for item in page.changes]
+ )
+ return state.model_copy(
+ update={
+ "pending": tuple(dict.fromkeys(identities)),
+ "token": page.nextPageToken,
+ "terminal": page.newStartPageToken if state.phase == "changes" else None,
+ "recent": recent,
+ "loaded": True,
+ }
+ )
diff --git a/backend/airweave/platform/sources/records/google_sheets_content.py b/backend/airweave/platform/sources/records/google_sheets_content.py
new file mode 100644
index 0000000000..1f3b251c28
--- /dev/null
+++ b/backend/airweave/platform/sources/records/google_sheets_content.py
@@ -0,0 +1,285 @@
+"""Bounded native grid acquisition; Drive owns inventory, identity and publication."""
+
+import asyncio
+import json
+from time import monotonic
+from urllib.parse import quote
+
+import httpx
+from pydantic import BaseModel, Field, JsonValue, TypeAdapter
+
+from airweave.domains.entities.canonical.requests import BlobReference, CaptureRecord
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.http_client.bounded_response import bounded_response_bytes
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.records.sheets_manifest import (
+ GridBounds,
+ GridGap,
+ GridPart,
+ SheetsState,
+ WorkspaceManifestV2,
+ parse_sheet_manifest,
+)
+from airweave.platform.sources.records.sheets_models import (
+ SheetProperties,
+ a1_range,
+ metadata_shape,
+ validate_grid,
+ validate_partition,
+)
+from airweave.platform.sources.records.workspace_manifest import (
+ MANIFEST_MIME,
+ ExportState,
+ canonical_json,
+)
+
+MAX_RESPONSE_BYTES = 4 * 1024 * 1024
+MAX_NATIVE_BYTES = 32 * 1024 * 1024
+MAX_GRID_REQUESTS = 32
+MAX_CAPTURE_SECONDS = 120
+_JSON = TypeAdapter(dict[str, JsonValue])
+
+
+async def read_native_sheet(
+ file_id: str,
+ *,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ range_name: str | None = None,
+ timeout: float = 60,
+) -> dict[str, JsonValue]:
+ """GET retains read-only scope compatibility; never interpret HTTP errors as gaps."""
+ url = "https://sheets.googleapis.com/v4/spreadsheets/" + quote(file_id, safe="")
+ params = {"commentsViewMode": "COMMENTS_VIEW_MODE_OMITTED"}
+ if range_name is not None:
+ params.update(ranges=range_name, includeGridData="true")
+ for attempt in range(2):
+ headers = await authorization_headers(auth, refresh=attempt == 1)
+ headers["Accept-Encoding"] = "identity"
+ async with client.stream(
+ "GET", url, headers=headers, params=params, timeout=timeout
+ ) as response:
+ if response.status_code == 401 and auth.supports_refresh and attempt == 0:
+ continue
+ if not response.is_success:
+ raise_for_status(
+ httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ content=b"",
+ request=response.request,
+ ),
+ source_short_name="google_drive",
+ token_provider_kind=auth.provider_kind,
+ )
+ if response.headers.get("content-encoding", "identity").lower() != "identity":
+ raise ValueError("Sheets capture requires identity response encoding")
+ body = await bounded_response_bytes(response, MAX_RESPONSE_BYTES, label="Sheets JSON")
+ data = _JSON.validate_python(json.loads(body))
+ metadata_shape(data, file_id)
+ return data
+ raise ValueError("Sheets authentication refresh did not complete")
+
+
+def remaining_ranges(properties: SheetProperties, row: int, column: int) -> tuple[GridBounds, ...]:
+ """Compact the unrequested suffix into at most two rectangles, not millions of tiles."""
+ grid = properties.gridProperties
+ if grid is None:
+ raise ValueError("GRID sheet has no dimensions")
+ common = {"sheet_id": properties.sheetId, "end_column": grid.columnCount}
+ result = []
+ end_row = min(row + 100, grid.rowCount)
+ if column < grid.columnCount and row < grid.rowCount:
+ result.append(GridBounds(**common, start_row=row, end_row=end_row, start_column=column))
+ if end_row < grid.rowCount:
+ result.append(
+ GridBounds(**common, start_row=end_row, end_row=grid.rowCount, start_column=0)
+ )
+ return tuple(result)
+
+
+def grid_tiles(properties: SheetProperties):
+ """Lazily plan fixed rectangles; large empty allocated grids are never materialized."""
+ grid = properties.gridProperties
+ if grid is None:
+ raise ValueError("GRID sheet has no dimensions")
+ for row in range(0, grid.rowCount, 100):
+ for column in range(0, grid.columnCount, 100):
+ yield GridBounds(
+ sheet_id=properties.sheetId,
+ start_row=row,
+ end_row=min(row + 100, grid.rowCount),
+ start_column=column,
+ end_column=min(column + 100, grid.columnCount),
+ )
+
+
+class GridAcquisition(BaseModel):
+ """One file's transient budget, retained descriptors and exact unknown ranges."""
+
+ started: float
+ used_bytes: int
+ requests: int = 0
+ blobs: list[BlobReference]
+ parts: list[GridPart] = Field(default_factory=list)
+ gaps: list[GridGap] = Field(default_factory=list)
+
+ async def capture_sheet(
+ self,
+ properties: SheetProperties,
+ *,
+ file_id: str,
+ files: FileService,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ ) -> None:
+ """Stop the native scan explicitly at the budget; never imply blank unknown ranges."""
+ if properties.sheetType != "GRID":
+ self.gaps.append(GridGap(sheet_id=properties.sheetId, reason="unsupported_sheet_type"))
+ return
+ for bounds in grid_tiles(properties):
+ remaining = MAX_CAPTURE_SECONDS - (monotonic() - self.started)
+ if (
+ self.requests >= MAX_GRID_REQUESTS
+ or self.used_bytes >= MAX_NATIVE_BYTES
+ or remaining <= 0
+ ):
+ self.gaps.extend(
+ GridGap(sheet_id=properties.sheetId, bounds=b, reason="capture_budget")
+ for b in remaining_ranges(properties, bounds.start_row, bounds.start_column)
+ )
+ return
+ self.requests += 1
+ try:
+ async with asyncio.timeout(remaining):
+ content = await read_grid_bytes(
+ file_id,
+ properties,
+ bounds,
+ client=client,
+ auth=auth,
+ timeout=min(60, remaining),
+ )
+ except TimeoutError:
+ self.gaps.extend(
+ GridGap(sheet_id=properties.sheetId, bounds=b, reason="capture_budget")
+ for b in remaining_ranges(properties, bounds.start_row, bounds.start_column)
+ )
+ return
+ except FileSkippedException:
+ self.gaps.append(
+ GridGap(sheet_id=properties.sheetId, bounds=bounds, reason="read_size_limit")
+ )
+ continue
+ if self.used_bytes + len(content) > MAX_NATIVE_BYTES:
+ self.gaps.append(
+ GridGap(sheet_id=properties.sheetId, bounds=bounds, reason="capture_budget")
+ )
+ self.used_bytes = MAX_NATIVE_BYTES
+ continue
+ blob = await files.store_canonical_blob(content, media_type="application/json")
+ self.used_bytes += len(content)
+ self.blobs.append(blob)
+ self.parts.append(GridPart(bounds=bounds, blob=blob.sha256))
+
+
+async def read_grid_bytes(
+ file_id: str,
+ properties: SheetProperties,
+ bounds: GridBounds,
+ *,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+ timeout: float,
+) -> bytes:
+ """Check stable native sheet identity as well as the outer Drive version fence."""
+ data = await read_native_sheet(
+ file_id,
+ client=client,
+ auth=auth,
+ range_name=a1_range(properties.title, bounds),
+ timeout=timeout,
+ )
+ response = metadata_shape(data, file_id)
+ if [sheet.properties.sheetId for sheet in response.sheets] != [properties.sheetId]:
+ raise ValueError("Requested grid returned different sheets")
+ if response.sheets[0].properties != properties:
+ raise ValueError("Sheet identity or dimensions changed during acquisition")
+ validate_grid(data, file_id, bounds)
+ return canonical_json(data)
+
+
+async def capture_spreadsheet_parts(
+ record: CaptureRecord,
+ *,
+ export: ExportState,
+ export_blob: BlobReference | None,
+ files: FileService,
+ client: AirweaveHttpClient,
+ auth: SourceAuthProvider,
+) -> CaptureRecord:
+ """Keep original responses and exact requested bounds; caller fences Drive version."""
+ started = monotonic()
+ metadata = await read_native_sheet(record.identity.native_id, client=client, auth=auth)
+ shape = metadata_shape(metadata, record.identity.native_id)
+ raw = canonical_json(metadata)
+ metadata_blob = await files.store_canonical_blob(raw, media_type="application/json")
+ blobs = [metadata_blob] + ([export_blob] if export_blob else [])
+ acquisition = GridAcquisition(started=started, used_bytes=len(raw), blobs=blobs)
+ for sheet in shape.sheets:
+ await acquisition.capture_sheet(
+ sheet.properties,
+ file_id=record.identity.native_id,
+ files=files,
+ client=client,
+ auth=auth,
+ )
+ parts, gaps, blobs = acquisition.parts, acquisition.gaps, acquisition.blobs
+ manifest = WorkspaceManifestV2(
+ file_id=record.identity.native_id,
+ drive_version=record.payload["version"],
+ export=export,
+ native=SheetsState(
+ status="partial" if gaps else "complete",
+ metadata_blob=metadata_blob.sha256,
+ parts=tuple(parts),
+ missing=tuple(gaps),
+ ),
+ )
+ validate_partition(manifest, metadata)
+ content = canonical_json(manifest.model_dump(mode="json"))
+ marked = (await files.store_canonical_blob(content, media_type=MANIFEST_MIME)).model_copy(
+ update={"role": "representation_manifest"}
+ )
+ unique: dict[str, BlobReference] = {}
+ for blob in blobs:
+ previous = unique.get(blob.sha256)
+ if previous is not None and (previous.key, previous.size_bytes) != (
+ blob.key,
+ blob.size_bytes,
+ ):
+ raise ValueError("Identical spreadsheet bytes have conflicting descriptors")
+ if previous is not None and previous.media_type != blob.media_type:
+ blob = blob.model_copy(update={"media_type": "application/octet-stream"})
+ unique[blob.sha256] = blob
+ retained = tuple(unique[key] for key in sorted(unique)) + (marked,)
+ parse_sheet_manifest(
+ content,
+ file_id=record.identity.native_id,
+ drive_version=record.payload["version"],
+ blobs=retained,
+ )
+ return CaptureRecord.model_validate(
+ record.model_dump()
+ | {
+ "blobs": retained,
+ "content_hash": marked.sha256,
+ "completeness": "partial" if gaps else "complete",
+ }
+ )
diff --git a/backend/airweave/platform/sources/records/sheets_manifest.py b/backend/airweave/platform/sources/records/sheets_manifest.py
new file mode 100644
index 0000000000..7aee0d25e3
--- /dev/null
+++ b/backend/airweave/platform/sources/records/sheets_manifest.py
@@ -0,0 +1,110 @@
+"""Versioned native spreadsheet coverage under a single Drive file identity."""
+
+import hashlib
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from airweave.domains.entities.canonical.requests import BlobReference
+from airweave.platform.sources.records.workspace_manifest import ExportState, Sha256
+
+SHEETS_MIME = "application/vnd.google-apps.spreadsheet"
+
+
+class GridBounds(BaseModel):
+ """Zero-based, exclusive bounds; missing native cells inside a read range are blank."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ sheet_id: int = Field(ge=0)
+ start_row: int = Field(ge=0)
+ end_row: int = Field(gt=0)
+ start_column: int = Field(ge=0)
+ end_column: int = Field(gt=0)
+
+ @model_validator(mode="after")
+ def nonempty(self):
+ """Reject reversed or empty native ranges."""
+ if self.end_row <= self.start_row or self.end_column <= self.start_column:
+ raise ValueError("Grid bounds must be nonempty")
+ return self
+
+
+class GridPart(BaseModel):
+ """One successfully retained response for a precise requested rectangle."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ bounds: GridBounds
+ blob: Sha256
+
+
+class GridGap(BaseModel):
+ """An explicitly unretained range or unsupported native sheet."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ sheet_id: int = Field(ge=0)
+ bounds: GridBounds | None = None
+ reason: Literal["capture_budget", "read_size_limit", "unsupported_sheet_type"]
+
+ @model_validator(mode="after")
+ def identity(self):
+ """Bind a missing range to its stated sheet."""
+ if self.reason == "unsupported_sheet_type":
+ if self.bounds is not None:
+ raise ValueError("Unsupported sheet has no grid range")
+ elif self.bounds is None or self.bounds.sheet_id != self.sheet_id:
+ raise ValueError("Missing range must belong to its sheet")
+ return self
+
+
+class SheetsState(BaseModel):
+ """Structured coverage independent of exported-file availability."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ kind: Literal["sheets"] = "sheets"
+ status: Literal["complete", "partial"]
+ metadata_blob: Sha256
+ parts: tuple[GridPart, ...] = ()
+ missing: tuple[GridGap, ...] = ()
+ comments: Literal["omitted"] = "omitted"
+ embedded_media: Literal["not_retained"] = "not_retained"
+ calculated_values: Literal["observed_at_capture"] = "observed_at_capture"
+
+ @model_validator(mode="after")
+ def coverage(self):
+ """Require explicit gaps for every partial capture."""
+ if (self.status == "partial") != bool(self.missing):
+ raise ValueError("Grid status must agree with missing ranges")
+ return self
+
+
+class WorkspaceManifestV2(BaseModel):
+ """Sheets extension; immutable Docs v1 and export-only originals remain readable."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ schema_version: Literal[2] = 2
+ file_id: str = Field(min_length=1)
+ drive_version: str = Field(pattern=r"^[0-9]+$")
+ export: ExportState
+ native: SheetsState
+
+
+def parse_sheet_manifest(
+ content: bytes, *, file_id: str, drive_version: str, blobs: tuple[BlobReference, ...]
+) -> WorkspaceManifestV2:
+ """All semantic descriptors must resolve to exactly the owning record's bytes."""
+ manifest = WorkspaceManifestV2.model_validate_json(content)
+ if (manifest.file_id, manifest.drive_version) != (file_id, drive_version):
+ raise ValueError("Spreadsheet manifest identity/version mismatch")
+ marked = [blob for blob in blobs if blob.role == "representation_manifest"]
+ digest = hashlib.sha256(content).hexdigest()
+ if len(marked) != 1 or (marked[0].sha256, marked[0].size_bytes) != (digest, len(content)):
+ raise ValueError("Spreadsheet manifest is not the uniquely retained manifest")
+ by_digest = {blob.sha256: blob for blob in blobs}
+ if len(by_digest) != len(blobs) or any(blob.source_path is not None for blob in blobs):
+ raise ValueError("Spreadsheet descriptors are ambiguous")
+ refs = {manifest.native.metadata_blob, *(part.blob for part in manifest.native.parts)}
+ if manifest.export.blob:
+ refs.add(manifest.export.blob)
+ if digest in refs or refs != set(by_digest) - {digest}:
+ raise ValueError("Spreadsheet manifest does not account for every retained part")
+ return manifest
diff --git a/backend/airweave/platform/sources/records/sheets_models.py b/backend/airweave/platform/sources/records/sheets_models.py
new file mode 100644
index 0000000000..7a93c6f875
--- /dev/null
+++ b/backend/airweave/platform/sources/records/sheets_models.py
@@ -0,0 +1,176 @@
+"""Typed native Sheets shape and exact retained-grid admission."""
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError
+
+from airweave.platform.sources.records.sheets_manifest import GridBounds, WorkspaceManifestV2
+
+
+class GridDimensions(BaseModel):
+ """Allocated grid dimensions from native sheet metadata."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ rowCount: int = Field(ge=0)
+ columnCount: int = Field(ge=0)
+
+
+class SheetProperties(BaseModel):
+ """Identity and dimensions needed for version-consistent range acquisition."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ sheetId: int = Field(ge=0)
+ title: str = Field(min_length=1)
+ sheetType: str = "GRID"
+ gridProperties: GridDimensions | None = None
+
+
+class SheetMetadata(BaseModel):
+ """Native sheet metadata boundary; unmodeled fields remain in retained JSON."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ properties: SheetProperties
+
+
+class SpreadsheetMetadata(BaseModel):
+ """Workbook identity and sheet inventory used to admit grid coverage."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ spreadsheetId: str
+ sheets: tuple[SheetMetadata, ...]
+
+
+def metadata_shape(data: dict[str, JsonValue], file_id: str) -> SpreadsheetMetadata:
+ """Validate native identity without stripping fields from retained JSON."""
+ try:
+ metadata = SpreadsheetMetadata.model_validate(data)
+ except ValidationError:
+ raise ValueError("Malformed native spreadsheet metadata") from None
+ if metadata.spreadsheetId != file_id:
+ raise ValueError("Spreadsheet response belongs to another file")
+ ids = [s.properties.sheetId for s in metadata.sheets]
+ if len(ids) != len(set(ids)):
+ raise ValueError("Spreadsheet contains duplicate sheet IDs")
+ return metadata
+
+
+class NativeRow(BaseModel):
+ """Native sparse cells, preserving all returned fields."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ values: tuple[dict[str, JsonValue], ...] = ()
+
+
+class NativeGrid(BaseModel):
+ """Native sparse block with explicit zero-based offsets."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ startRow: int = Field(default=0, ge=0)
+ startColumn: int = Field(default=0, ge=0)
+ rowData: tuple[NativeRow, ...] = ()
+
+
+class SheetGrid(BaseModel):
+ """Grid blocks supplied for one requested native sheet."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ data: tuple[NativeGrid, ...] = ()
+
+
+def validate_grid(data: dict[str, JsonValue], file_id: str, bounds: GridBounds) -> None:
+ """Provider sparsity is valid only inside the explicit requested rectangle."""
+ shape = metadata_shape(data, file_id)
+ if [sheet.properties.sheetId for sheet in shape.sheets] != [bounds.sheet_id]:
+ raise ValueError("Grid response must identify its exact requested sheet")
+ sheets = data["sheets"]
+ try:
+ sheet = SheetGrid.model_validate(sheets[0])
+ except ValidationError:
+ raise ValueError("Malformed native spreadsheet grid") from None
+ seen: set[tuple[int, int]] = set()
+ for block in sheet.data:
+ if not block.rowData:
+ continue
+ if not (
+ bounds.start_row <= block.startRow < bounds.end_row
+ and bounds.start_column <= block.startColumn < bounds.end_column
+ ):
+ raise ValueError("Grid response starts outside the requested bounds")
+ if block.startRow + len(block.rowData) > bounds.end_row:
+ raise ValueError("Grid response exceeds requested rows")
+ if any(block.startColumn + len(row.values) > bounds.end_column for row in block.rowData):
+ raise ValueError("Grid response exceeds requested columns")
+ coordinates = {
+ (row_index, column_index)
+ for row_index, row in enumerate(block.rowData, start=block.startRow)
+ for column_index in range(block.startColumn, block.startColumn + len(row.values))
+ }
+ if seen.intersection(coordinates):
+ raise ValueError("Grid response contains overlapping cells")
+ seen.update(coordinates)
+
+
+def validate_partition(manifest: WorkspaceManifestV2, metadata: dict[str, JsonValue]) -> None:
+ """No overlap, missing range, foreign sheet or optimistic complete declaration."""
+ shape = metadata_shape(metadata, manifest.file_id)
+ identities = {sheet.properties.sheetId for sheet in shape.sheets}
+ if any(p.bounds.sheet_id not in identities for p in manifest.native.parts):
+ raise ValueError("Retained grid references an absent sheet")
+ if any(g.sheet_id not in identities for g in manifest.native.missing):
+ raise ValueError("Grid gap references an absent sheet")
+ for sheet in shape.sheets:
+ props = sheet.properties
+ parts = [p.bounds for p in manifest.native.parts if p.bounds.sheet_id == props.sheetId]
+ gaps = [g for g in manifest.native.missing if g.sheet_id == props.sheetId]
+ if props.sheetType != "GRID":
+ if parts or len(gaps) != 1 or gaps[0].reason != "unsupported_sheet_type":
+ raise ValueError("Unsupported sheet must have explicit unavailable coverage")
+ continue
+ dims = props.gridProperties
+ if dims is None or any(g.bounds is None for g in gaps):
+ raise ValueError("GRID coverage requires dimensions and explicit ranges")
+ bounds = parts + [g.bounds for g in gaps]
+ _validate_rectangles(bounds, dims)
+
+
+def _validate_rectangles(bounds: list[GridBounds], dims: GridDimensions) -> None:
+ area = 0
+ for index, current in enumerate(bounds):
+ if current.end_row > dims.rowCount or current.end_column > dims.columnCount:
+ raise ValueError("Grid coverage exceeds sheet dimensions")
+ area += (current.end_row - current.start_row) * (current.end_column - current.start_column)
+ for previous in bounds[:index]:
+ if max(previous.start_row, current.start_row) < min(
+ previous.end_row, current.end_row
+ ) and max(previous.start_column, current.start_column) < min(
+ previous.end_column, current.end_column
+ ):
+ raise ValueError("Grid coverage overlaps")
+ if area != dims.rowCount * dims.columnCount:
+ raise ValueError("Grid coverage does not account for every configured cell")
+
+
+class SpreadsheetCell(BaseModel):
+ """One native cell with absolute zero-based coordinates."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ row: int = Field(ge=0)
+ column: int = Field(ge=0)
+ value: dict[str, JsonValue]
+
+
+def column_name(index: int) -> str:
+ """Convert a zero-based column to its native A1 name without a sheet-title parser."""
+ result = ""
+ number = index + 1
+ while number:
+ number, remainder = divmod(number - 1, 26)
+ result = chr(65 + remainder) + result
+ return result
+
+
+def a1_range(title: str, bounds: GridBounds) -> str:
+ """Quote the native title while stable sheet IDs remain the admission authority."""
+ escaped = title.replace("'", "''")
+ return (
+ f"'{escaped}'!{column_name(bounds.start_column)}{bounds.start_row + 1}:"
+ f"{column_name(bounds.end_column - 1)}{bounds.end_row}"
+ )
diff --git a/backend/airweave/platform/sources/records/workspace_manifest.py b/backend/airweave/platform/sources/records/workspace_manifest.py
new file mode 100644
index 0000000000..32f99b82d2
--- /dev/null
+++ b/backend/airweave/platform/sources/records/workspace_manifest.py
@@ -0,0 +1,169 @@
+"""Versioned Docs capture provenance. Drive remains the sole record/access owner."""
+
+import hashlib
+import json
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator
+
+from airweave.domains.entities.canonical.requests import BlobReference
+
+Sha256 = Annotated[str, Field(pattern=r"^[a-f0-9]{64}$")]
+DOCS_MIME = "application/vnd.google-apps.document"
+MANIFEST_MIME = "application/vnd.almanac.workspace-manifest+json"
+
+
+class ExportState(BaseModel):
+ """A retained export or a specific supported acquisition limitation."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ status: Literal["retained", "unavailable"]
+ blob: Sha256 | None = None
+ reason: Literal["export_size_limit", "unsupported", "read_size_limit"] | None = None
+
+ @model_validator(mode="after")
+ def coverage(self):
+ """Require explicit, internally consistent representation coverage."""
+ if self.status == "retained":
+ if self.blob is None or self.reason is not None:
+ raise ValueError("Retained export requires only a blob")
+ elif self.blob is not None or self.reason is None:
+ raise ValueError("Unavailable export requires only a reason")
+ return self
+
+
+class DocsGap(BaseModel):
+ """An explicit missing portion of this acquisition policy."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ source_path: str | None = Field(default=None, pattern=r"^/")
+ reason: Literal["read_size_limit", "unsupported_media", "capture_budget"]
+
+
+class DocsMediaPart(BaseModel):
+ """A native response location backed by retained immutable bytes."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ source_path: str = Field(pattern=r"^/")
+ blob: Sha256
+
+
+class DocsState(BaseModel):
+ """Declared request policy and native response availability, not revision history."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ kind: Literal["docs"] = "docs"
+ status: Literal["complete", "partial", "unavailable"]
+ document_blob: Sha256 | None = None
+ include_tabs_content: Literal[True] = True
+ suggestions_view: Literal["DEFAULT_FOR_CURRENT_ACCESS"] = "DEFAULT_FOR_CURRENT_ACCESS"
+ comments_view: Literal["COMMENTS_VIEW_MODE_OMITTED"] = "COMMENTS_VIEW_MODE_OMITTED"
+ embedded_media: Literal["not_retained", "retained", "partial"]
+ media: tuple[DocsMediaPart, ...] = ()
+ missing: tuple[DocsGap, ...] = ()
+
+ @model_validator(mode="after")
+ def coverage(self):
+ """Require explicit, internally consistent representation coverage."""
+ if (self.status == "unavailable") != (self.document_blob is None):
+ raise ValueError("Docs response availability must match its blob")
+ if self.status == "unavailable" and not self.missing:
+ raise ValueError("Unavailable Docs response requires an explicit gap")
+ if self.status == "complete" and any(g.reason != "unsupported_media" for g in self.missing):
+ raise ValueError("Complete Docs fields cannot contain non-media gaps")
+ if self.embedded_media == "not_retained" and self.media:
+ raise ValueError("Unretained Docs media cannot reference bytes")
+ if self.embedded_media == "retained" and any(
+ g.reason == "unsupported_media" for g in self.missing
+ ):
+ raise ValueError("Retained Docs media cannot contain missing media")
+ if len({p.source_path for p in self.media}) != len(self.media):
+ raise ValueError("Docs media locations must be unique")
+ return self
+
+
+class WorkspaceManifestV1(BaseModel):
+ """One file/version's supported representations; native Sheets is not implemented."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ schema_version: Literal[1] = 1
+ file_id: str = Field(min_length=1)
+ drive_version: str = Field(pattern=r"^[0-9]+$")
+ export: ExportState
+ native: DocsState
+
+
+def canonical_json(value: JsonValue) -> bytes:
+ """Stable JSON preserves all native values, independent of object key order."""
+ return json.dumps(
+ value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False
+ ).encode("utf-8")
+
+
+def validate_document(document: dict[str, JsonValue], *, file_id: str) -> None:
+ """Validate identity and all-tabs shape without dropping unknown native fields."""
+ if document.get("documentId") != file_id:
+ raise ValueError("Docs response belongs to another file")
+ tabs = document.get("tabs")
+ if not isinstance(tabs, list) or not tabs:
+ raise ValueError("Docs response is missing all-tabs content")
+ seen: set[str] = set()
+ pending = list(tabs)
+ while pending:
+ tab = pending.pop()
+ if not isinstance(tab, dict) or not isinstance(tab.get("documentTab"), dict):
+ raise ValueError("Docs response has an unsupported tab structure")
+ properties = tab.get("tabProperties")
+ tab_id = properties.get("tabId") if isinstance(properties, dict) else None
+ if not isinstance(tab_id, str) or not tab_id or tab_id in seen:
+ raise ValueError("Docs tab identity is missing or duplicated")
+ seen.add(tab_id)
+ children = tab.get("childTabs", [])
+ if not isinstance(children, list):
+ raise ValueError("Docs child tabs are malformed")
+ pending.extend(children)
+
+
+def document_media_gaps(document: dict[str, JsonValue]) -> tuple[DocsGap, ...]:
+ """Embedded object bytes are not downloaded through unaudited content URLs."""
+ gaps = []
+ pending: list[tuple[str, JsonValue]] = [("", document)]
+ while pending:
+ path, value = pending.pop()
+ if isinstance(value, dict):
+ for key, child in value.items():
+ location = path + "/" + key.replace("~", "~0").replace("/", "~1")
+ if key == "embeddedObject":
+ gaps.append(DocsGap(source_path=location, reason="unsupported_media"))
+ else:
+ pending.append((location, child))
+ elif isinstance(value, list):
+ pending.extend((f"{path}/{index}", child) for index, child in enumerate(value))
+ return tuple(sorted(gaps, key=lambda gap: gap.source_path or ""))
+
+
+def parse_manifest(
+ content: bytes, *, file_id: str, drive_version: str, blobs: tuple[BlobReference, ...]
+) -> WorkspaceManifestV1:
+ """Reject foreign/missing/recursive references; Workspace alone deduplicates by digest."""
+ manifest = WorkspaceManifestV1.model_validate_json(content)
+ if manifest.file_id != file_id or manifest.drive_version != drive_version:
+ raise ValueError("Workspace manifest identity/version does not match its original")
+ marked = [blob for blob in blobs if blob.role == "representation_manifest"]
+ digest = hashlib.sha256(content).hexdigest()
+ if len(marked) != 1 or marked[0].sha256 != digest or marked[0].size_bytes != len(content):
+ raise ValueError("Workspace manifest is not the uniquely retained manifest")
+ by_digest: dict[str, BlobReference] = {}
+ for blob in blobs:
+ if blob.sha256 in by_digest:
+ raise ValueError("Workspace descriptors must be unambiguous by digest")
+ if blob.source_path is not None:
+ raise ValueError("Workspace parts cannot claim locations in native Drive metadata")
+ by_digest[blob.sha256] = blob
+ references = {part.blob for part in manifest.native.media}
+ references.update(
+ value for value in (manifest.export.blob, manifest.native.document_blob) if value
+ )
+ if digest in references or references != set(by_digest) - {digest}:
+ raise ValueError("Workspace manifest does not exactly account for its retained parts")
+ return manifest
diff --git a/backend/airweave/platform/sources/salesforce.py b/backend/airweave/platform/sources/salesforce.py
index 866defdfcc..6bf0eb19cf 100644
--- a/backend/airweave/platform/sources/salesforce.py
+++ b/backend/airweave/platform/sources/salesforce.py
@@ -34,12 +34,12 @@
_parse_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/servicenow.py b/backend/airweave/platform/sources/servicenow.py
index a2ea6695b8..b8a639b03e 100644
--- a/backend/airweave/platform/sources/servicenow.py
+++ b/backend/airweave/platform/sources/servicenow.py
@@ -29,12 +29,12 @@
ServiceNowProblemEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
TABLE_API_PATH = "/api/now/table"
diff --git a/backend/airweave/platform/sources/sharepoint.py b/backend/airweave/platform/sources/sharepoint.py
index a7a68770bf..170225d7ab 100644
--- a/backend/airweave/platform/sources/sharepoint.py
+++ b/backend/airweave/platform/sources/sharepoint.py
@@ -48,13 +48,13 @@
_parse_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/sharepoint2019v2/client.py b/backend/airweave/platform/sources/sharepoint2019v2/client.py
index 2ae7ff7eea..fbb3ee1ac1 100644
--- a/backend/airweave/platform/sources/sharepoint2019v2/client.py
+++ b/backend/airweave/platform/sources/sharepoint2019v2/client.py
@@ -17,7 +17,7 @@
from httpx_ntlm import HttpNtlmAuth
from tenacity import retry, stop_after_attempt
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_ntlm_auth_or_rate_limit_or_timeout,
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
diff --git a/backend/airweave/platform/sources/sharepoint_online/client.py b/backend/airweave/platform/sources/sharepoint_online/client.py
index fcfd945ffc..e930f1f6b3 100644
--- a/backend/airweave/platform/sources/sharepoint_online/client.py
+++ b/backend/airweave/platform/sources/sharepoint_online/client.py
@@ -16,7 +16,7 @@
import httpx
from tenacity import retry, stop_after_attempt
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
diff --git a/backend/airweave/platform/sources/sharepoint_online/source.py b/backend/airweave/platform/sources/sharepoint_online/source.py
index 332b1fd5c6..9b65c8e756 100644
--- a/backend/airweave/platform/sources/sharepoint_online/source.py
+++ b/backend/airweave/platform/sources/sharepoint_online/source.py
@@ -55,13 +55,13 @@
SharePointOnlineFileDeletionEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.microsoft_sensitivity_labels import SensitivityLabelFilter
from airweave.platform.sources.sharepoint_online.acl import extract_access_control
from airweave.platform.sources.sharepoint_online.builders import (
build_drive_entity,
diff --git a/backend/airweave/platform/sources/shopify.py b/backend/airweave/platform/sources/shopify.py
index 3bcdc5282c..979a8b537b 100644
--- a/backend/airweave/platform/sources/shopify.py
+++ b/backend/airweave/platform/sources/shopify.py
@@ -59,12 +59,12 @@
_parse_shopify_ts,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
SHOPIFY_API_VERSION = "2024-01"
diff --git a/backend/airweave/platform/sources/slab.py b/backend/airweave/platform/sources/slab.py
index 10209c0dd5..3ee859244f 100644
--- a/backend/airweave/platform/sources/slab.py
+++ b/backend/airweave/platform/sources/slab.py
@@ -29,12 +29,12 @@
from airweave.platform.decorators import source
from airweave.platform.entities.slab import SlabCommentEntity, SlabPostEntity, SlabTopicEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
SLAB_GRAPHQL_URL = "https://api.slab.com/v1/graphql"
diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py
index 4e59082856..064a8e8c5d 100644
--- a/backend/airweave/platform/sources/slack.py
+++ b/backend/airweave/platform/sources/slack.py
@@ -1,17 +1,50 @@
-"""Slack source implementation for federated search."""
+"""Slack capture of accessible conversations, with optional live search."""
from __future__ import annotations
+import hashlib
+import json
import re
+from datetime import datetime, timedelta, timezone
+from functools import partial
from typing import Any, AsyncGenerator, Dict, List, Optional
-
-from tenacity import retry, stop_after_attempt
+from uuid import UUID
+
+from pydantic import (
+ BaseModel,
+ ConfigDict,
+ Field,
+ JsonValue,
+ StrictBool,
+ TypeAdapter,
+ ValidationError,
+ model_validator,
+)
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.exceptions import SourceAuthError
-from airweave.domains.sources.token_providers.protocol import TokenProviderProtocol
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import (
+ CapturePage,
+ InvalidScanContinuation,
+ ScopeAccessLost,
+)
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+ parent_container_key,
+)
+from airweave.domains.entities.canonical.scan_models import ChildScopeObservation, ScanContinuation
+from airweave.domains.sources.exceptions import SourceAuthError, SourceError, SourceRateLimitError
+from airweave.domains.sources.token_providers.protocol import (
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage.file_service import FileService
from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder
from airweave.domains.syncs.cursors.cursor import SyncCursor
@@ -21,14 +54,42 @@
from airweave.platform.entities._base import AirweaveSystemMetadata, BaseEntity, Breadcrumb
from airweave.platform.entities.slack import SlackMessageEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import _parse_retry_after, raise_for_status
+from airweave.platform.sources.slack_content import SlackMessageFiles, capture_slack_file
+from airweave.platform.sources.slack_errors import SlackApiError
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
+
+def _message_timestamp(value: object) -> datetime | None:
+ """Decode Slack decimal seconds exactly; unavailable native dates stay unknown."""
+ if not isinstance(value, str) or not re.fullmatch(r"[0-9]{1,12}(?:\.[0-9]{1,6})?", value):
+ return None
+ seconds, _, fraction = value.partition(".")
+ try:
+ return datetime(1970, 1, 1, tzinfo=timezone.utc) + timedelta(
+ seconds=int(seconds), microseconds=int(fraction.ljust(6, "0"))
+ )
+ except OverflowError:
+ return None
+
+
+def _conversation_timestamp(value: object, *, milliseconds: bool = False) -> datetime | None:
+ """Conversation creation uses seconds; settings updates use milliseconds."""
+ if type(value) is not int or value < 0:
+ return None
+ try:
+ return datetime(1970, 1, 1, tzinfo=timezone.utc) + timedelta(
+ microseconds=value * (1000 if milliseconds else 1_000_000)
+ )
+ except OverflowError:
+ return None
+
+
# Pattern for Slack mrkdwn special sequences:
# <@U1234|display_name> → @display_name (user mention with label)
# <@U1234> → @U1234 (user mention without label)
@@ -68,6 +129,64 @@ def _replace(m: re.Match[str]) -> str:
return re.sub(r" +", " ", cleaned).strip()
+class SlackPrincipal(BaseModel):
+ """Native authorized workspace/user identity, with no display-name inference."""
+
+ model_config = ConfigDict(extra="ignore", frozen=True)
+ ok: StrictBool
+ team_id: str = Field(min_length=1, pattern=r"^\S+$")
+ user_id: str = Field(min_length=1, pattern=r"^\S+$")
+
+ @model_validator(mode="after")
+ def successful(self):
+ """Only an explicit successful auth.test response attests identity."""
+ if not self.ok:
+ raise ValueError("Slack identity request did not succeed")
+ return self
+
+
+class SlackHistoryCoverage(BaseModel):
+ """Slack omits this flag unless older history is unavailable under its limit."""
+
+ model_config = ConfigDict(extra="ignore")
+ is_limited: StrictBool = False
+
+
+class SlackPageMetadata(BaseModel):
+ """Only pagination metadata is parsed; record JSON stays untouched."""
+
+ next_cursor: str = Field(default="", max_length=4096)
+
+
+class SlackRootContinuation(BaseModel):
+ """Bounded conversation-list pagination state."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ cursor: str = Field(default="", max_length=4096)
+ recent: tuple[str, ...] = Field(default=(), max_length=16)
+
+
+class SlackMessageContinuation(BaseModel):
+ """One history page's thread queue, never the whole conversation in memory."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ history_cursor: str = Field(default="", max_length=4096)
+ history_recent: tuple[str, ...] = Field(default=(), max_length=16)
+ history_done: bool = False
+ pending_threads: tuple[str, ...] = Field(default=(), max_length=100)
+ reply_cursor: str = Field(default="", max_length=4096)
+ reply_recent: tuple[str, ...] = Field(default=(), max_length=16)
+
+
+class SlackFileContinuation(BaseModel):
+ """Position only in the retained parent inventory, never a mutable provider page."""
+
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ parent_id: UUID
+ inventory: str = Field(pattern=r"^[a-f0-9]{64}$")
+ next_index: int = Field(ge=0, le=100)
+
+
@source(
name="Slack",
short_name="slack",
@@ -81,28 +200,83 @@ def _replace(m: re.Match[str]) -> str:
config_class=SlackConfig,
labels=["Communication", "Messaging"],
supports_continuous=False,
- federated_search=True,
+ federated_search=False,
rate_limit_level=RateLimitLevel.ORG,
)
class SlackSource(BaseSource):
- """Slack source connector using federated search.
+ """Capture all accessible conversation history and replies without search truncation."""
+
+ # Class-level capabilities remain available to the search source registry.
+ canonical_record_types = ("channel", "message", "file")
+ canonical_container_parents = {"message": "channel", "file": "message"}
+ _slack_config: SlackConfig | None = None
+
+ @property
+ def slack_config(self) -> SlackConfig | None:
+ """The source configuration fixes its capture topology before a cycle starts."""
+ return self._slack_config
+
+ @slack_config.setter
+ def slack_config(self, config: SlackConfig) -> None:
+ self._slack_config = config
+ self.canonical_record_types = (
+ ("channel", "message", "file") if config.capture_files else ("channel", "message")
+ )
+ self.canonical_container_parents = (
+ {"message": "channel", "file": "message"}
+ if config.capture_files
+ else {"message": "channel"}
+ )
- Instead of syncing all messages and files, this source searches Slack at query time
- using the search.all API endpoint. This is necessary because Slack's rate limits
- are too restrictive for full synchronization.
- """
+ _verified_principal: SlackPrincipal | None = None
+
+ def _require_principal(self) -> SlackPrincipal:
+ """Unbound legacy search cannot bypass canonical capture identity checks."""
+ principal = self._verified_principal
+ if (
+ principal is None
+ or self.slack_config is None
+ or principal.team_id != self.slack_config.expected_team_id
+ or principal.user_id != self.slack_config.expected_user_id
+ ):
+ raise ValueError("Owned Slack capture requires an attested workspace and user")
+ return principal
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Existing cycle ownership includes the trusted visibility principal."""
+ principal = self._require_principal()
+ material = {"version": 2, "team_id": principal.team_id, "user_id": principal.user_id}
+ if self.slack_config.capture_files:
+ material["capture_files"] = "message_owned_children_v1"
+ fingerprint = hashlib.sha256(
+ json.dumps(
+ material,
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+ return CycleConfiguration.from_source(
+ fingerprint=fingerprint,
+ record_types=self.canonical_record_types,
+ container_parents=self.canonical_container_parents,
+ exact_parent_validation=("message",) if self.slack_config.capture_files else (),
+ )
@classmethod
async def create(
cls,
*,
- auth: TokenProviderProtocol,
+ auth: SourceAuthProvider,
logger: ContextualLogger,
http_client: AirweaveHttpClient,
config: SlackConfig,
) -> SlackSource:
"""Create a new Slack source."""
- return cls(auth=auth, logger=logger, http_client=http_client)
+ instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance.slack_config = config
+ if config.expected_team_id is not None:
+ await instance.validate()
+ return instance
# ------------------------------------------------------------------
# HTTP
@@ -110,19 +284,17 @@ async def create(
@retry(
stop=stop_after_attempt(5),
- retry=retry_if_rate_limit_or_timeout,
- wait=wait_rate_limit_with_backoff,
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
reraise=True,
)
async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict:
"""Make authenticated GET request to Slack API with token refresh support."""
- token = await self.auth.get_token()
- headers = {"Authorization": f"Bearer {token}"}
+ headers = await authorization_headers(self.auth)
response = await self.http_client.get(url, headers=headers, params=params)
if response.status_code == 401 and self.auth.supports_refresh:
- new_token = await self.auth.force_refresh()
- headers = {"Authorization": f"Bearer {new_token}"}
+ headers = await authorization_headers(self.auth, refresh=True)
response = await self.http_client.get(url, headers=headers, params=params)
raise_for_status(
@@ -130,7 +302,351 @@ async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict:
source_short_name=self.short_name,
token_provider_kind=self.auth.provider_kind,
)
- return response.json()
+ payload = response.json()
+ if not payload.get("ok"):
+ error = payload.get("error", "unknown_error")
+ if error in {"ratelimited", "rate_limited"}:
+ raise SourceRateLimitError(
+ source_short_name="slack", retry_after=_parse_retry_after(response, default=60)
+ )
+ if error in {"invalid_auth", "not_authed", "token_revoked", "account_inactive"}:
+ raise SourceAuthError(
+ "Slack authentication is unavailable",
+ source_short_name="slack",
+ status_code=401,
+ token_provider_kind=self.auth.provider_kind,
+ )
+ # Do not turn denied channels, missing scopes or provider errors into empty pages.
+ raise SlackApiError(error)
+ return payload
+
+ @staticmethod
+ def _capture_message(message: dict, channel_id: str) -> CaptureRecord:
+ """Keep native response untouched; opt-in acquisition subsequently adds blob evidence."""
+ edited = message.get("edited")
+ edited_ts = edited.get("ts") if isinstance(edited, dict) else None
+ return CaptureRecord(
+ identity=RecordIdentity(
+ record_type="message", native_id=message["ts"], container_id=channel_id
+ ),
+ payload=message,
+ observed_at=datetime.now(timezone.utc),
+ completeness="partial" if message.get("files") else "complete",
+ source_created_at=_message_timestamp(message["ts"]),
+ source_updated_at=_message_timestamp(edited_ts),
+ )
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Keep the established native container identity for this flat provider."""
+ return self._child_scope(parent.identity, record_type)
+
+ def _child_scope(self, parent: RecordIdentity, record_type: str) -> CompletedScope:
+ """Enumeration and fresh-page receipts use exactly the same scope identity."""
+ if self.canonical_container_parents.get(record_type) != parent.record_type:
+ raise ValueError("Unsupported Slack child scope")
+ return CompletedScope(
+ record_type=record_type,
+ container_id=parent_container_key(parent)
+ if record_type == "file"
+ else parent.native_id,
+ parent=parent,
+ )
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Fetch one page; record and nested reply progress are committed by the pipeline."""
+ self._require_principal()
+ try:
+ if scope.record_type == "file":
+ return await self._file_page(scope, continuation, files, parent)
+ page = await self._validated_capture_page(scope, continuation)
+ if self.slack_config.capture_files and scope.record_type == "message":
+ records = tuple(self._attachment_owner(record) for record in page.records)
+ page = page.model_copy(
+ update={
+ "records": records,
+ "child_scope_observations": tuple(
+ ChildScopeObservation(scope=self._child_scope(record.identity, "file"))
+ for record in records
+ ),
+ }
+ )
+ return page
+ except ValidationError:
+ # Validation diagnostics may contain private provider values.
+ raise ValueError("Slack capture returned invalid page or continuation data") from None
+
+ def _attachment_owner(self, record: CaptureRecord) -> CaptureRecord:
+ """Validate the native inventory shared by enumeration and exact owner reads."""
+ self._file_inventory(record.payload)
+ return record.model_copy(
+ update={
+ "payload_schema_version": 2,
+ "descendant_visibility_fields": ("files",),
+ "completeness": "complete",
+ }
+ )
+
+ async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord:
+ """Recheck an attachment owner without treating an ambiguous miss as deletion."""
+ self._require_principal()
+ if not self.slack_config.capture_files or record.identity.record_type != "message":
+ raise ValueError("Slack exact owner refresh requires attachment capture and a message")
+ try:
+ message = await self._read_known_message(record)
+ if message is None:
+ raise ValueError("Slack message access remains unconfirmed; capture is incomplete")
+ captured = self._capture_message(message, record.identity.container_id)
+ return self._attachment_owner(captured).model_copy(update={"parent": record.parent})
+ except ValidationError:
+ raise ValueError("Slack exact message read returned invalid provider data") from None
+
+ @staticmethod
+ def _file_inventory(payload: dict[str, JsonValue]) -> list[dict[str, JsonValue]]:
+ metadata = SlackMessageFiles.model_validate(payload).files
+ if len({file.id for file in metadata}) != len(metadata):
+ raise ValueError("Slack message contains duplicate file identities")
+ return payload.get("files", [])
+
+ async def _file_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ files: FileService,
+ parent: SourceRecord | None,
+ ) -> CapturePage:
+ if (
+ not self.slack_config.capture_files
+ or parent is None
+ or parent.identity.record_type != "message"
+ or parent.payload_schema_version != 2
+ or parent.content_access != "available"
+ or parent.deleted_at is not None
+ or scope != self.child_scope(parent, "file")
+ ):
+ raise ValueError("Slack file scope requires its current retained message")
+ inventory = self._file_inventory(parent.payload)
+ digest = hashlib.sha256(
+ json.dumps(inventory, sort_keys=True, separators=(",", ":")).encode()
+ ).hexdigest()
+ state = (
+ SlackFileContinuation.model_validate(continuation.value)
+ if continuation.value
+ else SlackFileContinuation(parent_id=parent.id, inventory=digest, next_index=0)
+ )
+ if state.parent_id != parent.id or state.inventory != digest:
+ raise InvalidScanContinuation("Slack retained file inventory changed")
+ if state.next_index > len(inventory):
+ raise InvalidScanContinuation("Slack file continuation exceeds retained inventory")
+ if state.next_index == len(inventory):
+ return CapturePage(records=(), continuation=continuation, final=True)
+ native = inventory[state.next_index]
+ item = CaptureRecord(
+ identity=RecordIdentity(
+ record_type="file", native_id=native["id"], container_id=scope.container_id
+ ),
+ parent=parent.identity,
+ payload=native,
+ observed_at=datetime.now(timezone.utc),
+ )
+ captured = await capture_slack_file(self, item, files)
+ following = state.model_copy(update={"next_index": state.next_index + 1})
+ return CapturePage(
+ records=(captured,),
+ continuation=ScanContinuation(value=following.model_dump(mode="json")),
+ final=following.next_index == len(inventory),
+ )
+
+ async def _validated_capture_page(
+ self, scope: CompletedScope, continuation: ScanContinuation
+ ) -> CapturePage:
+ if scope.record_type == "channel" and scope.container_id is None:
+ state = SlackRootContinuation.model_validate(continuation.value)
+ items, following, recent = await self._capture_page_response(
+ "conversations.list",
+ "channels",
+ {"types": "public_channel,private_channel,im,mpim", "exclude_archived": "false"},
+ state.cursor,
+ state.recent,
+ )
+ records = tuple(
+ CaptureRecord(
+ identity=RecordIdentity(record_type="channel", native_id=item["id"]),
+ payload=item,
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=_conversation_timestamp(item.get("created")),
+ source_updated_at=_conversation_timestamp(
+ item.get("updated"), milliseconds=True
+ ),
+ )
+ for item in items
+ )
+ return CapturePage(
+ records=records,
+ final=not following,
+ continuation=ScanContinuation(
+ value=SlackRootContinuation(cursor=following, recent=recent).model_dump(
+ mode="json"
+ )
+ ),
+ )
+ if scope.record_type != "message" or scope.container_id is None:
+ raise ValueError("Slack page capture requires a declared whole conversation scope")
+ state = SlackMessageContinuation.model_validate(continuation.value)
+ channel_id = scope.container_id
+ try:
+ if state.pending_threads:
+ items, following, recent = await self._capture_page_response(
+ "conversations.replies",
+ "messages",
+ {"channel": channel_id, "ts": state.pending_threads[0]},
+ state.reply_cursor,
+ state.reply_recent,
+ )
+ pending = state.pending_threads if following else state.pending_threads[1:]
+ state = state.model_copy(
+ update={
+ "pending_threads": pending,
+ "reply_cursor": following,
+ "reply_recent": recent if following else (),
+ }
+ )
+ else:
+ if state.history_done:
+ raise ValueError("Completed Slack continuation cannot fetch another page")
+ items, following, recent = await self._capture_page_response(
+ "conversations.history",
+ "messages",
+ {"channel": channel_id},
+ state.history_cursor,
+ state.history_recent,
+ )
+ pending = tuple(
+ dict.fromkeys(item["ts"] for item in items if item.get("reply_count", 0))
+ )
+ state = SlackMessageContinuation(
+ history_cursor=following,
+ history_recent=recent,
+ history_done=not following,
+ pending_threads=pending,
+ )
+ except SlackApiError as exc:
+ if exc.code in {"channel_not_found", "not_in_channel"}:
+ raise ScopeAccessLost("Slack conversation access was lost") from exc
+ raise
+ return CapturePage(
+ records=tuple(self._capture_message(item, channel_id) for item in items),
+ continuation=ScanContinuation(value=state.model_dump(mode="json")),
+ final=state.history_done and not state.pending_threads,
+ )
+
+ async def _capture_page_response(
+ self,
+ operation: str,
+ item_key: str,
+ params: dict[str, JsonValue],
+ cursor: str,
+ recent: tuple[str, ...],
+ ) -> tuple[list[dict[str, JsonValue]], str, tuple[str, ...]]:
+ """Validate provider pagination without rewriting native record payloads."""
+ try:
+ payload = await self._get(
+ f"https://slack.com/api/{operation}",
+ {**params, "limit": 100, **({"cursor": cursor} if cursor else {})},
+ )
+ except SlackApiError as exc:
+ if exc.code == "invalid_cursor":
+ raise InvalidScanContinuation("Slack page cursor expired") from exc
+ if operation == "conversations.replies" and exc.code == "thread_not_found":
+ # History queued a thread that no longer exists. Restart inventory;
+ # missing replies alone cannot authorize deletion or scope completion.
+ raise InvalidScanContinuation("Slack queued thread is no longer available") from exc
+ raise
+ if operation == "conversations.history":
+ coverage = SlackHistoryCoverage.model_validate(payload)
+ if coverage.is_limited:
+ raise SourceError(
+ "Slack history is limited by workspace visibility. "
+ "Capture remains incomplete; older retained messages were not reconciled.",
+ source_short_name="slack",
+ )
+ items = TypeAdapter(list[dict[str, JsonValue]]).validate_python(payload.get(item_key))
+ metadata = SlackPageMetadata.model_validate(payload.get("response_metadata") or {})
+ following = metadata.next_cursor.strip()
+ if not following and payload.get("has_more"):
+ raise ValueError("Slack capture has more records without a cursor")
+ if following:
+ digest = hashlib.sha256(following.encode()).hexdigest()
+ if following == cursor or digest in recent:
+ raise ValueError("Slack capture repeated a pagination cursor")
+ recent = (*recent[-15:], digest)
+ return items, following, recent
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Require provider confirmation before hiding an omitted prior conversation."""
+ self._require_principal()
+ if record.identity.record_type == "message":
+ await self._confirm_message_omission(record)
+ return
+ if record.identity.record_type != "channel" or record.identity.container_id is not None:
+ raise ValueError("Slack omission confirmation requires a channel or message")
+ native_id = record.identity.native_id
+ try:
+ await self._get("https://slack.com/api/conversations.info", {"channel": native_id})
+ except SlackApiError as exc:
+ if exc.code not in {"channel_not_found", "not_in_channel"}:
+ raise
+ return
+ raise ValueError("Slack omitted an accessible prior channel; retry enumeration")
+
+ async def _confirm_message_omission(self, record: SourceRecord) -> None:
+ """An exact read can disprove omission; an empty read does not prove deletion."""
+ message = await self._read_known_message(record)
+ if message is not None:
+ raise ValueError("Slack omitted an accessible prior message; retry enumeration")
+ raise ValueError("Slack message omission remains unconfirmed; capture is incomplete")
+
+ async def _read_known_message(self, record: SourceRecord) -> dict[str, JsonValue] | None:
+ """Use native timestamp bounds and the retained thread route for one exact owner."""
+ identity = record.identity
+ if (
+ identity.container_id is None
+ or record.parent
+ != RecordIdentity(record_type="channel", native_id=identity.container_id)
+ or record.payload.get("ts") != identity.native_id
+ ):
+ raise ValueError("Slack message omission requires its retained channel identity")
+ thread_ts = record.payload.get("thread_ts")
+ if thread_ts is not None and (not isinstance(thread_ts, str) or not thread_ts):
+ raise ValueError("Slack retained message has invalid thread identity")
+ params = {
+ "channel": identity.container_id,
+ "oldest": identity.native_id,
+ "latest": identity.native_id,
+ "inclusive": "true",
+ "limit": 1,
+ }
+ operation = "conversations.history"
+ if thread_ts and thread_ts != identity.native_id:
+ operation = "conversations.replies"
+ params["ts"] = thread_ts
+ payload = await self._get(f"https://slack.com/api/{operation}", params)
+ messages = TypeAdapter(list[dict[str, JsonValue]]).validate_python(payload.get("messages"))
+ matches = [message for message in messages if message.get("ts") == identity.native_id]
+ if len(matches) > 1:
+ raise ValueError("Slack exact message read returned duplicate identities")
+ if not matches:
+ return None
+ message = matches[0]
+ if message.get("thread_ts", identity.native_id) != (thread_ts or identity.native_id):
+ raise ValueError("Slack exact message read changed its retained thread identity")
+ return message
# ------------------------------------------------------------------
# Federated search
@@ -292,5 +808,17 @@ async def generate_entities(
yield # make this a generator # noqa: RUF027
async def validate(self) -> None:
- """Validate credentials by calling Slack auth.test."""
- await self._get("https://slack.com/api/auth.test")
+ """Re-attest the exact bound workspace/user, clearing old evidence first."""
+ self._verified_principal = None
+ raw = await self._get("https://slack.com/api/auth.test")
+ try:
+ principal = SlackPrincipal.model_validate(raw)
+ except ValidationError:
+ raise ValueError("Slack returned an invalid native identity") from None
+ if self.slack_config is not None and self.slack_config.expected_team_id is not None:
+ if (
+ principal.team_id != self.slack_config.expected_team_id
+ or principal.user_id != self.slack_config.expected_user_id
+ ):
+ raise ValueError("Slack identity does not match the trusted binding")
+ self._verified_principal = principal
diff --git a/backend/airweave/platform/sources/slack_content.py b/backend/airweave/platform/sources/slack_content.py
new file mode 100644
index 0000000000..70ffab8aa4
--- /dev/null
+++ b/backend/airweave/platform/sources/slack_content.py
@@ -0,0 +1,177 @@
+"""Acquire Slack originals before the existing page transaction advances its cursor."""
+
+from typing import TYPE_CHECKING
+from urllib.parse import urlsplit
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from airweave.domains.entities.canonical.requests import BlobReference, CaptureRecord
+from airweave.domains.entities.canonical.slack_files import (
+ SlackFileManifest,
+ SlackFileOutcome,
+ SlackFileReason,
+)
+from airweave.domains.sources.exceptions import (
+ SourceEntityForbiddenError,
+ SourceEntityNotFoundError,
+ SourceServerError,
+)
+from airweave.domains.storage.exceptions import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.slack_errors import SlackApiError
+
+if TYPE_CHECKING:
+ from airweave.platform.sources.slack import SlackSource
+
+
+class SlackFileMetadata(BaseModel):
+ """Parse only acquisition fields; original native JSON remains authoritative."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ id: str = Field(pattern=r"^F[A-Z0-9]+$", max_length=128)
+ name: str | None = None
+ size: int | None = Field(default=None, ge=0)
+ mimetype: str | None = Field(default=None, max_length=256)
+ is_external: bool = False
+ mode: str | None = None
+ url_private_download: str | None = Field(default=None, max_length=16384)
+ url_private: str | None = Field(default=None, max_length=16384)
+
+
+class SlackMessageFiles(BaseModel):
+ """Malformed file lists fail capture instead of quietly dropping originals."""
+
+ model_config = ConfigDict(extra="ignore", strict=True)
+ files: list[SlackFileMetadata] = Field(default_factory=list, max_length=100)
+
+
+def _original_url(file: SlackFileMetadata) -> str | None:
+ url = file.url_private_download or file.url_private
+ if not url:
+ return None
+ try:
+ parsed = urlsplit(url)
+ valid = (
+ parsed.scheme == "https"
+ and parsed.hostname in {"slack.com", "files.slack.com"}
+ and parsed.port in {None, 443}
+ and parsed.username is None
+ and parsed.password is None
+ and not parsed.fragment
+ )
+ except ValueError:
+ valid = False
+ if not valid:
+ raise ValueError("Slack original URL is outside the permitted origin")
+ return url
+
+
+async def capture_slack_file(
+ source: "SlackSource", record: CaptureRecord, files: FileService
+) -> CaptureRecord:
+ """Acquire one message-owned file; the caller commits its child page atomically."""
+ file = SlackFileMetadata.model_validate(record.payload)
+ if record.identity.record_type != "file" or record.identity.native_id != file.id:
+ raise ValueError("Slack file payload does not match its record identity")
+ file, native, reason = await _enrich(source, file)
+ blobs = []
+ if reason is None:
+ if file.is_external or file.mode == "external":
+ reason = "unsupported_external"
+ elif file.size is not None and file.size > files.MAX_FILE_SIZE_BYTES:
+ reason = "oversized"
+ elif file.size is None or not file.mimetype or (url := _original_url(file)) is None:
+ reason = "missing_metadata"
+ else:
+ blob, reason = await _download(source, file, url, files)
+ if blob is not None:
+ blobs.append(blob.model_copy(update={"source_path": ""}))
+ manifest = SlackFileManifest(
+ files=(
+ SlackFileOutcome(
+ index=0,
+ native_id=file.id,
+ outcome="captured" if reason is None else "unavailable",
+ reason=reason,
+ file=native,
+ ),
+ )
+ )
+ reference = await files.store_canonical_blob(
+ manifest.model_dump_json().encode(), media_type="application/json"
+ )
+ blobs.append(reference.model_copy(update={"role": "representation_manifest"}))
+ return record.model_copy(
+ update={
+ "blobs": tuple(blobs),
+ "completeness": "partial" if reason else "complete",
+ }
+ )
+
+
+async def _download(
+ source: "SlackSource", file: SlackFileMetadata, url: str, files: FileService
+) -> tuple[BlobReference | None, SlackFileReason | None]:
+ try:
+ blob = await files.capture_canonical_url(
+ url,
+ source.http_client,
+ source.auth,
+ source.logger,
+ media_type=file.mimetype,
+ expected_media_type=file.mimetype,
+ follow_redirects=False,
+ )
+ except FileSkippedException:
+ return None, "oversized"
+ except httpx.HTTPStatusError as exc:
+ if exc.response.status_code in {403, 404}:
+ return None, "access_denied" if exc.response.status_code == 403 else "not_found"
+ # Never put private download URLs or response bodies in domain diagnostics.
+ safe = httpx.Response(
+ exc.response.status_code,
+ headers={"Retry-After": exc.response.headers["Retry-After"]}
+ if "Retry-After" in exc.response.headers
+ else {},
+ request=httpx.Request("GET", "https://slack.com/"),
+ )
+ raise_for_status(
+ safe, source_short_name="slack", token_provider_kind=source.auth.provider_kind
+ )
+ raise
+ except httpx.RequestError:
+ raise SourceServerError(
+ "Slack original download failed", source_short_name="slack"
+ ) from None
+ if file.size is not None and blob.size_bytes != file.size:
+ raise ValueError("Slack original size changed during capture; retry the page")
+ return blob.model_copy(update={"filename": file.name or None}), None
+
+
+async def _enrich(
+ source: "SlackSource", file: SlackFileMetadata
+) -> tuple[SlackFileMetadata, dict[str, JsonValue] | None, SlackFileReason | None]:
+ if file.is_external or file.mode == "external":
+ return file, None, None
+ if _original_url(file) and file.size is not None and file.mimetype:
+ return file, None, None
+ try:
+ response = await source._get("https://slack.com/api/files.info", {"file": file.id})
+ except SourceEntityForbiddenError:
+ return file, None, "access_denied"
+ except SourceEntityNotFoundError:
+ return file, None, "not_found"
+ except SlackApiError as exc:
+ if exc.code in {"access_denied", "not_visible"}:
+ return file, None, "access_denied"
+ if exc.code in {"file_not_found", "file_deleted"}:
+ return file, None, "not_found"
+ # Missing scope is a connection/configuration failure, not evidence of file denial.
+ raise
+ native = response.get("file")
+ enriched = SlackFileMetadata.model_validate(native)
+ if enriched.id != file.id:
+ raise ValueError("Slack file metadata identity changed")
+ return enriched, native, None
diff --git a/backend/airweave/platform/sources/slack_errors.py b/backend/airweave/platform/sources/slack_errors.py
new file mode 100644
index 0000000000..fd2ce8a6b7
--- /dev/null
+++ b/backend/airweave/platform/sources/slack_errors.py
@@ -0,0 +1,10 @@
+"""Native Slack application error codes shared by message and file acquisition."""
+
+
+class SlackApiError(ValueError):
+ """A native application failure cannot imply a successful empty page."""
+
+ def __init__(self, code: str):
+ """Expose the bounded provider code to explicit access interpretation."""
+ self.code = code
+ super().__init__(f"Slack request failed: {code}")
diff --git a/backend/airweave/platform/sources/slite.py b/backend/airweave/platform/sources/slite.py
index 90e38011ba..140ce3cc91 100644
--- a/backend/airweave/platform/sources/slite.py
+++ b/backend/airweave/platform/sources/slite.py
@@ -22,12 +22,12 @@
from airweave.platform.decorators import source
from airweave.platform.entities.slite import SliteNoteEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod
SLITE_API_BASE = "https://api.slite.com/v1"
diff --git a/backend/airweave/platform/sources/stripe.py b/backend/airweave/platform/sources/stripe.py
index ae57060373..c172c41748 100644
--- a/backend/airweave/platform/sources/stripe.py
+++ b/backend/airweave/platform/sources/stripe.py
@@ -19,15 +19,27 @@
from __future__ import annotations
import base64
+import json
+from contextlib import aclosing
from datetime import datetime
+from functools import partial
from typing import AsyncGenerator, Optional
-from tenacity import retry, stop_after_attempt
+import httpx
+from pydantic import JsonValue
+from tenacity import retry, retry_if_exception_type, stop_after_attempt
from airweave.core.logging import ContextualLogger
from airweave.core.shared_models import RateLimitLevel
+from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError
from airweave.domains.browse_tree.types import NodeSelectionData
-from airweave.domains.sources.token_providers.protocol import AuthProviderKind, SourceAuthProvider
+from airweave.domains.sources.exceptions import SourceError
+from airweave.domains.sources.token_providers.protocol import (
+ AuthProviderKind,
+ ManagedAuthProvider,
+ SourceAuthProvider,
+ authorization_headers,
+)
from airweave.domains.storage.file_service import FileService
from airweave.domains.syncs.cursors.cursor import SyncCursor
from airweave.platform.configs.auth import StripeAuthConfig
@@ -49,12 +61,13 @@
_parse_stripe_ts,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.stripe_capture import StripeCapture
from airweave.schemas.source_connection import AuthenticationMethod
_parse_unix_timestamp = _parse_stripe_ts
@@ -80,6 +93,9 @@ class StripeSource(BaseSource):
including transactions, customers, subscriptions, and account analytics.
"""
+ canonical_record_types = StripeCapture.canonical_record_types
+ MAX_CAPTURE_RESPONSE_BYTES = 16 * 1024 * 1024
+
@classmethod
async def create(
cls,
@@ -91,12 +107,76 @@ async def create(
) -> StripeSource:
"""Create a new Stripe source instance."""
instance = cls(auth=auth, logger=logger, http_client=http_client)
+ if config.original_capture is not None:
+ instance._capture_page_source = await StripeCapture.create(
+ instance._read_original, config.original_capture
+ )
+ return instance
+ if isinstance(auth, ManagedAuthProvider):
+ raise SourceError(
+ "Managed Stripe requires an explicitly bound original capture configuration",
+ source_short_name="stripe",
+ )
if auth.provider_kind == AuthProviderKind.CREDENTIAL:
instance._api_key = auth.credentials.api_key
else:
instance._api_key = await auth.get_token()
return instance
+ @retry(
+ stop=stop_after_attempt(5),
+ retry=retry_if_rate_limit_or_timeout | retry_if_exception_type(AuthProviderRateLimitError),
+ wait=partial(wait_rate_limit_with_backoff, max_rate_limit_wait=None),
+ reraise=True,
+ )
+ async def _read_original(
+ self, path: str, *, params: dict[str, str | int], headers: dict[str, str]
+ ) -> dict[str, JsonValue]:
+ """Read the fixed Stripe origin using the existing managed/direct transport."""
+ if not path.startswith("/v1/") or "?" in path or "#" in path or ".." in path:
+ raise SourceError("Invalid Stripe capture path", source_short_name="stripe")
+ if self.auth.provider_kind == AuthProviderKind.CREDENTIAL:
+ credentials = StripeAuthConfig.model_validate(self.auth.credentials.model_dump())
+ auth_headers = {"Authorization": f"Bearer {credentials.api_key}"}
+ else:
+ auth_headers = await authorization_headers(self.auth)
+ async with self.http_client.stream(
+ "GET",
+ "https://api.stripe.com" + path,
+ params=params,
+ headers={**headers, **auth_headers},
+ timeout=20.0,
+ follow_redirects=False,
+ ) as response:
+ if not response.is_success:
+ # Do not read provider error bodies, which may echo private fields.
+ safe_response = httpx.Response(
+ response.status_code,
+ headers={k: v for k, v in response.headers.items() if k == "retry-after"},
+ request=httpx.Request("GET", "https://api.stripe.com/v1/"),
+ )
+ raise_for_status(
+ safe_response,
+ source_short_name="stripe",
+ token_provider_kind=self.auth.provider_kind,
+ )
+ content = bytearray()
+ async with aclosing(response.aiter_bytes(chunk_size=64 * 1024)) as chunks:
+ async for chunk in chunks:
+ if len(content) + len(chunk) > self.MAX_CAPTURE_RESPONSE_BYTES:
+ raise SourceError(
+ "Stripe capture response exceeds the 16 MiB decoded-body limit; "
+ "the page was not committed",
+ source_short_name="stripe",
+ )
+ content.extend(chunk)
+ try:
+ return json.loads(content)
+ except (ValueError, UnicodeError):
+ raise SourceError(
+ "Stripe returned invalid JSON", source_short_name="stripe"
+ ) from None
+
@retry(
stop=stop_after_attempt(5),
retry=retry_if_rate_limit_or_timeout,
@@ -551,4 +631,7 @@ async def generate_entities( # noqa: C901
async def validate(self) -> None:
"""Verify Stripe API key by pinging a lightweight endpoint (/v1/balance)."""
+ if self._capture_page_source is not None:
+ await self._capture_page_source.prepare_cycle(None)
+ return
await self._get("https://api.stripe.com/v1/balance")
diff --git a/backend/airweave/platform/sources/stripe_capture.py b/backend/airweave/platform/sources/stripe_capture.py
new file mode 100644
index 0000000000..85f275dc11
--- /dev/null
+++ b/backend/airweave/platform/sources/stripe_capture.py
@@ -0,0 +1,276 @@
+"""Account-bound Stripe v1 originals; full discovery, never absence reconciliation.
+
+The caller owns credentials/HTTP errors and must apply every supplied header. No
+activation is implied by constructing this adapter. Nested collections and linked
+files are retained as returned, not claimed to be fully acquired.
+"""
+
+import hashlib
+from datetime import datetime, timezone
+from typing import Annotated, Literal, Protocol
+
+from pydantic import (
+ BaseModel,
+ ConfigDict,
+ Field,
+ JsonValue,
+ StrictBool,
+ TypeAdapter,
+ ValidationError,
+)
+
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import CapturePage, CapturePlan
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.configs.config import StripeCaptureConfig
+
+NativeID = Annotated[str, Field(pattern=r"^[A-Za-z0-9_]{1,255}$")]
+_JSON = TypeAdapter(dict[str, JsonValue])
+_ENDPOINTS = {
+ "balance": "balance",
+ "balance_transaction": "balance_transactions",
+ "charge": "charges",
+ "customer": "customers",
+ "event": "events",
+ "invoice": "invoices",
+ "payment_intent": "payment_intents",
+ "payment_method": "payment_methods",
+ "payout": "payouts",
+ "refund": "refunds",
+ "subscription": "subscriptions",
+}
+
+
+class StripeCaptureError(ValueError):
+ """Sanitized boundary failure; the current page must not commit."""
+
+
+class StripeRead(Protocol):
+ """Authenticated transport; fixed relative v1 paths and no provider URL following."""
+
+ async def __call__(
+ self, path: str, *, params: dict[str, str | int], headers: dict[str, str]
+ ) -> dict[str, JsonValue]:
+ """Raise existing auth/rate/transport errors; return only successful native JSON."""
+ ...
+
+
+class _Account(BaseModel):
+ id: NativeID
+ object: Literal["account"]
+
+
+class _Balance(BaseModel):
+ object: Literal["balance"]
+ livemode: StrictBool
+
+
+class _Object(BaseModel):
+ id: NativeID
+ object: str
+ livemode: StrictBool | None = None
+ created: Annotated[int, Field(strict=True, ge=0, le=253402300799)] | None = None
+
+
+class _Page(BaseModel):
+ object: Literal["list"]
+ data: list[dict[str, JsonValue]] = Field(max_length=100)
+ has_more: StrictBool
+
+
+class _Continuation(BaseModel):
+ model_config = ConfigDict(extra="forbid", frozen=True)
+ fingerprint: str
+ record_type: str
+ starting_after: NativeID | None = None
+ done: StrictBool = False
+
+
+class StripeCapture:
+ """Eleven independent account-owned roots, preserving legacy resource breadth.
+
+ Payment methods use the unfiltered list (custom types are excluded by Stripe).
+ Subscriptions request all statuses; test-clock inventories and full event history
+ remain unqualified. Older API versions rejecting unfiltered lists fail explicitly;
+ no card-only fallback silently narrows capture.
+ """
+
+ canonical_record_types = tuple(_ENDPOINTS)
+ canonical_container_parents: dict[str, str] = {}
+
+ def __init__(self, read: StripeRead, config: StripeCaptureConfig) -> None:
+ """Use create to attest access before exposing capture configuration."""
+ self._read = read
+ self.config = config
+ self._verified = False
+ self.fingerprint = hashlib.sha256(
+ ("stripe-full-v2:" + config.model_dump_json()).encode()
+ ).hexdigest()
+
+ @classmethod
+ async def create(cls, read: StripeRead, config: StripeCaptureConfig) -> "StripeCapture":
+ """Attest the effective account and mode through the capture request context."""
+ instance = cls(read, config)
+ await instance.verify_principal()
+ return instance
+
+ async def _get(self, path: str, params: dict[str, str | int]) -> dict[str, JsonValue]:
+ headers = {"Stripe-Version": self.config.api_version}
+ if self.config.connected_account_id is not None:
+ headers["Stripe-Account"] = self.config.connected_account_id
+ raw = await self._read(path, params=params, headers=headers)
+ try:
+ return _JSON.validate_python(raw)
+ except ValidationError:
+ raise StripeCaptureError("Stripe returned invalid native JSON") from None
+
+ async def verify_principal(self) -> None:
+ """Reset prior attestation before any validation/reconnect attempt."""
+ self._verified = False
+ try:
+ account = _Account.model_validate(await self._get("/v1/account", {}))
+ if account.id != self.config.expected_account_id or (
+ self.config.connected_account_id is not None
+ and account.id != self.config.connected_account_id
+ ):
+ raise StripeCaptureError("Stripe account does not match trusted binding")
+ balance = _Balance.model_validate(await self._get("/v1/balance", {}))
+ if balance.livemode != self.config.livemode:
+ raise StripeCaptureError("Stripe mode does not match trusted binding")
+ except ValidationError:
+ raise StripeCaptureError("Stripe identity response is invalid") from None
+ self._verified = True
+
+ def _require_principal(self) -> None:
+ if not self._verified:
+ raise StripeCaptureError("Stripe principal is not attested")
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Completed list traversal never proves deletion or exhaustive provider coverage."""
+ self._require_principal()
+ return CycleConfiguration.from_source(
+ fingerprint=self.fingerprint,
+ record_types=self.canonical_record_types,
+ container_parents={},
+ completion_policies={kind: "discovery_only" for kind in self.canonical_record_types},
+ )
+
+ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan:
+ """Re-attest each full cycle; event capture is not an incremental checkpoint."""
+ await self.verify_principal()
+ return CapturePlan()
+
+ def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation:
+ """The engine binds cycle configuration; per-kind cursor binding occurs on first page."""
+ if cycle.configuration != self.capture_cycle_configuration:
+ raise StripeCaptureError("Stripe capture configuration changed")
+ return ScanContinuation()
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Relationships to other Stripe objects are references, not visibility parents."""
+ raise StripeCaptureError("Stripe capture has no child scopes")
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Aging events, list omission and generic HTTP errors never establish deletion."""
+ raise StripeCaptureError("Stripe discovery cannot confirm absence")
+
+ def _record(self, raw: dict[str, JsonValue], kind: str) -> CaptureRecord:
+ if kind == "balance":
+ balance = _Balance.model_validate(raw)
+ native_id, mode, created = "balance", balance.livemode, None
+ else:
+ obj = _Object.model_validate(raw)
+ if "livemode" in raw and obj.livemode is None:
+ raise StripeCaptureError("Stripe object mode is invalid")
+ # Stripe documents object as the kind discriminator; ID examples are
+ # not an exhaustive prefix contract (including historical IDs).
+ if obj.object != kind:
+ raise StripeCaptureError("Stripe object kind does not match scope")
+ native_id, mode = obj.id, obj.livemode
+ created = (
+ datetime.fromtimestamp(obj.created, timezone.utc)
+ if obj.created is not None
+ else None
+ )
+ if mode is not None and mode != self.config.livemode:
+ raise StripeCaptureError("Stripe object mode does not match source")
+ return CaptureRecord(
+ identity=RecordIdentity(record_type=kind, native_id=native_id),
+ payload=raw,
+ completeness="partial",
+ source_created_at=created,
+ observed_at=datetime.now(timezone.utc),
+ )
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """One list response and its exact cursor are committed by the shared engine."""
+ self._require_principal()
+ kind = scope.record_type
+ if (
+ kind not in _ENDPOINTS
+ or scope.container_id is not None
+ or scope.parent is not None
+ or parent is not None
+ ):
+ raise StripeCaptureError("Unsupported Stripe capture scope")
+ try:
+ progress = (
+ _Continuation.model_validate(continuation.value)
+ if continuation.value
+ else _Continuation(fingerprint=self.fingerprint, record_type=kind)
+ )
+ if (
+ progress.fingerprint != self.fingerprint
+ or progress.record_type != kind
+ or progress.done
+ ):
+ raise StripeCaptureError("Stripe continuation binding is invalid")
+ params: dict[str, str | int] = {} if kind == "balance" else {"limit": 100}
+ if progress.starting_after is not None:
+ if kind == "balance":
+ raise StripeCaptureError("Stripe balance cannot be paginated")
+ params["starting_after"] = progress.starting_after
+ if kind == "subscription":
+ params["status"] = "all"
+ raw = await self._get("/v1/" + _ENDPOINTS[kind], params)
+ if kind == "balance":
+ objects, has_more = [raw], False
+ else:
+ page = _Page.model_validate(raw)
+ objects, has_more = page.data, page.has_more
+ records = tuple(self._record(item, kind) for item in objects)
+ ids = [record.identity.native_id for record in records]
+ if (
+ len(set(ids)) != len(ids)
+ or progress.starting_after in ids
+ or (has_more and not ids)
+ ):
+ raise StripeCaptureError("Stripe pagination did not make progress")
+ following = _Continuation(
+ fingerprint=self.fingerprint,
+ record_type=kind,
+ starting_after=ids[-1] if has_more else None,
+ done=not has_more,
+ )
+ return CapturePage(
+ records=records,
+ continuation=ScanContinuation(value=following.model_dump(mode="json")),
+ final=not has_more,
+ )
+ except (ValidationError, OverflowError, OSError):
+ raise StripeCaptureError("Stripe capture response or continuation is invalid") from None
diff --git a/backend/airweave/platform/sources/teams.py b/backend/airweave/platform/sources/teams.py
index 4ec2cf7074..bb23aa3342 100644
--- a/backend/airweave/platform/sources/teams.py
+++ b/backend/airweave/platform/sources/teams.py
@@ -39,12 +39,12 @@
_parse_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/tests/test_calendar_identity.py b/backend/airweave/platform/sources/tests/test_calendar_identity.py
new file mode 100644
index 0000000000..8e597a9b83
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_calendar_identity.py
@@ -0,0 +1,67 @@
+"""Native primary resource attestation gates capture independently of lifecycle wiring."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import GoogleCalendarConfig
+from airweave.platform.sources.google_calendar import GoogleCalendarSource
+
+
+@pytest.mark.asyncio
+async def test_principal_change_invalidates_attestation_and_capture_identity():
+ native_id = "opaque-primary-id"
+
+ def profile(request):
+ assert request.url.path.endswith("/calendars/primary")
+ return httpx.Response(200, json={"id": native_id})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(profile)) as client:
+ source = await GoogleCalendarSource.create(
+ auth=StaticTokenProvider("same-broker-account"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GoogleCalendarConfig(expected_primary_calendar_id=native_id),
+ )
+ fingerprint = source.capture_cycle_configuration.fingerprint
+ native_id = "different-primary"
+ with pytest.raises(ValueError, match="does not match"):
+ await source.validate()
+ with pytest.raises(ValueError, match="attested primary"):
+ await source.capture_page(
+ CompletedScope(record_type="calendar"),
+ ScanContinuation(value={}),
+ files=MagicMock(),
+ )
+ source.calendar_config = GoogleCalendarConfig(expected_primary_calendar_id=native_id)
+ await source.validate()
+ assert source.capture_cycle_configuration.fingerprint != fingerprint
+ native_id = "private invalid value"
+ with pytest.raises(ValueError, match="invalid primary identity") as error:
+ await source.validate()
+ assert native_id not in str(error.value)
+
+
+@pytest.mark.asyncio
+async def test_unbound_legacy_source_cannot_start_or_resume_owned_capture():
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(lambda _: pytest.fail("Unexpected native call"))
+ ) as client:
+ source = await GoogleCalendarSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GoogleCalendarConfig(),
+ )
+ with pytest.raises(ValueError, match="attested primary"):
+ await source.prepare_cycle(None)
+ with pytest.raises(ValueError, match="attested primary"):
+ await source.capture_page(
+ CompletedScope(record_type="calendar"),
+ ScanContinuation(value={}),
+ files=MagicMock(),
+ )
diff --git a/backend/airweave/platform/sources/tests/test_drive_identity.py b/backend/airweave/platform/sources/tests/test_drive_identity.py
new file mode 100644
index 0000000000..2e76066990
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_drive_identity.py
@@ -0,0 +1,95 @@
+"""Native Drive principal attestation precedes every owned capture entrypoint."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import GoogleDriveConfig
+from airweave.platform.sources.google_drive import GoogleDriveSource
+
+
+async def connector(client, expected="principal-a"):
+ return await GoogleDriveSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GoogleDriveConfig(expected_permission_id=expected),
+ )
+
+
+@pytest.mark.parametrize(
+ "user",
+ [
+ {"permissionId": "different", "me": True},
+ {"permissionId": "principal-a", "me": False},
+ {"permissionId": "principal-a", "me": "true"},
+ {"permissionId": 123, "me": True},
+ {"permissionId": "", "me": True},
+ {"permissionId": "principal-a"},
+ None,
+ ],
+)
+async def test_invalid_principal_never_reaches_capture(user):
+ calls = []
+
+ def respond(request):
+ calls.append(request.url.path)
+ assert request.url.params["fields"] == "user(permissionId,me)"
+ return httpx.Response(200, json={"user": user})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+ with pytest.raises(ValueError) as error:
+ await connector(client)
+ assert calls == ["/drive/v3/about"]
+ assert "principal-a" not in str(error.value) and "different" not in str(error.value)
+
+
+async def test_revalidation_failure_revokes_direct_capture_and_resume():
+ replies = ["principal-a", "different"]
+
+ def respond(request):
+ assert request.url.path == "/drive/v3/about"
+ return httpx.Response(200, json={"user": {"permissionId": replies.pop(0), "me": True}})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+ source = await connector(client)
+ original = source.capture_cycle_configuration
+ with pytest.raises(ValueError, match="trusted binding"):
+ await source.validate()
+ with pytest.raises(ValueError, match="attested"):
+ _ = source.capture_cycle_configuration
+ with pytest.raises(ValueError, match="attested"):
+ await source.prepare_cycle(None)
+ with pytest.raises(ValueError, match="attested"):
+ source.initial_continuation(MagicMock(configuration=original))
+ with pytest.raises(ValueError, match="attested"):
+ await source.capture_page(None, None, files=MagicMock())
+ with pytest.raises(ValueError, match="attested"):
+ await source.refresh_known(None, files=MagicMock())
+ assert not replies
+
+
+async def test_legacy_creation_is_allowed_but_canonical_capture_requires_binding():
+ def respond(request):
+ raise AssertionError("Legacy creation should not require a new permission binding")
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+ source = await connector(client, None)
+ with pytest.raises(ValueError, match="attested"):
+ await source.prepare_cycle(None)
+
+
+async def test_fingerprint_and_resume_are_principal_bound():
+ principals = iter(["principal-a", "principal-b"])
+
+ def respond(request):
+ return httpx.Response(200, json={"user": {"permissionId": next(principals), "me": True}})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+ first = await connector(client)
+ second = await connector(client, "principal-b")
+ assert first.capture_cycle_configuration != second.capture_cycle_configuration
+ with pytest.raises(ValueError, match="trusted capture configuration"):
+ second.initial_continuation(MagicMock(configuration=first.capture_cycle_configuration))
diff --git a/backend/airweave/platform/sources/tests/test_gmail_capture.py b/backend/airweave/platform/sources/tests/test_gmail_capture.py
new file mode 100644
index 0000000000..ac72942306
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py
@@ -0,0 +1,522 @@
+"""Gmail native pages and MIME retention; SQL recovery is tested separately."""
+
+from uuid import uuid4
+
+import httpx
+import pytest
+
+from airweave.domains.entities.canonical.cycle_models import (
+ CaptureCycle,
+ CycleConfiguration,
+ CycleVersion,
+)
+from airweave.domains.entities.canonical.page_source import (
+ InvalidCaptureCheckpoint,
+ InvalidScanContinuation,
+)
+from airweave.platform.sources.gmail_capture import BASE, GmailCapture
+from airweave.platform.sources.gmail_pages import GmailPages
+
+
+def message(mid, **extra):
+ return {"id": mid, "threadId": "t", "internalDate": "1000", "payload": {}, **extra}
+
+
+def missing():
+ response = httpx.Response(404, request=httpx.Request("GET", BASE))
+ return httpx.HTTPStatusError("Gone", request=response.request, response=response)
+
+
+class Provider:
+ def __init__(self, steps):
+ self.steps = list(steps)
+ self.calls = []
+
+ async def get(self, url, params=None):
+ self.calls.append((url, params))
+ expected, response = self.steps.pop(0)
+ assert url == BASE + expected
+ if isinstance(response, Exception):
+ raise response
+ return response
+
+
+async def initial(pages):
+ config = CycleConfiguration(fingerprint="a" * 64, parents={"message": (None,)})
+ plan = await pages.prepare(None, config)
+ cycle = CaptureCycle(
+ version=CycleVersion(cycle_id=uuid4(), revision=1),
+ configuration=config,
+ mode=plan.mode,
+ starting_checkpoint=plan.starting_checkpoint,
+ )
+ return pages.initial(cycle)
+
+
+@pytest.mark.asyncio
+async def test_baseline_interleaves_history_before_next_list_page():
+ provider = Provider(
+ [
+ ("/profile", {"historyId": "10"}),
+ ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "p2"}),
+ ("/messages/a", message("a")),
+ (
+ "/history",
+ {"historyId": "11", "history": [{"labelsRemoved": [{"message": {"id": "a"}}]}]},
+ ),
+ ("/messages/a", message("a", labelIds=[])),
+ ("/messages", {"messages": [{"id": "b"}]}),
+ ("/messages/b", message("b")),
+ (
+ "/history",
+ {"historyId": "12", "history": [{"messagesAdded": [{"message": {"id": "c"}}]}]},
+ ),
+ ("/messages/c", message("c")),
+ ]
+ )
+ pages = GmailPages(GmailCapture(provider.get, None))
+ continuation = await initial(pages)
+ results = []
+ for _ in range(4):
+ page = await pages.page(continuation)
+ results.append(page)
+ continuation = page.continuation
+ assert [r.identity.native_id for p in results for r in p.records] == ["a", "a", "b", "c"]
+ assert [p.final for p in results] == [False, False, False, True]
+ assert all(p.provider_checkpoint is None for p in results[:-1])
+ assert results[-1].provider_checkpoint.value == {"history_id": "12"}
+ assert provider.calls[3][1]["startHistoryId"] == "10"
+ assert provider.calls[5][1]["pageToken"] == "p2"
+ assert provider.calls[7][1]["startHistoryId"] == "11"
+ assert provider.calls[1][1]["includeSpamTrash"] == "true"
+ assert "q" not in provider.calls[1][1]
+
+
+@pytest.mark.asyncio
+async def test_filtered_query_is_list_only_without_provider_checkpoint():
+ provider = Provider([("/messages", {"messages": [{"id": "a"}]}), ("/messages/a", message("a"))])
+ pages = GmailPages(GmailCapture(provider.get, "in:inbox"))
+ page = await pages.page(await initial(pages))
+ assert page.final and page.provider_checkpoint is None
+ assert provider.calls[0][1]["q"] == "in:inbox"
+ assert not provider.steps
+
+
+@pytest.mark.asyncio
+async def test_expired_history_requires_whole_cycle_recovery():
+ provider = Provider(
+ [("/profile", {"historyId": "10"}), ("/messages", {}), ("/history", missing())]
+ )
+ pages = GmailPages(GmailCapture(provider.get, None))
+ first = await pages.page(await initial(pages))
+ with pytest.raises(InvalidCaptureCheckpoint):
+ await pages.page(first.continuation)
+ assert first.provider_checkpoint is None
+
+
+@pytest.mark.asyncio
+async def test_history_replay_hydrates_bounded_offsets_and_rejects_changed_page():
+ raw = {"historyId": "11", "history": [{"messages": [{"id": str(i)} for i in range(51)]}]}
+ steps = [("/profile", {"historyId": "10"}), ("/messages", {}), ("/history", raw)]
+ steps.extend((f"/messages/{i}", message(str(i))) for i in range(50))
+ steps.append(("/history", {**raw, "historyId": "12"}))
+ provider = Provider(steps)
+ pages = GmailPages(GmailCapture(provider.get, None))
+ listing = await pages.page(await initial(pages))
+ half = await pages.page(listing.continuation)
+ assert len(half.records) == 50 and not half.final and half.provider_checkpoint is None
+ assert half.continuation.value["history_offset"] == 50
+ with pytest.raises(InvalidScanContinuation):
+ await pages.page(half.continuation)
+ assert provider.calls[-1][0].endswith("/history")
+
+
+@pytest.mark.asyncio
+async def test_history_pagination_keeps_start_boundary_until_terminal():
+ provider = Provider(
+ [
+ ("/profile", {"historyId": "10"}),
+ ("/messages", {}),
+ ("/history", {"historyId": "999", "nextPageToken": "next"}),
+ ("/history", {"historyId": "1000"}),
+ ]
+ )
+ pages = GmailPages(GmailCapture(provider.get, None))
+ listing = await pages.page(await initial(pages))
+ intermediate = await pages.page(listing.continuation)
+ assert intermediate.continuation.value["history_boundary"] == "10"
+ assert intermediate.provider_checkpoint is None
+ terminal = await pages.page(intermediate.continuation)
+ assert provider.calls[-1][1]["startHistoryId"] == "10"
+ assert terminal.provider_checkpoint.value == {"history_id": "1000"}
+
+
+@pytest.mark.asyncio
+async def test_repeated_listing_cursor_fails_without_a_completed_page():
+ provider = Provider(
+ [("/messages", {"nextPageToken": "same"}), ("/messages", {"nextPageToken": "same"})]
+ )
+ pages = GmailPages(GmailCapture(provider.get, "in:inbox"))
+ first = await pages.page(await initial(pages))
+ with pytest.raises(ValueError, match="repeated"):
+ await pages.page(first.continuation)
+ assert not first.final
+
+
+@pytest.mark.asyncio
+async def test_external_mime_bytes_are_stored_before_complete_record_without_payload_mutation():
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.storage.file_service import FileService
+
+ storage = AsyncMock()
+ files = FileService(uuid4(), storage, sync_id=uuid4())
+ raw = message(
+ "a",
+ payload={
+ "parts": [
+ {
+ "mimeType": "text/plain",
+ "filename": "नोट्स.txt",
+ "body": {
+ "attachmentId": "body-id",
+ "size": 3,
+ },
+ }
+ ]
+ },
+ )
+ provider = Provider([("/messages/a", raw)])
+ attachment_get = AsyncMock(return_value={"data": "YWJj", "size": 3})
+ record = await GmailCapture(
+ provider.get, None, files=files, attachment_get=attachment_get
+ ).message("a")
+ assert record.completeness == "complete"
+ assert record.payload == raw
+ assert record.blobs[0].source_path == "/payload/parts/0/body"
+ assert record.blobs[0].filename == "नोट्स.txt"
+ assert record.blobs[0].size_bytes == 3
+ storage.write_file.assert_awaited_once_with(record.blobs[0].key, b"abc")
+
+
+@pytest.mark.asyncio
+async def test_blob_storage_failure_propagates_and_prevents_record_checkpoint():
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.storage.file_service import FileService
+
+ storage = AsyncMock()
+ storage.write_file.side_effect = OSError("storage unavailable")
+ files = FileService(uuid4(), storage, sync_id=uuid4())
+ provider = Provider(
+ [
+ ("/history", {"historyId": "21", "history": [{"messages": [{"id": "a"}]}]}),
+ ("/messages/a", message("a", payload={"body": {"attachmentId": "body-id", "size": 3}})),
+ ]
+ )
+ attachment_get = AsyncMock(return_value={"data": "YWJj", "size": 3})
+ capture = GmailCapture(provider.get, None, files=files, attachment_get=attachment_get)
+ page = await capture.history_page("20")
+ with pytest.raises(OSError):
+ await capture.hydrate_history_page(page)
+
+
+@pytest.mark.asyncio
+async def test_oversize_declared_mime_is_partial_without_fetching_or_writing():
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.storage.file_service import FileService
+
+ storage = AsyncMock()
+ files = FileService(uuid4(), storage, sync_id=uuid4())
+ files.MAX_FILE_SIZE_BYTES = 2
+ provider = Provider(
+ [
+ (
+ "/messages/a",
+ message(
+ "a",
+ payload={
+ "body": {
+ "attachmentId": "body-id",
+ "size": 3,
+ }
+ },
+ ),
+ )
+ ]
+ )
+ attachment_get = AsyncMock()
+ record = await GmailCapture(
+ provider.get, None, files=files, attachment_get=attachment_get
+ ).message("a")
+ assert record.completeness == "partial"
+ attachment_get.assert_not_called()
+ storage.write_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_mime_size_mismatch_does_not_store_unverified_bytes():
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.storage.file_service import FileService
+
+ storage = AsyncMock()
+ files = FileService(uuid4(), storage, sync_id=uuid4())
+ provider = Provider(
+ [
+ (
+ "/messages/a",
+ message(
+ "a",
+ payload={
+ "body": {
+ "attachmentId": "body-id",
+ "size": 4,
+ }
+ },
+ ),
+ )
+ ]
+ )
+ attachment_get = AsyncMock(return_value={"data": "YWJj", "size": 3})
+ with pytest.raises(ValueError, match="size"):
+ await GmailCapture(provider.get, None, files=files, attachment_get=attachment_get).message(
+ "a"
+ )
+ storage.write_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_mime_json_stream_size_limit_closes_response():
+ import logging
+
+ from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider
+ from airweave.domains.storage.exceptions import FileSkippedException
+ from airweave.platform.configs.config import GmailConfig
+ from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+ from airweave.platform.sources.gmail import GmailSource
+
+ response = httpx.Response(200, json={"data": "YWJj", "size": 3})
+ async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: response)) as client:
+ source = await GmailSource.create(
+ auth=ManagedAuthProvider(
+ api_key="not-used",
+ connected_account_id="test",
+ allowed_hosts=frozenset({"gmail.googleapis.com"}),
+ ),
+ logger=logging.getLogger("test"),
+ http_client=AirweaveHttpClient(client, uuid4(), "gmail", feature_flag_enabled=False),
+ config=GmailConfig(),
+ )
+ with pytest.raises(FileSkippedException):
+ await source._get_mime_json(BASE + "/messages/a/attachments/b", max_bytes=5)
+ assert response.is_closed
+
+
+@pytest.mark.asyncio
+async def test_large_originals_split_batch_without_relisting_or_truncation():
+ raw = message("large", payload={"body": {"data": "x" * (9 * 1024 * 1024)}})
+ provider = Provider(
+ [
+ ("/messages", {"messages": [{"id": "large"}, {"id": "small"}]}),
+ ("/messages/large", raw),
+ ("/messages/small", message("small")),
+ ]
+ )
+ pages = GmailPages(GmailCapture(provider.get, "in:inbox"))
+ first = await pages.page(await initial(pages))
+ assert not first.final and first.records[0].payload == raw
+ assert first.continuation.value["pending_ids"] == ["small"]
+ # A new adapter instance models loss of all in-process state.
+ final = await GmailPages(GmailCapture(provider.get, "in:inbox")).page(first.continuation)
+ assert final.final and final.records[0].identity.native_id == "small"
+ assert len([url for url, _ in provider.calls if url.endswith("/messages")]) == 1
+
+
+def test_combined_cursor_rings_fit_durable_continuation_and_recent_loops_fail():
+ import hashlib
+
+ from airweave.platform.sources.gmail_pages import _next_token, _Progress
+
+ values = tuple(hashlib.sha256(str(i).encode()).hexdigest() for i in range(128))
+ state = _Progress(
+ mode="full",
+ phase="history",
+ history_boundary="h" * 1024,
+ list_token="l" * 8192,
+ history_token="h" * 8192,
+ list_tokens=values,
+ history_tokens=values,
+ )
+ assert len(state.continuation().model_dump_json().encode()) < 65536
+ rotated = values
+ for i in range(128, 1000):
+ rotated = _next_token(str(i), rotated)
+ assert len(rotated) == 128
+ with pytest.raises(ValueError, match="repeated"):
+ _next_token("999", rotated)
+
+
+def test_only_explicit_page_token_parameter_errors_authorize_reset():
+ from airweave.platform.sources.gmail_pages import invalid_page_token
+
+ assert invalid_page_token(
+ {
+ "error": {
+ "errors": [
+ {"reason": "badRequest", "location": "pageToken", "locationType": "parameter"}
+ ]
+ }
+ }
+ )
+ assert not invalid_page_token(
+ {
+ "error": {
+ "errors": [{"reason": "badRequest", "location": "q", "locationType": "parameter"}]
+ }
+ }
+ )
+ assert not invalid_page_token({"error": {"message": "Invalid request"}})
+
+
+@pytest.mark.asyncio
+async def test_bounded_native_reads_refresh_once_and_identify_only_explicit_bad_token():
+ from unittest.mock import MagicMock
+
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.platform.configs.config import GmailConfig
+ from airweave.platform.sources.gmail import GmailSource
+
+ class RefreshingToken(StaticTokenProvider):
+ refreshes = 0
+
+ @property
+ def supports_refresh(self):
+ return True
+
+ async def force_refresh(self):
+ self.refreshes += 1
+ self._token = "refreshed"
+ return self._token
+
+ auth = RefreshingToken("initial")
+ requests = []
+ responses = [
+ httpx.Response(401, json={"error": {"message": "expired"}}),
+ httpx.Response(200, json={"historyId": "current"}),
+ httpx.Response(
+ 400,
+ json={
+ "error": {
+ "errors": [
+ {
+ "reason": "badRequest",
+ "location": "pageToken",
+ "locationType": "parameter",
+ }
+ ]
+ }
+ },
+ ),
+ ]
+
+ async def handler(request):
+ requests.append(request)
+ return responses.pop(0)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+ source = await GmailSource.create(
+ auth=auth, logger=MagicMock(), http_client=client, config=GmailConfig()
+ )
+ assert await source._get_capture_json(BASE + "/profile") == {"historyId": "current"}
+ with pytest.raises(InvalidScanContinuation):
+ await source._get_capture_json(BASE + "/messages", params={"pageToken": "expired"})
+ assert auth.refreshes == 1
+ assert requests[0].headers["authorization"] == "Bearer initial"
+ assert requests[1].headers["authorization"] == "Bearer refreshed"
+
+
+@pytest.mark.asyncio
+async def test_bounded_transport_stops_before_reading_entire_oversized_body():
+ from unittest.mock import MagicMock
+
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.domains.storage.exceptions import FileSkippedException
+ from airweave.platform.configs.config import GmailConfig
+ from airweave.platform.sources.gmail import GmailSource
+
+ class Chunks(httpx.AsyncByteStream):
+ observed = 0
+
+ closed = False
+
+ def __aiter__(self):
+ self.iterator = self.chunks()
+ return self.iterator
+
+ async def chunks(self):
+ for _ in range(10):
+ self.observed += 1
+ yield b"x" * 5
+
+ async def aclose(self):
+ self.closed = True
+ await self.iterator.aclose()
+
+ chunks = Chunks()
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(lambda request: httpx.Response(200, stream=chunks))
+ ) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ with pytest.raises(FileSkippedException):
+ await source._get_mime_json(BASE + "/messages/a", max_bytes=8)
+ assert chunks.observed == 2 and chunks.closed
+
+
+@pytest.mark.asyncio
+async def test_managed_proxy_decoded_headers_and_identity_request_match_bounded_reader():
+ import json
+ from unittest.mock import MagicMock
+
+ from airweave.domains.sources.token_providers.static import StaticTokenProvider
+ from airweave.platform.configs.config import GmailConfig
+ from airweave.platform.http_client.composio_transport import ComposioTransport
+ from airweave.platform.sources.gmail import GmailSource
+
+ async def proxy(request):
+ payload = json.loads(request.content)
+ assert any(
+ item["name"].lower() == "accept-encoding" and item["value"] == "identity"
+ for item in payload["parameters"]
+ )
+ return httpx.Response(
+ 200,
+ json={
+ "status": 200,
+ "data": {"historyId": "synthetic"},
+ "headers": {"content-encoding": "gzip"},
+ },
+ )
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as upstream:
+ async with httpx.AsyncClient(
+ transport=ComposioTransport(
+ client=upstream,
+ api_key="fixture",
+ connected_account_id="fixture",
+ allowed_hosts={"gmail.googleapis.com"},
+ )
+ ) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ assert await source._get_capture_json(BASE + "/profile") == {"historyId": "synthetic"}
diff --git a/backend/airweave/platform/sources/tests/test_gmail_errors.py b/backend/airweave/platform/sources/tests/test_gmail_errors.py
new file mode 100644
index 0000000000..e65709c193
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_gmail_errors.py
@@ -0,0 +1,181 @@
+"""Native quota responses through both Gmail request boundaries."""
+
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+
+from airweave.domains.sources.exceptions import SourceEntityForbiddenError
+from airweave.domains.sources.exceptions.classifier import classify_error
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import GmailConfig
+from airweave.platform.sources.gmail import GmailSource
+from airweave.platform.sources.gmail_errors import GmailThrottleError
+
+
+def envelope(reason, *, domain="usageLimits", code=403):
+ return {"error": {"code": code, "errors": [{"domain": domain, "reason": reason}]}}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "method,reason,delay",
+ [
+ ("_get", "rateLimitExceeded", None),
+ ("_get_capture_json", "userRateLimitExceeded", "3"),
+ ("_get_capture_json", "rateLimitExceeded", "180"),
+ ("_get_capture_json", "rateLimitExceeded", "invalid"),
+ ("_get_capture_json", "rateLimitExceeded", "NaN"),
+ ],
+)
+async def test_native_quota_retry_is_bounded_and_preserves_evidence(method, reason, delay):
+ calls = 0
+
+ def handler(request):
+ nonlocal calls
+ calls += 1
+ return httpx.Response(
+ 403, json=envelope(reason), headers={} if delay is None else {"Retry-After": delay}
+ )
+
+ sleep = AsyncMock()
+ async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("fixture"),
+ logger=Mock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ call = getattr(GmailSource, method).retry_with(sleep=sleep)
+ with pytest.raises(GmailThrottleError) as caught:
+ await call(source, "https://gmail.googleapis.com/gmail/v1/users/me/profile")
+ assert caught.value.status_code == 403
+ assert caught.value.retry_after == (float(delay) if delay in {"3", "180"} else None)
+ assert calls == (1 if delay == "180" else 5)
+ assert sleep.await_count == (0 if delay == "180" else 4)
+ if delay == "3":
+ assert [call.args[0] for call in sleep.await_args_list] == [3] * 4
+ elif delay != "180":
+ assert [call.args[0] for call in sleep.await_args_list] == [2, 2, 4, 8]
+ assert classify_error(caught.value).category is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "body",
+ [
+ envelope("domainPolicy", domain="global"),
+ envelope("dailyLimitExceeded"),
+ envelope("unknown"),
+ envelope("rateLimitExceeded", domain="unknown"),
+ envelope("rateLimitExceeded", code=400),
+ {"error": {"code": 403, "errors": []}},
+ {"error": {"code": 403, "errors": [{"reason": "rateLimitExceeded"}]}},
+ {
+ "error": {
+ "code": 403,
+ "errors": [
+ {"domain": "usageLimits", "reason": "rateLimitExceeded"},
+ {"domain": "global", "reason": "domainPolicy"},
+ ],
+ }
+ },
+ "not json",
+ b"{broken",
+ ],
+)
+async def test_permission_unknown_and_malformed_errors_do_not_retry(body):
+ calls = 0
+
+ def handler(request):
+ nonlocal calls
+ calls += 1
+ return (
+ httpx.Response(403, content=body)
+ if isinstance(body, bytes)
+ else httpx.Response(403, json=body)
+ )
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("fixture"),
+ logger=Mock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ with pytest.raises(SourceEntityForbiddenError):
+ await source._get_capture_json("https://gmail.googleapis.com/gmail/v1/users/me/profile")
+ assert calls == 1
+
+
+def test_exhausted_throttle_uses_transient_validation_response():
+ from airweave.domains.sources.http_translation import http_exception_for_credential_validation
+
+ error = GmailThrottleError(None)
+ response = http_exception_for_credential_validation(error, source_short_name="gmail")
+ assert response.status_code == 502
+ assert not response.headers
+ assert "try again later" in response.detail
+
+
+@pytest.mark.parametrize("header", ["Wed, 01 Jan 2020 00:00:00 GMT", "-1", "inf"])
+def test_retry_timing_does_not_invent_provider_delay(header):
+ from airweave.platform.sources.gmail_errors import raise_gmail_throttle
+
+ with pytest.raises(GmailThrottleError) as caught:
+ raise_gmail_throttle(
+ httpx.Response(403, json=envelope("rateLimitExceeded"), headers={"Retry-After": header})
+ )
+ assert caught.value.retry_after == (None if header == "inf" else 0)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["_get", "_get_capture_json"])
+async def test_quota_then_success_returns_provider_result(method):
+ calls = 0
+
+ def handler(request):
+ nonlocal calls
+ calls += 1
+ return (
+ httpx.Response(403, json=envelope("userRateLimitExceeded"))
+ if calls == 1
+ else httpx.Response(200, json={"historyId": "verified"})
+ )
+
+ sleep = AsyncMock()
+ async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("fixture"),
+ logger=Mock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ call = getattr(GmailSource, method).retry_with(sleep=sleep)
+ assert await call(source, "https://gmail.googleapis.com/gmail/v1/users/me/profile") == {
+ "historyId": "verified"
+ }
+ assert calls == 2
+ sleep.assert_awaited_once_with(2)
+
+
+def test_future_http_date_retry_after_uses_actual_delay(monkeypatch):
+ from datetime import datetime, timezone
+
+ from airweave.platform.sources import gmail_errors
+
+ class Clock(datetime):
+ @classmethod
+ def now(cls, tz=None):
+ return datetime(2026, 9, 30, 20, 0, 0, tzinfo=timezone.utc)
+
+ monkeypatch.setattr(gmail_errors, "datetime", Clock)
+ with pytest.raises(GmailThrottleError) as caught:
+ gmail_errors.raise_gmail_throttle(
+ httpx.Response(
+ 403,
+ json=envelope("rateLimitExceeded"),
+ headers={"Retry-After": "Wed, 30 Sep 2026 20:01:00 GMT"},
+ )
+ )
+ assert caught.value.retry_after == 60
diff --git a/backend/airweave/platform/sources/tests/test_gmail_history_pages.py b/backend/airweave/platform/sources/tests/test_gmail_history_pages.py
new file mode 100644
index 0000000000..64619734cb
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_gmail_history_pages.py
@@ -0,0 +1,206 @@
+"""Bounded history replay uses current native message state, never event arrival state."""
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from airweave.domains.entities.canonical.page_source import InvalidScanContinuation
+from airweave.platform.sources.gmail_capture import BASE, GmailCapture, parse_history_page
+from airweave.platform.sources.http_helpers import raise_for_status
+from airweave.platform.sources.tests.test_gmail_capture import Provider, message
+
+
+def history(ids, **extra):
+ return {
+ "historyId": "opaque:boundary",
+ "history": [{"messages": [{"id": i} for i in ids]}],
+ **extra,
+ }
+
+
+def translated_error(status):
+ response = httpx.Response(status, request=httpx.Request("GET", BASE))
+ try:
+ raise_for_status(response, source_short_name="gmail", token_provider_kind="composio")
+ except Exception as error:
+ return error
+ raise AssertionError("Expected source exception")
+
+
+@pytest.mark.asyncio
+async def test_one_history_entry_can_hydrate_more_than_500_messages_in_bounded_batches():
+ ids = [f"m{i}" for i in range(1003)]
+ provider = Provider([(f"/messages/{mid}", message(mid)) for mid in ids])
+ capture = GmailCapture(provider.get, None)
+ page = parse_history_page(history(ids))
+ first = await capture.hydrate_history_page(page)
+ assert len(first.records) == 500 and not first.complete
+ second = await capture.hydrate_history_page(
+ page, offset=first.next_offset, expected_fingerprint=page.fingerprint
+ )
+ final = await capture.hydrate_history_page(
+ page, offset=second.next_offset, expected_fingerprint=page.fingerprint
+ )
+ assert len(second.records) == 500 and len(final.records) == 3
+ assert final.next_offset == 1003 and final.complete
+ assert [r.identity.native_id for r in (*first.records, *second.records, *final.records)] == ids
+ assert not provider.steps
+
+
+def test_all_native_arrays_ordered_dedup_and_opaque_boundary():
+ raw = {
+ "historyId": "opaque:end",
+ "history": [
+ {
+ "messages": [{"id": "a"}],
+ "messagesAdded": [{"message": {"id": "a"}}, {"message": {"id": "b"}}],
+ "messagesDeleted": [{"message": {"id": "c"}}],
+ "labelsAdded": [{"message": {"id": "d"}, "labelIds": ["L"]}],
+ "labelsRemoved": [{"message": {"id": "e"}}],
+ }
+ ],
+ }
+ page = parse_history_page(raw)
+ assert page.message_ids == ("a", "b", "c", "d", "e")
+ assert page.history_id == "opaque:end"
+ assert parse_history_page({**raw, "historyId": "new-current"}).fingerprint != page.fingerprint
+ intermediate = {**raw, "nextPageToken": "next"}
+ assert (
+ parse_history_page(intermediate).fingerprint
+ == parse_history_page({**intermediate, "historyId": "new-current"}).fingerprint
+ )
+ assert parse_history_page({**raw, "nextPageToken": "next"}).fingerprint != page.fingerprint
+
+
+@pytest.mark.asyncio
+async def test_changed_replayed_ids_reject_offset_before_any_reads():
+ provider = Provider([])
+ capture = GmailCapture(provider.get, None)
+ old = parse_history_page(history(["a", "b"]))
+ changed = parse_history_page(history(["b", "a"]))
+ with pytest.raises(InvalidScanContinuation):
+ await capture.hydrate_history_page(changed, offset=1, expected_fingerprint=old.fingerprint)
+ with pytest.raises(ValueError, match="fingerprint"):
+ await capture.hydrate_history_page(old, offset=1)
+ assert not provider.calls
+
+
+@pytest.mark.asyncio
+async def test_current_state_overrides_deleted_event_and_typed_404_tombstones():
+ provider = Provider(
+ [
+ ("/messages/alive", message("alive", labelIds=["TRASH"])),
+ ("/messages/gone", translated_error(404)),
+ ]
+ )
+ page = parse_history_page(
+ {
+ "historyId": "opaque",
+ "history": [
+ {"messagesDeleted": [{"message": {"id": "alive"}}, {"message": {"id": "gone"}}]}
+ ],
+ }
+ )
+ batch = await GmailCapture(provider.get, None).hydrate_history_page(page)
+ assert batch.records[0].kind == "upsert"
+ assert batch.records[0].payload["labelIds"] == ["TRASH"]
+ assert batch.records[1].kind == "delete"
+ assert batch.records[1].removal_reason == "provider_deleted"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status", [401, 403, 429, 500])
+async def test_non_missing_errors_propagate_from_message_and_history(status):
+ error = translated_error(status)
+ provider = Provider([("/messages/a", error), ("/history", error)])
+ capture = GmailCapture(provider.get, None)
+ with pytest.raises(type(error)):
+ await capture.hydrate_history_page(parse_history_page(history(["a"])))
+ with pytest.raises(type(error)):
+ await capture.history_page("opaque")
+
+
+@pytest.mark.asyncio
+async def test_page_fetch_preserves_boundary_and_token_without_advancing_cursor():
+ provider = Provider([("/history", history([], nextPageToken="next"))])
+ page = await GmailCapture(provider.get, None).history_page(
+ "opaque:start", "token", max_results=1
+ )
+ assert provider.calls[0][1] == {
+ "startHistoryId": "opaque:start",
+ "pageToken": "token",
+ "maxResults": 1,
+ }
+ assert page.next_page_token == "next" and page.message_ids == ()
+
+
+@pytest.mark.parametrize(
+ "raw",
+ [
+ {"historyId": 123},
+ {"historyId": ""},
+ {"historyId": "x", "history": [{"messages": [{"id": 7}]}]},
+ {"historyId": "x", "history": [{"messagesDeleted": [{}]}]},
+ ],
+)
+def test_malformed_provider_history_never_becomes_empty_success(raw):
+ with pytest.raises(ValidationError):
+ parse_history_page(raw)
+
+
+@pytest.mark.asyncio
+async def test_history_hydration_preserves_native_payload_and_verified_mime_bytes():
+ from unittest.mock import AsyncMock
+ from uuid import uuid4
+
+ from airweave.domains.storage.file_service import FileService
+
+ storage = AsyncMock()
+ files = FileService(uuid4(), storage, sync_id=uuid4())
+ raw = message(
+ "a",
+ payload={
+ "mimeType": "text/plain",
+ "body": {
+ "attachmentId": "body-id",
+ "size": 3,
+ },
+ },
+ )
+ provider = Provider([("/messages/a", raw)])
+ capture = GmailCapture(
+ provider.get,
+ None,
+ files=files,
+ attachment_get=AsyncMock(return_value={"data": "YWJj", "size": 3}),
+ )
+ batch = await capture.hydrate_history_page(parse_history_page(history(["a"])))
+ (record,) = batch.records
+ assert record.payload == raw and record.completeness == "complete"
+ assert record.blobs[0].source_path == "/payload/body"
+ storage.write_file.assert_awaited_once_with(record.blobs[0].key, b"abc")
+
+
+@pytest.mark.asyncio
+async def test_new_event_for_same_message_invalidates_partial_hydration():
+ initial = {
+ "historyId": "first",
+ "history": [
+ {"id": "event-one", "labelsAdded": [{"message": {"id": "a"}, "labelIds": ["INBOX"]}]}
+ ],
+ }
+ changed = {
+ "historyId": "second",
+ "history": [
+ *initial["history"],
+ {"id": "event-two", "labelsRemoved": [{"message": {"id": "a"}, "labelIds": ["INBOX"]}]},
+ ],
+ }
+ old, new = parse_history_page(initial), parse_history_page(changed)
+ assert old.message_ids == new.message_ids == ("a",)
+ provider = Provider([])
+ with pytest.raises(InvalidScanContinuation):
+ await GmailCapture(provider.get, None).hydrate_history_page(
+ new, offset=1, expected_fingerprint=old.fingerprint
+ )
+ assert not provider.calls
diff --git a/backend/airweave/platform/sources/tests/test_gmail_identity.py b/backend/airweave/platform/sources/tests/test_gmail_identity.py
new file mode 100644
index 0000000000..29eb131078
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_gmail_identity.py
@@ -0,0 +1,70 @@
+"""Mailbox expectations belong to the binding, never the observed response."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import GmailConfig
+from airweave.platform.sources.gmail import GmailSource
+
+
+@pytest.mark.asyncio
+async def test_attestation_and_changed_expectation():
+ requests = []
+ mailbox = "person@example.test"
+
+ def profile(request):
+ requests.append(request)
+ return httpx.Response(200, json={"emailAddress": mailbox})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(profile)) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("same-broker-account"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GmailConfig(expected_mailbox="person@example.test"),
+ )
+ original = source.capture_cycle_configuration.fingerprint
+ assert requests[0].url.params["fields"] == "emailAddress"
+ source.capture_config = GmailConfig(expected_mailbox="other@example.test")
+ with pytest.raises(ValueError, match="attested mailbox"):
+ _ = source.capture_cycle_configuration
+ with pytest.raises(ValueError, match="does not match"):
+ await source.validate()
+ source.capture_config = GmailConfig(expected_mailbox="person@example.test")
+ with pytest.raises(ValueError, match="attested mailbox"):
+ _ = source.capture_cycle_configuration
+ await source.validate()
+ assert source.capture_cycle_configuration.fingerprint == original
+ # A separately attested identity cannot reuse the old capture fingerprint.
+ mailbox = "other@example.test"
+ source.capture_config = GmailConfig(expected_mailbox=mailbox)
+ await source.validate()
+ assert source.capture_cycle_configuration.fingerprint != original
+ mailbox = "private malformed value"
+ with pytest.raises(ValueError, match="invalid mailbox identity") as error:
+ await source.validate()
+ assert mailbox not in str(error.value)
+
+
+@pytest.mark.asyncio
+async def test_direct_resume_without_expectation_cannot_fetch():
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(lambda _: pytest.fail("Unexpected provider call"))
+ ) as client:
+ source = await GmailSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=GmailConfig(),
+ )
+ with pytest.raises(ValueError, match="attested mailbox"):
+ await source.capture_page(
+ CompletedScope(record_type="message"),
+ ScanContinuation(value={}),
+ files=MagicMock(),
+ )
diff --git a/backend/airweave/platform/sources/tests/test_gmail_retry.py b/backend/airweave/platform/sources/tests/test_gmail_retry.py
new file mode 100644
index 0000000000..895ad84f05
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_gmail_retry.py
@@ -0,0 +1,127 @@
+"""Exercise the decorated request boundary, not just its exception predicate."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+from tenacity import wait_none
+
+from airweave.domains.auth_provider.exceptions import (
+ AuthProviderAuthError,
+ AuthProviderRateLimitError,
+ AuthProviderServerError,
+)
+from airweave.domains.sources.exceptions import (
+ SourceAuthError,
+ SourceRateLimitError,
+ SourceServerError,
+)
+from airweave.domains.sources.token_providers.protocol import AuthProviderKind
+from airweave.platform.sources.gmail import GmailSource
+
+
+def request(method, effects):
+ upstream = AsyncMock(side_effect=effects)
+ source = SimpleNamespace(
+ _get_mime_json=upstream,
+ _authed_headers=AsyncMock(return_value={}),
+ http_client=SimpleNamespace(get=upstream),
+ logger=Mock(),
+ )
+ decorated = getattr(GmailSource, method).retry_with(wait=wait_none())
+ return source, upstream, decorated
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["_get", "_get_capture_json"])
+@pytest.mark.parametrize(
+ "error",
+ [
+ SourceServerError(status_code=503),
+ AuthProviderServerError(status_code=503),
+ AuthProviderRateLimitError(retry_after=60),
+ ],
+)
+async def test_transient_server_error_retries_actual_request(method, error):
+ # For _get, return a real response; canonical JSON returns an already parsed dict.
+ success = (
+ httpx.Response(200, json={"ok": True}, request=httpx.Request("GET", "https://example.test"))
+ if method == "_get"
+ else {"ok": True}
+ )
+ source, upstream, call = request(method, [error, success])
+ source.short_name = "gmail"
+ source.auth = SimpleNamespace(provider_kind=None)
+ assert await call(source, "https://example.test") == {"ok": True}
+ assert upstream.await_count == 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["_get", "_get_capture_json"])
+@pytest.mark.parametrize(
+ "error",
+ [
+ SourceServerError(status_code=503),
+ AuthProviderServerError(status_code=503),
+ AuthProviderRateLimitError(),
+ ],
+)
+async def test_server_error_exhausts_exactly_five_attempts(method, error):
+ source, upstream, call = request(method, error)
+ with pytest.raises(type(error)) as caught:
+ await call(source, "https://example.test")
+ assert caught.value is error
+ assert upstream.await_count == 5
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["_get", "_get_capture_json"])
+@pytest.mark.parametrize(
+ "error",
+ [
+ SourceAuthError(
+ "fixture auth",
+ source_short_name="gmail",
+ status_code=401,
+ token_provider_kind=AuthProviderKind.STATIC,
+ ),
+ AuthProviderAuthError("fixture proxy authentication"),
+ AuthProviderRateLimitError(retry_after=180),
+ SourceRateLimitError(retry_after=180),
+ httpx.HTTPStatusError(
+ "rate limit",
+ request=httpx.Request("GET", "https://example.test"),
+ response=httpx.Response(429, headers={"Retry-After": "180"}),
+ ),
+ httpx.HTTPStatusError(
+ "unauthorized",
+ request=httpx.Request("GET", "https://example.test"),
+ response=httpx.Response(401),
+ ),
+ RuntimeError("budget"),
+ asyncio.CancelledError(),
+ ],
+)
+async def test_auth_budget_and_cancellation_do_not_retry(method, error):
+ source, upstream, call = request(method, error)
+ with pytest.raises(type(error)):
+ await call(source, "https://example.test")
+ assert upstream.await_count == 1
+
+
+@pytest.mark.asyncio
+async def test_proxy_retry_after_within_bound_is_honored():
+ source, upstream, _ = request(
+ "_get_capture_json",
+ [
+ AuthProviderRateLimitError(retry_after=60),
+ {"ok": True},
+ ],
+ )
+ sleep = AsyncMock()
+ call = GmailSource._get_capture_json.retry_with(sleep=sleep)
+ assert await call(source, "https://example.test") == {"ok": True}
+ sleep.assert_awaited_once_with(60)
+ assert upstream.await_count == 2
diff --git a/backend/airweave/platform/sources/tests/test_outlook_calendar_capture.py b/backend/airweave/platform/sources/tests/test_outlook_calendar_capture.py
new file mode 100644
index 0000000000..393ad119fa
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_outlook_calendar_capture.py
@@ -0,0 +1,281 @@
+"""Synthetic Graph interpretation, not live provider or SQL recovery qualification."""
+
+from datetime import datetime, timezone
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+import httpx
+import pytest
+
+from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleVersion
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError, SourceServerError
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import CalendarOccurrenceWindow, OutlookCalendarConfig
+from airweave.platform.sources.outlook_calendar_capture import BASE, OutlookCalendarCapture
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient
+
+WINDOW = CalendarOccurrenceWindow(start="2026-10-01T00:00:00Z", end="2026-11-01T00:00:00Z")
+ROOT = CompletedScope(record_type="calendar")
+CALENDAR = {"id": "calendar", "name": "Original", "canViewPrivateItems": True}
+EVENT = {
+ "id": "single",
+ "changeKey": "v1",
+ "type": "singleInstance",
+ "start": {"dateTime": "2026-10-02T10:00:00.0000000", "timeZone": "UTC"},
+ "end": {"dateTime": "2026-10-02T11:00:00.0000000", "timeZone": "UTC"},
+ "isAllDay": False,
+ "isCancelled": False,
+ "hasAttachments": False,
+ "originalStartTimeZone": "India Standard Time",
+ "nativeExtension": {"keep": [1, True, None]},
+}
+
+
+def retained(identity, payload, parent=None):
+ return SourceRecord(
+ id=uuid4(),
+ sync_id=uuid4(),
+ identity=identity,
+ parent=parent,
+ revision=1,
+ payload=payload,
+ payload_schema_version=1,
+ capture_hash="a" * 64,
+ content_hash=None,
+ completeness="partial",
+ observed_at=datetime.now(timezone.utc),
+ source_created_at=None,
+ source_updated_at=None,
+ deleted_at=None,
+ removal_reason=None,
+ blobs=(),
+ indexed_revision=None,
+ indexed_pipeline_version=None,
+ )
+
+
+PARENT = retained(RecordIdentity(record_type="calendar", native_id="calendar"), CALENDAR)
+
+
+async def source_for(client):
+ return await OutlookCalendarCapture.create(
+ graph=OutlookGraphClient(
+ StaticTokenProvider("secret"), client, "outlook_calendar", "principal"
+ ),
+ config=OutlookCalendarConfig(
+ capture_originals=True, expected_principal_id="principal", occurrence_window=WINDOW
+ ),
+ )
+
+
+async def initial(source, parent=PARENT):
+ prepared = await source.prepare_cycle(None)
+ assert prepared.mode == "full"
+ assert not source.capture_cycle_configuration.scope_changes
+ cycle = CaptureCycle(
+ version=CycleVersion(cycle_id=uuid4(), revision=1),
+ configuration=source.capture_cycle_configuration,
+ mode=prepared.mode,
+ source_plan=prepared.source_plan,
+ )
+ scope = source.child_scope(parent, "event") if parent else ROOT
+ plan = await source.prepare_scope(scope, cycle, None, parent=parent, force_full=False)
+ return scope, source.initial_scope_continuation(scope, cycle, plan)
+
+
+@pytest.mark.asyncio
+async def test_one_event_owner_across_full_and_expanded_pages_preserves_native_cancelled_json():
+ reads = []
+
+ def graph(request):
+ path = request.url.path
+ if path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ reads.append(path)
+ assert request.headers["Prefer"] == 'IdType="ImmutableId"'
+ if path.endswith("/events"):
+ assert "startDateTime" not in request.url.params
+ return httpx.Response(200, json={"value": [{"id": "single"}, {"id": "master"}]})
+ if path.endswith("/calendarView"):
+ assert request.url.params["startDateTime"] == WINDOW.start.isoformat()
+ return httpx.Response(200, json={"value": [{"id": "single"}, {"id": "exception"}]})
+ native_id = path.rsplit("/", 1)[1]
+ event = {**EVENT, "id": native_id}
+ if native_id == "master":
+ event.update(type="seriesMaster", recurrence={"pattern": {"type": "weekly"}})
+ if native_id == "exception":
+ event.update(
+ type="exception",
+ seriesMasterId="master",
+ isCancelled=True,
+ originalStart="2026-10-01T10:00:00Z",
+ )
+ return httpx.Response(200, json=event)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ scope, cursor = await initial(source)
+ first = await source.capture_page(scope, cursor, files=MagicMock(), parent=PARENT)
+ second = await source.capture_page(
+ scope, first.continuation, files=MagicMock(), parent=PARENT
+ )
+ assert not first.final and second.final
+ assert first.records[0].identity == second.records[0].identity
+ assert first.records[0].payload == EVENT
+ assert second.records[1].kind == "upsert" and second.records[1].payload["isCancelled"] is True
+ assert {r.parent for r in (*first.records, *second.records)} == {PARENT.identity}
+ assert all(r.completeness == "partial" for r in (*first.records, *second.records))
+ assert reads.count("/v1.0/me/calendars/calendar/events/single") == 2
+
+
+@pytest.mark.asyncio
+async def test_late_failure_retries_same_page_and_resumed_process_keeps_frozen_window():
+ fail = True
+ requests = []
+ next_link = f"{BASE}/calendars/calendar/events?$skiptoken=second"
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ requests.append(str(request.url))
+ if request.url.path.endswith("/events"):
+ if "$skiptoken" not in request.url.params:
+ return httpx.Response(
+ 200, json={"value": [{"id": "first"}], "@odata.nextLink": next_link}
+ )
+ return httpx.Response(200, json={"value": [{"id": "second"}, {"id": "third"}]})
+ if request.url.path.endswith("/third") and fail:
+ return httpx.Response(503)
+ return httpx.Response(200, json={**EVENT, "id": request.url.path.rsplit("/", 1)[1]})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ scope, cursor = await initial(source)
+ first = await source.capture_page(scope, cursor, files=MagicMock(), parent=PARENT)
+ serialized = first.continuation.model_dump_json()
+ with pytest.raises(SourceServerError):
+ await source.capture_page(scope, first.continuation, files=MagicMock(), parent=PARENT)
+ assert first.continuation.model_dump_json() == serialized
+ fail = False
+ resumed = await source_for(client)
+ second = await resumed.capture_page(
+ scope,
+ ScanContinuation.model_validate_json(serialized),
+ files=MagicMock(),
+ parent=PARENT,
+ )
+ assert [r.identity.native_id for r in second.records] == ["second", "third"]
+ assert second.continuation.value["phase"] == "calendarView"
+ assert requests.count(next_link) == 2
+ assert second.continuation.value["context"]["window"] == WINDOW.model_dump(mode="json")
+
+
+@pytest.mark.asyncio
+async def test_exact_omissions_keep_historical_occurrences_and_ambiguous404_fails():
+ missing = False
+ historical = {
+ **EVENT,
+ "id": "old",
+ "type": "occurrence",
+ "seriesMasterId": "master",
+ "start": {"dateTime": "2020-01-01T00:00:00", "timeZone": "UTC"},
+ }
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ return httpx.Response(404) if missing else httpx.Response(200, json=historical)
+
+ record = retained(
+ RecordIdentity(record_type="event", native_id="old", container_id="calendar"),
+ historical,
+ PARENT.identity,
+ )
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ observation = await source.refresh_known(record, files=MagicMock())
+ assert observation.kind == "upsert" and observation.payload == historical
+ missing = True
+ with pytest.raises(SourceEntityNotFoundError):
+ await source.refresh_known(record, files=MagicMock())
+
+
+@pytest.mark.asyncio
+async def test_inventory_permission_metadata_and_foreign_link_rejected_before_hydration():
+ foreign = False
+ calls = []
+
+ def graph(request):
+ calls.append(request.url.path)
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ if request.url.path.endswith("/calendars"):
+ page = {"value": [{"id": "calendar"}]}
+ if foreign:
+ page["@odata.nextLink"] = f"{BASE}/calendars/another/events?$skiptoken=private"
+ return httpx.Response(200, json=page)
+ return httpx.Response(200, json=CALENDAR)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ scope, cursor = await initial(source, parent=None)
+ page = await source.capture_page(scope, cursor, files=MagicMock())
+ assert page.final and page.records[0].payload == CALENDAR
+ assert "canViewPrivateItems" in page.records[0].descendant_visibility_fields
+ foreign = True
+ before = calls.count("/v1.0/me/calendars/calendar")
+ with pytest.raises(OutlookBoundaryError, match="changed its native collection"):
+ await source.capture_page(scope, cursor, files=MagicMock())
+ assert calls.count("/v1.0/me/calendars/calendar") == before
+
+
+@pytest.mark.asyncio
+async def test_repeated_token_and_removed_full_listing_never_complete():
+ removed = False
+ next_link = f"{BASE}/calendars/calendar/events?$skiptoken=loop"
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ if removed:
+ return httpx.Response(
+ 200, json={"value": [{"id": "x", "@removed": {"reason": "deleted"}}]}
+ )
+ return httpx.Response(200, json={"value": [], "@odata.nextLink": next_link})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ scope, cursor = await initial(source)
+ page = await source.capture_page(scope, cursor, files=MagicMock(), parent=PARENT)
+ with pytest.raises(OutlookBoundaryError, match="repeated"):
+ await source.capture_page(scope, page.continuation, files=MagicMock(), parent=PARENT)
+ removed = True
+ with pytest.raises(OutlookBoundaryError, match="Invalid Outlook calendar collection"):
+ await source.capture_page(scope, cursor, files=MagicMock(), parent=PARENT)
+
+
+@pytest.mark.asyncio
+async def test_principal_and_scope_binding_reject_before_any_provider_read():
+ calls = []
+
+ def graph(request):
+ calls.append(str(request.url))
+ return httpx.Response(200, json={"id": "principal"})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await source_for(client)
+ scope, cursor = await initial(source)
+ before = len(calls)
+ wrong = cursor.model_dump(mode="json")
+ wrong["value"]["context"]["calendar_id"] = "another"
+ with pytest.raises(OutlookBoundaryError, match="another scope"):
+ await source.capture_page(
+ scope, ScanContinuation.model_validate(wrong), files=MagicMock(), parent=PARENT
+ )
+ source.graph.verified_principal_id = None
+ with pytest.raises(OutlookBoundaryError, match="not attested"):
+ await source.capture_page(scope, cursor, files=MagicMock(), parent=PARENT)
+ assert len(calls) == before
diff --git a/backend/airweave/platform/sources/tests/test_outlook_identity.py b/backend/airweave/platform/sources/tests/test_outlook_identity.py
new file mode 100644
index 0000000000..3b550f9758
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_outlook_identity.py
@@ -0,0 +1,343 @@
+"""Outlook native identity and transport boundaries, without provider qualification."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+from pydantic import SecretStr
+
+from airweave.domains.auth_provider.exceptions import AuthProviderAuthError
+from airweave.domains.sources.exceptions import SourceAuthError
+from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import OutlookCalendarConfig, OutlookMailConfig
+from airweave.platform.http_client.composio_transport import ComposioTransport
+from airweave.platform.sources.outlook_calendar import OutlookCalendarSource
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError
+from airweave.platform.sources.outlook_mail import OutlookMailSource
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "source_cls,config_cls",
+ [
+ (OutlookMailSource, OutlookMailConfig),
+ (OutlookCalendarSource, OutlookCalendarConfig),
+ ],
+)
+async def test_managed_attestation_reset_blocks_list_and_attachment_reads(source_cls, config_cls):
+ principal = "opaque-principal"
+ requests = []
+
+ def graph(request):
+ requests.append(request)
+ assert "authorization" not in request.headers
+ assert "prefer" not in request.headers # No silent legacy ID migration.
+ return httpx.Response(
+ 200, json={"id": principal} if request.url.path == "/v1.0/me" else {"value": []}
+ )
+
+ auth = ManagedAuthProvider(
+ api_key=SecretStr("private"),
+ connected_account_id="account",
+ allowed_hosts=frozenset({"graph.microsoft.com"}),
+ )
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ with pytest.raises(OutlookBoundaryError, match="trusted expected"):
+ await source_cls.create(
+ auth=auth, logger=MagicMock(), http_client=client, config=config_cls()
+ )
+ assert not requests
+ source = await source_cls.create(
+ auth=auth,
+ logger=MagicMock(),
+ http_client=client,
+ config=config_cls(expected_principal_id=principal),
+ )
+ await source._get("https://graph.microsoft.com/v1.0/me/messages/id/attachments")
+ principal = "wrong-private-principal"
+ with pytest.raises(OutlookBoundaryError, match="does not match") as error:
+ await source.validate()
+ assert principal not in str(error.value)
+ count = len(requests)
+ with pytest.raises(OutlookBoundaryError, match="attested"):
+ await source._get("https://graph.microsoft.com/v1.0/me/events/id/attachments")
+ assert len(requests) == count
+ with pytest.raises(OutlookBoundaryError, match="does not match"):
+ await anext(source.generate_entities())
+
+
+@pytest.mark.asyncio
+async def test_managed_proxy_and_native_401_keep_different_error_authority():
+ native_status = 200
+ proxy_status = 200
+
+ def proxy(request):
+ import json
+
+ payload = json.loads(request.content)
+ assert payload["connected_account_id"] == "bound-account"
+ assert all(p["name"].lower() != "authorization" for p in payload["parameters"])
+ return httpx.Response(
+ proxy_status,
+ json={
+ "status": native_status,
+ "data": {"id": "principal", "error": "private-token"},
+ },
+ )
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as proxy_client:
+ transport = ComposioTransport(
+ api_key="private",
+ connected_account_id="bound-account",
+ allowed_hosts={"graph.microsoft.com"},
+ client=proxy_client,
+ )
+ auth = ManagedAuthProvider(
+ api_key=SecretStr("private"),
+ connected_account_id="bound-account",
+ allowed_hosts=frozenset({"graph.microsoft.com"}),
+ )
+ async with httpx.AsyncClient(transport=transport) as client:
+ source = await OutlookMailSource.create(
+ auth=auth,
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookMailConfig(expected_principal_id="principal"),
+ )
+ native_status = 401
+ with pytest.raises(SourceAuthError) as error:
+ await source._get("https://graph.microsoft.com/v1.0/me/messages")
+ assert "private-token" not in str(error.value)
+ assert source.graph.verified_principal_id is None
+ proxy_status = 401
+ with pytest.raises(AuthProviderAuthError):
+ await source.validate()
+
+
+@pytest.mark.asyncio
+async def test_unsafe_continuations_are_fatal_before_request_even_with_direct_credentials():
+ calls = []
+
+ def graph(request):
+ calls.append(request)
+ return httpx.Response(200, json={"value": []})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await OutlookMailSource.create(
+ auth=StaticTokenProvider("private"),
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookMailConfig(),
+ )
+ for url in (
+ "https://evil.example/v1.0/me/messages?token=private",
+ "https://graph.microsoft.com/v1.0/users/other/messages",
+ "https://graph.microsoft.com/beta/me/messages",
+ "https://graph.microsoft.com/v1.0/me/%2e%2e/users/other",
+ "https://graph.microsoft.com/v1.0/me/%252e%252e/users/other",
+ "https://graph.microsoft.com/v1.0/me/messages#private",
+ ):
+ with pytest.raises(OutlookBoundaryError) as error:
+ await source._get(url)
+ assert "private" not in str(error.value)
+ assert not calls
+ await source.validate() # No /me permission newly required for unbound direct connections.
+ assert len(calls) == 1
+ assert calls[0].url.path == "/v1.0/me/mailFolders"
+ assert calls[0].headers["authorization"] == "Bearer private"
+
+
+@pytest.mark.asyncio
+async def test_direct_refresh_rechecks_principal_before_retrying_attachment():
+ class RefreshingAuth(StaticTokenProvider):
+ @property
+ def supports_refresh(self):
+ return True
+
+ async def force_refresh(self):
+ self._token = "replacement"
+ return self._token
+
+ paths = []
+
+ def graph(request):
+ paths.append(request.url.path)
+ if request.url.path == "/v1.0/me":
+ identity = (
+ "original" if request.headers["authorization"] == "Bearer initial" else "other"
+ )
+ return httpx.Response(200, json={"id": identity})
+ return httpx.Response(401, json={"error": {"message": "private-details"}})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await OutlookMailSource.create(
+ auth=RefreshingAuth("initial"),
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookMailConfig(expected_principal_id="original"),
+ )
+ with pytest.raises(OutlookBoundaryError, match="does not match"):
+ await source._get("https://graph.microsoft.com/v1.0/me/messages/id/attachments")
+ assert paths == ["/v1.0/me", "/v1.0/me/messages/id/attachments", "/v1.0/me"]
+ assert source.graph.verified_principal_id is None
+
+
+@pytest.mark.asyncio
+async def test_broker_outlook_mapping_reuses_bound_graph_transport(monkeypatch):
+ from unittest.mock import AsyncMock
+
+ from airweave.domains.auth_provider.exceptions import AuthProviderConfigError
+ from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider
+
+ provider = await ComposioAuthProvider.create(
+ credentials={"api_key": "secret"}, config={"account_id": "selected", "user_id": "owner"}
+ )
+ account = {
+ "id": "selected",
+ "user_id": "owner",
+ "toolkit": {"slug": "outlook"},
+ "status": "ACTIVE",
+ }
+ monkeypatch.setattr(provider, "_get_with_auth", AsyncMock(return_value=account))
+ for name in ("outlook_mail", "outlook_calendar"):
+ result = await provider.get_auth_result(name, [])
+ assert result.credentials is None
+ assert result.managed_auth.allowed_hosts == frozenset({"graph.microsoft.com"})
+ assert result.managed_auth.connected_account_id == "selected"
+ account["toolkit"]["slug"] = "other"
+ with pytest.raises(AuthProviderConfigError, match="toolkit"):
+ await provider.get_auth_result("outlook_mail", [])
+
+
+@pytest.mark.asyncio
+async def test_malformed_principal_and_error_envelope_do_not_become_empty_success():
+ data = {"id": "principal"}
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(lambda _: httpx.Response(200, json=data))
+ ) as client:
+ source = await OutlookCalendarSource.create(
+ auth=StaticTokenProvider("secret"),
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookCalendarConfig(expected_principal_id="principal"),
+ )
+ for invalid_body in (None, [], "private-body"):
+ data = invalid_body
+ with pytest.raises(OutlookBoundaryError, match="invalid JSON"):
+ await source._get("https://graph.microsoft.com/v1.0/me/calendars")
+ data = {"error": {"message": "private-message"}}
+ with pytest.raises(OutlookBoundaryError, match="error envelope") as error:
+ await source._get("https://graph.microsoft.com/v1.0/me/calendars")
+ assert "private-message" not in str(error.value)
+ with pytest.raises(OutlookBoundaryError, match="invalid principal"):
+ await source.validate()
+ assert source.graph.verified_principal_id is None
+
+
+@pytest.mark.asyncio
+async def test_attachment_pagination_cannot_swallow_unsafe_continuation():
+ calls = []
+
+ def graph(request):
+ calls.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "value": [],
+ "@odata.nextLink": "https://graph.microsoft.com/v1.0/users/other/messages?token=private",
+ },
+ )
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await OutlookMailSource.create(
+ auth=StaticTokenProvider("secret"),
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookMailConfig(),
+ )
+ with pytest.raises(OutlookBoundaryError, match="outside the bound mailbox"):
+ await anext(source._process_attachments("message-id", [], None))
+ assert len(calls) == 1
+ assert calls[0].url.path == "/v1.0/me/messages/message-id/attachments"
+ assert "private" not in str(source.logger.method_calls)
+
+
+@pytest.mark.parametrize("capture_originals", [False, True])
+async def test_source_creation_selects_capture_adapter_without_replacing_auth(capture_originals):
+ from airweave.domains.entities.canonical.page_source import CanonicalPageSource
+
+ def graph(request):
+ assert request.url.path == "/v1.0/me"
+ return httpx.Response(200, json={"id": "principal"})
+
+ auth = StaticTokenProvider("synthetic")
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await OutlookMailSource.create(
+ auth=auth,
+ logger=MagicMock(),
+ http_client=client,
+ config=OutlookMailConfig(
+ capture_originals=capture_originals,
+ expected_principal_id="principal",
+ included_folders=[],
+ excluded_folders=[],
+ ),
+ )
+ assert source.auth is auth
+ if capture_originals:
+ assert isinstance(source.capture_page_source, CanonicalPageSource)
+ assert source.capture_page_source.graph is source.graph
+ else:
+ assert source.capture_page_source is None
+
+
+async def test_delta_expiry_classification_requires_structured_code_and_delta_endpoint():
+ from airweave.domains.sources.exceptions import SourceEntityNotFoundError
+ from airweave.platform.sources.outlook_graph import OutlookDeltaExpiredError, OutlookGraphClient
+
+ code = "syncStateNotFound"
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ return httpx.Response(404, json={"error": {"code": code, "message": "private-value"}})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = OutlookGraphClient(
+ StaticTokenProvider("secret"), client, "outlook_mail", "principal"
+ )
+ await source.verify_principal()
+ delta = "https://graph.microsoft.com/v1.0/me/mailFolders/inbox/messages/delta"
+ with pytest.raises(OutlookDeltaExpiredError) as error:
+ await source.get(delta)
+ assert "private-value" not in str(error.value)
+ with pytest.raises(SourceEntityNotFoundError):
+ await source.get("https://graph.microsoft.com/v1.0/me/messages/message")
+ code = "ErrorItemNotFound"
+ with pytest.raises(SourceEntityNotFoundError):
+ await source.get(delta)
+
+
+def test_owned_calendar_config_requires_principal_and_preserves_explicit_window():
+ """Owned activation needs identity; recurrence bounds are explicit and reproducible."""
+ from datetime import datetime, timezone
+
+ from pydantic import ValidationError
+
+ from airweave.platform.configs.config import CalendarOccurrenceWindow
+
+ with pytest.raises(ValidationError, match="expected_principal_id"):
+ OutlookCalendarConfig(capture_originals=True)
+ start = datetime(2026, 10, 1, tzinfo=timezone.utc)
+ end = datetime(2026, 11, 1, tzinfo=timezone.utc)
+ explicit = CalendarOccurrenceWindow(start=start, end=end)
+ config = OutlookCalendarConfig(
+ capture_originals=True,
+ expected_principal_id="principal",
+ occurrence_window=explicit,
+ )
+ assert config.resolved_window() == explicit
+ assert not OutlookCalendarConfig().capture_originals
+ with pytest.raises(ValidationError):
+ OutlookCalendarConfig(occurrence_window={"start": end, "end": start})
diff --git a/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py b/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py
new file mode 100644
index 0000000000..3279531c4e
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py
@@ -0,0 +1,459 @@
+"""Synthetic native interpretation; SQL page durability is tested by the shared engine suite."""
+
+from unittest.mock import AsyncMock, MagicMock
+from uuid import uuid4
+
+import httpx
+import pytest
+
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.requests import BlobReference, CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceEntityNotFoundError
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import OutlookMailConfig
+from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient
+from airweave.platform.sources.outlook_mail_capture import OutlookMailCapture
+from airweave.platform.sources.outlook_mail_models import OutlookMailContinuation
+
+MIME = b"From: from@example.com\r\nTo: to@example.com\r\nSubject: Original\r\n\r\nFull body\r\n"
+MESSAGE = {
+ "id": "immutable-message",
+ "changeKey": "version-one",
+ "parentFolderId": "folder-b",
+ "subject": "Original",
+ "from": {"emailAddress": {"address": "from@example.com"}},
+ "sender": {"emailAddress": {"address": "delegate@example.com"}},
+ "body": {"contentType": "text", "content": "Full body"},
+ "bodyPreview": "Full",
+ "hasAttachments": False,
+ "nativeExtension": {"keep": [1, True]},
+}
+ROOT = CompletedScope(record_type="message")
+
+
+def files(maximum=10000):
+ result = MagicMock()
+ result.MAX_FILE_SIZE_BYTES = maximum
+ result.store_canonical_blob = AsyncMock(
+ return_value=BlobReference(
+ key="original", sha256="a" * 64, size_bytes=len(MIME), media_type="message/rfc822"
+ )
+ )
+ return result
+
+
+async def capture(client, *, config=None):
+ return await OutlookMailCapture.create(
+ graph=OutlookGraphClient(
+ StaticTokenProvider("secret"), client, "outlook_mail", "principal"
+ ),
+ config=config
+ or OutlookMailConfig(
+ expected_principal_id="principal", included_folders=[], excluded_folders=[]
+ ),
+ )
+
+
+def initial(source):
+ return ScanContinuation(
+ value=OutlookMailContinuation(fingerprint=source.fingerprint).model_dump(mode="json")
+ )
+
+
+async def first_message_page(source, retained):
+ progress = initial(source)
+ while True:
+ page = await source.capture_page(ROOT, progress, files=retained)
+ if page.records or page.final:
+ return page
+ progress = page.continuation
+
+
+def topology(request):
+ if request.url.path == "/v1.0/me/mailFolders":
+ assert request.url.params["includeHiddenFolders"] == "true"
+ return httpx.Response(200, json={"value": [{"id": "folder-b"}]})
+ if request.url.path.endswith("/childFolders"):
+ assert request.url.params["includeHiddenFolders"] == "true"
+ return httpx.Response(200, json={"value": []})
+ return None
+
+
+@pytest.mark.asyncio
+async def test_full_page_retains_exact_json_and_mime_and_resumes_pending_id():
+ requested = []
+ fail_second = True
+
+ def graph(request):
+ requested.append(request.url.path)
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ assert request.headers["Prefer"] == 'IdType="ImmutableId"'
+ folder_page = topology(request)
+ if folder_page is not None:
+ return folder_page
+ if request.url.path.endswith("/messages/delta"):
+ return httpx.Response(
+ 200,
+ json={
+ "value": [{"id": "immutable-message"}, {"id": "second"}],
+ "@odata.deltaLink": str(request.url.copy_with(query=b"$deltatoken=end")),
+ },
+ )
+ if request.url.path.endswith("/$value"):
+ if fail_second and "/second/" in request.url.path:
+ raise httpx.ReadError("private-url-failure")
+ return httpx.Response(200, content=MIME, headers={"content-type": "message/rfc822"})
+ mid = request.url.path.rsplit("/", 1)[-1]
+ return httpx.Response(200, json={**MESSAGE, "id": mid})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ retained = files()
+ first = await first_message_page(source, retained)
+ assert first.records[0].payload == MESSAGE
+ assert first.records[0].parent is None
+ assert first.records[0].completeness == "partial"
+ assert first.records[0].blobs[0].source_path is None
+ assert not first.final
+ retained.store_canonical_blob.assert_awaited_once_with(MIME, media_type="message/rfc822")
+ saved = first.continuation.model_copy(deep=True)
+ with pytest.raises(httpx.RequestError) as error:
+ await source.capture_page(ROOT, saved, files=retained)
+ assert "private-url-failure" not in str(error.value)
+ assert saved == first.continuation
+ fail_second = False
+ recreated = await capture(client)
+ second = await recreated.capture_page(ROOT, saved, files=retained)
+ assert second.final and second.records[0].identity.native_id == "second"
+ assert requested.count("/v1.0/me/mailFolders/folder-b/messages/delta") == 1
+ assert second.provider_checkpoint is not None
+
+
+@pytest.mark.asyncio
+async def test_current_folder_decides_move_membership_and_ambiguous_404_stops():
+ message = dict(MESSAGE)
+ unavailable = False
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ if "/mailFolders/" in request.url.path:
+ return httpx.Response(200, json={"id": request.url.path.rsplit("/", 1)[-1]})
+ if unavailable:
+ return httpx.Response(404, json={"error": {"code": "ErrorItemNotFound"}})
+ if request.url.path.endswith("/$value"):
+ return httpx.Response(200, content=MIME, headers={"content-type": "text/plain"})
+ return httpx.Response(200, json=message)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(
+ client,
+ config=OutlookMailConfig(
+ expected_principal_id="principal",
+ included_folders=["folder-a", "folder-b"],
+ excluded_folders=["excluded"],
+ ),
+ )
+ retained = files()
+ # Duplicate discoveries before/after a folder move keep the same root identity.
+ first = await source.message(message["id"], files=retained)
+ message["parentFolderId"] = "folder-a"
+ second = await source.message(message["id"], files=retained)
+ assert first.identity == second.identity and first.kind == second.kind == "upsert"
+ assert first.parent is second.parent is None
+ message["parentFolderId"] = "excluded"
+ outside = await source.message(message["id"], files=retained)
+ assert outside.removal_reason == "scope_removed"
+ assert retained.store_canonical_blob.await_count == 2
+ unavailable = True
+ previous = SourceRecord(
+ id=uuid4(),
+ sync_id=uuid4(),
+ identity=first.identity,
+ revision=1,
+ payload=first.payload,
+ payload_schema_version=1,
+ capture_hash="a" * 64,
+ content_hash=None,
+ completeness=first.completeness,
+ observed_at=first.observed_at,
+ source_created_at=None,
+ source_updated_at=None,
+ deleted_at=None,
+ removal_reason=None,
+ blobs=first.blobs,
+ indexed_revision=None,
+ indexed_pipeline_version=None,
+ )
+ with pytest.raises(SourceEntityNotFoundError):
+ await source.refresh_known(previous, files=retained)
+
+
+@pytest.mark.asyncio
+async def test_cursor_binding_and_wrong_collection_fail_before_message_reads():
+ requested = []
+ next_link = "https://graph.microsoft.com/v1.0/me/mailFolders/folder-b/messages/delta?$skip=next"
+
+ def graph(request):
+ requested.append(request.url.path)
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ folder_page = topology(request)
+ if folder_page is not None:
+ return folder_page
+ return httpx.Response(200, json={"value": [], "@odata.nextLink": next_link})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ wrong = initial(source).model_copy(update={"value": {"fingerprint": "a" * 64}})
+ with pytest.raises(OutlookBoundaryError, match="another capture"):
+ await source.capture_page(ROOT, wrong, files=files())
+ assert requested == ["/v1.0/me"]
+ next_link = "https://graph.microsoft.com/v1.0/me/events"
+ with pytest.raises(OutlookBoundaryError, match="changed the native collection"):
+ await first_message_page(source, files())
+ next_link = (
+ "https://graph.microsoft.com/v1.0/me/mailFolders/folder-b/messages/delta?token="
+ + "x" * 66000
+ )
+ with pytest.raises(OutlookBoundaryError, match="64 KiB"):
+ await first_message_page(source, files())
+
+
+@pytest.mark.asyncio
+async def test_mime_size_omission_is_explicit_and_change_key_race_cannot_commit():
+ changed = False
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ if request.url.path.endswith("/$value"):
+ return httpx.Response(
+ 200, content=MIME, headers={"content-type": "Message/RFC822; charset=utf-8"}
+ )
+ value = dict(MESSAGE)
+ if changed and "$select" in request.url.params:
+ value["changeKey"] = "new-version"
+ return httpx.Response(200, json=value)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ bounded = files(maximum=8)
+ original = await source.message(MESSAGE["id"], files=bounded)
+ assert original.completeness == "metadata_only" and original.blobs == ()
+ assert original.payload == MESSAGE
+ bounded.store_canonical_blob.assert_not_awaited()
+ changed = True
+ with pytest.raises(OutlookBoundaryError, match="changed during"):
+ await source.message(MESSAGE["id"], files=files())
+
+
+@pytest.mark.asyncio
+async def test_missing_received_date_cannot_broaden_date_selection():
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ assert not request.url.path.endswith("/$value")
+ return httpx.Response(200, json=MESSAGE)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(
+ client,
+ config=OutlookMailConfig(
+ expected_principal_id="principal",
+ after_date="2026/01/01",
+ included_folders=[],
+ excluded_folders=[],
+ ),
+ )
+ with pytest.raises(OutlookBoundaryError, match="requires a received timestamp"):
+ await source.message(MESSAGE["id"], files=files())
+
+
+@pytest.mark.parametrize("order", [("old", "new"), ("new", "old")])
+async def test_delta_move_orders_hydrate_mailbox_and_changes_preserve_scope_removal(order):
+ """Removed events are invalidations, never deletion authority."""
+ folder = "new"
+ unavailable = False
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ assert request.headers["Prefer"] == 'IdType="ImmutableId"'
+ if request.url.path.endswith("/messages/delta"):
+ event = {"id": MESSAGE["id"]}
+ if "/old/" in request.url.path:
+ event["@removed"] = {"reason": "deleted"}
+ return httpx.Response(
+ 200,
+ json={
+ "value": [event],
+ "@odata.deltaLink": str(request.url.copy_with(query=b"$deltatoken=next")),
+ },
+ )
+ if unavailable:
+ return httpx.Response(404, json={"error": {"code": "ErrorItemNotFound"}})
+ if request.url.path.endswith("/$value"):
+ return httpx.Response(200, content=MIME, headers={"content-type": "message/rfc822"})
+ return httpx.Response(200, json={**MESSAGE, "parentFolderId": folder})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ source.excluded_ids = frozenset({"excluded"})
+ state = OutlookMailContinuation(
+ fingerprint=source.fingerprint,
+ mode="changes",
+ phase="messages",
+ folders_to_visit=(),
+ remaining_folders=order,
+ )
+ cursor = source._continuation(state)
+ first = await source.capture_page(ROOT, cursor, files=files())
+ second = await source.capture_page(ROOT, first.continuation, files=files())
+ assert first.records[0].identity == second.records[0].identity
+ assert first.records[0].kind == second.records[0].kind == "upsert"
+ assert first.provider_checkpoint is None and second.provider_checkpoint is not None
+ assert second.final
+ folder = "excluded"
+ withdrawn = await source.capture_page(ROOT, cursor, files=files())
+ assert withdrawn.records[0].removal_reason == "scope_removed"
+ unavailable = True
+ with pytest.raises(SourceEntityNotFoundError):
+ await source.capture_page(ROOT, cursor, files=files())
+ assert cursor == source._continuation(state)
+
+
+async def test_recursive_hidden_topology_new_folder_and_disappeared_folder_reset():
+ from airweave.domains.entities.canonical.page_source import InvalidCaptureCheckpoint
+
+ calls = []
+ unknown_folder = False
+
+ def graph(request):
+ path = request.url.path
+ calls.append(path)
+ if path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ assert request.url.params["includeHiddenFolders"] == "true"
+ if path.endswith("/mailFolders"):
+ value = [{"id": "parent"}]
+ elif path.endswith("/parent/childFolders"):
+ value = [{"id": "hidden-child", "isHidden": True}]
+ else:
+ value = []
+ if unknown_folder and value:
+ value[0]["@odata.type"] = "#microsoft.graph.unknownFolder"
+ return httpx.Response(200, json={"value": value})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ cursor = initial(source)
+ for _ in range(3):
+ page = await source.capture_page(ROOT, cursor, files=files())
+ cursor = page.continuation
+ state = OutlookMailContinuation.model_validate(cursor.value)
+ assert state.phase == "messages"
+ assert set(state.remaining_folders) == {"parent", "hidden-child"}
+ assert calls[-1].endswith("/hidden-child/childFolders")
+ missing = OutlookMailContinuation(
+ fingerprint=source.fingerprint,
+ mode="changes",
+ folder_links={"removed": source._delta_url("removed") + "?$deltatoken=old"},
+ )
+ cursor = source._continuation(missing)
+ for _ in range(2):
+ cursor = (await source.capture_page(ROOT, cursor, files=files())).continuation
+ with pytest.raises(InvalidCaptureCheckpoint, match="disappeared"):
+ await source.capture_page(ROOT, cursor, files=files())
+ unknown_folder = True
+ with pytest.raises(OutlookBoundaryError, match="not qualified"):
+ await source.capture_page(ROOT, initial(source), files=files())
+
+
+async def test_expired_delta_restarts_without_treating_other_errors_as_empty():
+ from airweave.domains.entities.canonical.page_source import InvalidCaptureCheckpoint
+ from airweave.domains.sources.exceptions import SourceError
+
+ code = "syncStateNotFound"
+
+ def graph(request):
+ if request.url.path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ return httpx.Response(400, json={"error": {"code": code, "message": "private-error"}})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(client)
+ cursor = source._continuation(
+ OutlookMailContinuation(
+ fingerprint=source.fingerprint,
+ mode="changes",
+ phase="messages",
+ folders_to_visit=(),
+ remaining_folders=("folder",),
+ folder_links={"folder": source._delta_url("folder") + "?$deltatoken=old"},
+ )
+ )
+ with pytest.raises(InvalidCaptureCheckpoint):
+ await source.capture_page(ROOT, cursor, files=files())
+ code = "ErrorInvalidRequest"
+ with pytest.raises(SourceError) as error:
+ await source.capture_page(ROOT, cursor, files=files())
+ assert "private-error" not in str(error.value)
+
+
+@pytest.mark.parametrize("selected_view", [False, True])
+async def test_virtual_folder_traverses_physical_child_but_never_owns_message_delta(selected_view):
+ requests = []
+
+ def graph(request):
+ path = request.url.path
+ requests.append(path)
+ assert request.headers["Accept"] == "application/json;odata.metadata=minimal"
+ if path == "/v1.0/me":
+ return httpx.Response(200, json={"id": "principal"})
+ if path.endswith("/mailFolders/view"):
+ return httpx.Response(200, json={"id": "view"})
+ if path.endswith("/messages/delta"):
+ assert "/view/" not in path
+ return httpx.Response(
+ 200,
+ json={
+ "value": [],
+ "@odata.deltaLink": str(request.url.copy_with(query=b"$deltatoken=end")),
+ },
+ )
+ assert request.url.params["includeHiddenFolders"] == "true"
+ if path.endswith("/mailFolders"):
+ value = [
+ {"id": "physical"},
+ {"id": "view", "@odata.type": "#microsoft.graph.mailSearchFolder"},
+ ]
+ elif path.endswith("/view/childFolders"):
+ value = [{"id": "nested-physical", "isHidden": True}]
+ else:
+ value = []
+ return httpx.Response(200, json={"value": value})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client:
+ source = await capture(
+ client,
+ config=OutlookMailConfig(
+ expected_principal_id="principal",
+ included_folders=["view"] if selected_view else [],
+ excluded_folders=[],
+ ),
+ )
+ if selected_view:
+ with pytest.raises(OutlookBoundaryError, match="search-folder selections"):
+ await source.capture_page(ROOT, initial(source), files=files())
+ return
+ terminal = await first_message_page(source, files())
+ assert terminal.final and terminal.records == ()
+ assert set(terminal.provider_checkpoint.value["folder_links"]) == {
+ "physical",
+ "nested-physical",
+ }
+ assert "/v1.0/me/mailFolders/view/childFolders" in requests
+ assert "/v1.0/me/mailFolders/view/messages/delta" not in requests
diff --git a/backend/airweave/platform/sources/tests/test_slack_history_limit.py b/backend/airweave/platform/sources/tests/test_slack_history_limit.py
new file mode 100644
index 0000000000..71f0e1b931
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_slack_history_limit.py
@@ -0,0 +1,29 @@
+"""Unusable coverage flags must never become a complete channel inventory."""
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackSource
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("flag", ["false", None])
+async def test_malformed_history_limit_cannot_authorize_completion(flag):
+ source = SlackSource(
+ auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=MagicMock()
+ )
+ source.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U1")
+ source._get = AsyncMock(return_value={"ok": True, "team_id": "T1", "user_id": "U1"})
+ await source.validate()
+ source._get = AsyncMock(return_value={"messages": [], "is_limited": flag})
+ with pytest.raises(ValueError, match="invalid page"):
+ await source.capture_page(
+ CompletedScope(record_type="message", container_id="C1"),
+ ScanContinuation(value={}),
+ files=MagicMock(),
+ )
diff --git a/backend/airweave/platform/sources/tests/test_slack_identity.py b/backend/airweave/platform/sources/tests/test_slack_identity.py
new file mode 100644
index 0000000000..de67a61a10
--- /dev/null
+++ b/backend/airweave/platform/sources/tests/test_slack_identity.py
@@ -0,0 +1,66 @@
+"""Private-channel visibility is bound to both native workspace and user."""
+
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+from pydantic import ValidationError
+
+from airweave.domains.entities.canonical.requests import CompletedScope
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.token_providers.static import StaticTokenProvider
+from airweave.platform.configs.config import SlackConfig
+from airweave.platform.sources.slack import SlackSource
+
+
+@pytest.mark.asyncio
+async def test_native_user_change_clears_attestation_and_changes_fingerprint():
+ profile = {"ok": True, "team_id": "T1", "user_id": "U1"}
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(lambda request: httpx.Response(200, json=profile))
+ ) as client:
+ source = await SlackSource.create(
+ auth=StaticTokenProvider("same-broker"),
+ logger=MagicMock(),
+ http_client=client,
+ config=SlackConfig(expected_team_id="T1", expected_user_id="U1"),
+ )
+ original = source.capture_cycle_configuration.fingerprint
+ profile["user_id"] = "U2"
+ with pytest.raises(ValueError, match="does not match"):
+ await source.validate()
+ with pytest.raises(ValueError, match="attested workspace"):
+ await source.capture_page(
+ CompletedScope(record_type="channel"), ScanContinuation(), files=MagicMock()
+ )
+ source.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U2")
+ await source.validate()
+ assert source.capture_cycle_configuration.fingerprint != original
+ profile["ok"] = "true"
+ with pytest.raises(ValueError, match="invalid native identity"):
+ await source.validate()
+ with pytest.raises(ValueError, match="attested workspace"):
+ _ = source.capture_cycle_configuration
+
+
+@pytest.mark.asyncio
+async def test_missing_pair_cannot_start_capture():
+ with pytest.raises(ValidationError, match="both team and user"):
+ SlackConfig(expected_team_id="T1")
+ async with httpx.AsyncClient(
+ transport=httpx.MockTransport(
+ lambda request: pytest.fail("Unbound capture fetched provider data")
+ )
+ ) as client:
+ source = await SlackSource.create(
+ auth=StaticTokenProvider("synthetic"),
+ logger=MagicMock(),
+ http_client=client,
+ config=SlackConfig(),
+ )
+ with pytest.raises(ValueError, match="attested workspace"):
+ _ = source.capture_cycle_configuration
+ with pytest.raises(ValueError, match="attested workspace"):
+ await source.capture_page(
+ CompletedScope(record_type="channel"), ScanContinuation(), files=MagicMock()
+ )
diff --git a/backend/airweave/platform/sources/todoist.py b/backend/airweave/platform/sources/todoist.py
index 421f974ae8..3f3af64c7b 100644
--- a/backend/airweave/platform/sources/todoist.py
+++ b/backend/airweave/platform/sources/todoist.py
@@ -24,12 +24,12 @@
TodoistTaskEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/trello.py b/backend/airweave/platform/sources/trello.py
index bce66022c8..74568d7d0a 100644
--- a/backend/airweave/platform/sources/trello.py
+++ b/backend/airweave/platform/sources/trello.py
@@ -32,12 +32,12 @@
TrelloMemberEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py
new file mode 100644
index 0000000000..c703856b21
--- /dev/null
+++ b/backend/airweave/platform/sources/wispr.py
@@ -0,0 +1,539 @@
+"""Capture Wispr meeting and scratchpad representations through account-bound sessions."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from datetime import datetime, timezone
+from typing import Literal
+from urllib.parse import quote
+
+import httpx
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool
+
+from airweave.core.logging import ContextualLogger
+from airweave.domains.entities.canonical.cycle_models import CycleConfiguration
+from airweave.domains.entities.canonical.models import SourceRecord
+from airweave.domains.entities.canonical.page_source import CapturePage
+from airweave.domains.entities.canonical.requests import (
+ CaptureRecord,
+ CompletedScope,
+ RecordIdentity,
+)
+from airweave.domains.entities.canonical.scan_models import ScanContinuation
+from airweave.domains.sources.exceptions import SourceError, SourceServerError
+from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider
+from airweave.domains.storage import FileSkippedException
+from airweave.domains.storage.file_service import FileService
+from airweave.platform.decorators import source
+from airweave.platform.http_client.airweave_client import AirweaveHttpClient
+from airweave.platform.http_client.bounded_response import bounded_response_bytes
+from airweave.platform.http_client.composio_transport import ComposioTransport
+from airweave.platform.sources._base import BaseSource
+from airweave.schemas.source_connection import AuthenticationMethod
+
+# Decoded response and retained archive limits; overflow fails without advancing this body.
+MAX_RESPONSE_BYTES = 8 * 1024 * 1024
+MAX_ERROR_BYTES = 64 * 1024
+MAX_BODY_BYTES = 32 * 1024 * 1024
+
+
+class SessionResponse(BaseModel):
+ """Only the session identifier is needed; other routing metadata is ignored."""
+
+ session_id: str
+
+
+class ToolResponse(BaseModel):
+ """Preserve the entire native tool result, without persisting Composio log IDs."""
+
+ data: dict[str, JsonValue]
+ error: JsonValue = None
+
+
+class MeetingPage(BaseModel):
+ """Wispr's real search response pagination fields."""
+
+ model_config = ConfigDict(extra="ignore")
+ meetings: list[dict[str, JsonValue]] = Field(max_length=200)
+ has_more: StrictBool
+ next_cursor: str | None = None
+ truncated: StrictBool = False
+
+
+class ScratchpadPage(BaseModel):
+ """Scratchpad search uses notes and modification-time filtering."""
+
+ model_config = ConfigDict(extra="ignore")
+ notes: list[dict[str, JsonValue]] = Field(max_length=200)
+ has_more: StrictBool
+ next_cursor: str | None = None
+ truncated: StrictBool = False
+
+
+class _Listing(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ id: str = Field(min_length=1)
+
+
+class _Version(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ modified_at: AwareDatetime | None = None
+
+
+class _Window(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ since: AwareDatetime | None = None
+ until: AwareDatetime | None = None
+ depth: int = Field(default=0, ge=0, le=32)
+
+
+class _MeetingStart(BaseModel):
+ model_config = ConfigDict(extra="ignore")
+ start: AwareDatetime | None = None
+
+
+class _Progress(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ cursor: str | None = Field(default=None, min_length=1, max_length=8192)
+ seen: list[str] = Field(default_factory=list, max_length=1000)
+ count: int = Field(default=0, ge=0, le=1000)
+ window: _Window = Field(default_factory=_Window)
+ pending: list[_Window] = Field(default_factory=list, max_length=33)
+ dates_complete: bool = True
+ earliest: AwareDatetime | None = None
+ latest: AwareDatetime | None = None
+
+
+@source(
+ name="Wispr",
+ short_name="wispr",
+ auth_methods=[AuthenticationMethod.AUTH_PROVIDER],
+ labels=["Meetings"],
+ supports_continuous=False,
+)
+class WisprSource(BaseSource):
+ """Capture available notes, summaries and transcript ranges; never infer deletions.
+
+ The provider omits its raw editor JSON and has no deletion feed or snapshot cursor.
+ Captures remain explicitly partial, even after every exposed text range is fetched.
+ """
+
+ canonical_record_types = ("meeting_listing", "meeting", "scratchpad_listing", "scratchpad_note")
+ canonical_container_parents = {
+ "meeting": "meeting_listing",
+ "scratchpad_note": "scratchpad_listing",
+ }
+
+ @property
+ def capture_cycle_configuration(self) -> CycleConfiguration:
+ """Bind durable discovery policy to this account and representation version."""
+ return self._cycle_configuration
+
+ @classmethod
+ async def create(
+ cls,
+ *,
+ auth: ManagedToolAuthProvider,
+ logger: ContextualLogger,
+ http_client: AirweaveHttpClient,
+ config: BaseModel | None = None,
+ ) -> WisprSource:
+ """Require an explicit tool session capability, never a raw token."""
+ if not isinstance(auth, ManagedToolAuthProvider):
+ raise TypeError("Wispr requires managed tool authentication")
+ instance = cls(auth=auth, logger=logger, http_client=http_client)
+ instance._tool_auth = auth
+ instance._session_id: str | None = None
+ fingerprint = hashlib.sha256(
+ json.dumps(
+ {
+ "version": 3,
+ "account": auth.connected_account_id,
+ "kinds": cls.canonical_record_types,
+ "policy": "discovery_only",
+ "listing_limit": 200,
+ },
+ sort_keys=True,
+ ).encode()
+ ).hexdigest()
+ instance._cycle_configuration = CycleConfiguration.from_source(
+ fingerprint=fingerprint,
+ record_types=cls.canonical_record_types,
+ container_parents=cls.canonical_container_parents,
+ completion_policies={
+ "meeting_listing": "discovery_only",
+ "scratchpad_listing": "discovery_only",
+ },
+ )
+ return instance
+
+ async def _post(self, path: str, body: dict) -> dict:
+ async with self.http_client.stream(
+ "POST",
+ "https://backend.composio.dev/api/v3.1/tool_router/" + path,
+ headers={
+ "x-api-key": self._tool_auth.api_key.get_secret_value(),
+ "Accept-Encoding": "identity",
+ },
+ json=body,
+ ) as response:
+ if response.headers.get("content-encoding", "identity").lower() != "identity":
+ if not response.is_success:
+ ComposioTransport._check_proxy_response(
+ httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ request=response.request,
+ content=b"",
+ ),
+ response.request,
+ )
+ raise SourceError(
+ "Wispr did not honor bounded identity encoding", source_short_name="wispr"
+ )
+ maximum = MAX_RESPONSE_BYTES if response.is_success else MAX_ERROR_BYTES
+ try:
+ raw = await bounded_response_bytes(response, maximum, label="Wispr response")
+ except FileSkippedException as error:
+ # Preserve status-based authentication/throttle failures without retaining an
+ # oversized error body. Successful overflow is never a storage skip.
+ if not response.is_success:
+ ComposioTransport._check_proxy_response(
+ httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ request=response.request,
+ content=b"",
+ ),
+ response.request,
+ )
+ raise SourceError(
+ "Wispr response exceeds byte limit; capture is incomplete",
+ source_short_name="wispr",
+ ) from error
+ buffered = httpx.Response(
+ response.status_code,
+ headers=response.headers,
+ request=response.request,
+ content=raw,
+ )
+ ComposioTransport._check_proxy_response(buffered, response.request)
+ value = buffered.json()
+ if not isinstance(value, dict):
+ raise SourceError("Wispr response must be an object", source_short_name="wispr")
+ return value
+
+ async def _execute(self, slug: str, arguments: dict) -> dict[str, JsonValue]:
+ if self._session_id is None:
+ session = SessionResponse.model_validate(
+ await self._post(
+ "session",
+ {
+ "user_id": self._tool_auth.user_id,
+ "connected_accounts": {
+ "wispr_flow_mcp": [self._tool_auth.connected_account_id]
+ },
+ "toolkits": {"enable": ["wispr_flow_mcp"]},
+ "manage_connections": {"enable": False},
+ },
+ )
+ )
+ self._session_id = session.session_id
+ result = ToolResponse.model_validate(
+ await self._post(
+ "session/" + quote(self._session_id, safe="") + "/execute",
+ {
+ "tool_slug": slug,
+ "arguments": arguments,
+ },
+ )
+ )
+ if isinstance(result.error, str) and any(
+ phrase in result.error.lower() for phrase in ("rate limit", "too many requests")
+ ):
+ # HTTP 200 tool text is a diagnostic signal, not a structured 429 or
+ # a provider Retry-After contract. Keep the failure terminal here.
+ raise SourceServerError(
+ "Managed Wispr tool reported a rate-limit signal; "
+ "the limiting service and retry delay are unknown. Capture is incomplete.",
+ source_short_name="wispr",
+ )
+ if result.error is not None:
+ raise ValueError("Wispr tool execution failed; capture is incomplete")
+ return result.data
+
+ async def validate(self) -> None:
+ """Verify the bound session can discover both supported Wispr inventories."""
+ MeetingPage.model_validate(
+ await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", {"limit": 1})
+ )
+ ScratchpadPage.model_validate(
+ await self._execute("WISPR_FLOW_MCP_SEARCH_SCRATCHPAD_NOTES", {"limit": 1})
+ )
+
+ @staticmethod
+ def _next_offset(text: JsonValue, field: str) -> int | None:
+ if not isinstance(text, str):
+ raise ValueError("Wispr requested text is not a string")
+ matches = list(
+ re.finditer(
+ rf"\(\.\.\.truncated, \d+ chars remaining; continue with "
+ rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$",
+ text,
+ re.MULTILINE,
+ )
+ )
+ if len(matches) > 1:
+ raise ValueError("Wispr returned ambiguous continuation markers")
+ return int(matches[0].group(1)) if matches else None
+
+ @staticmethod
+ def _require_range_version(version: JsonValue) -> None:
+ """Missing versions cannot establish consistency across separate native reads."""
+ if not isinstance(version, str) or not version.strip():
+ raise ValueError("Wispr paginated body has no usable version; capture is incomplete")
+
+ @staticmethod
+ def _archive_size(value: JsonValue, previous: int = 0) -> int:
+ """Count complete native JSON and range descriptors once per acquired part."""
+ size = previous + len(
+ json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+ )
+ if size > MAX_BODY_BYTES:
+ raise SourceError(
+ "Wispr body exceeds byte limit; capture is incomplete", source_short_name="wispr"
+ )
+ return size
+
+ async def _body(
+ self, kind: Literal["meeting", "scratchpad_note"], identity: str
+ ) -> dict[str, JsonValue]:
+ """Retain complete returned ranges for the two concrete Wispr body tools."""
+ if kind == "meeting":
+ slug, id_field, offsets = (
+ "WISPR_FLOW_MCP_GET_MEETING",
+ "meeting_id",
+ {"content": 0, "transcript": 0},
+ )
+ else:
+ slug, id_field, offsets = (
+ "WISPR_FLOW_MCP_GET_SCRATCHPAD_NOTE",
+ "note_id",
+ {"content": 0},
+ )
+ responses: list[JsonValue] = []
+ archive_bytes = len(b'{"responses":[]}')
+ for _ in range(100):
+ arguments = {
+ id_field: identity,
+ **{
+ "view_" + field: {"char_limit": 40000, "start_char": offset}
+ for field, offset in offsets.items()
+ },
+ }
+ response = await self._execute(slug, arguments)
+ if response.get("id") != identity:
+ raise ValueError("Wispr returned a different body identity")
+ if responses and response.get("modified_at") != responses[0]["response"].get(
+ "modified_at"
+ ):
+ raise ValueError("Wispr body changed during paginated capture")
+ retained = {"requested_ranges": arguments, "response": response}
+ archive_bytes = self._archive_size(retained, archive_bytes + bool(responses))
+ responses.append(retained)
+ following = {}
+ for field, offset in offsets.items():
+ if (
+ field == "transcript"
+ and offset == 0
+ and response.get("has_transcript") is False
+ and "transcript" in response
+ and response["transcript"] is None
+ ):
+ # Explicit absence is not a failed range; retain native null unchanged.
+ continue
+ next_offset = self._next_offset(response.get(field), field)
+ if next_offset is not None:
+ if next_offset <= offset:
+ raise ValueError("Wispr continuation made no progress")
+ following[field] = next_offset
+ if not following:
+ return {"responses": responses}
+ self._require_range_version(response.get("modified_at"))
+ # Request only continuing ranges; complete ranges already remain in responses.
+ offsets = following
+ raise ValueError("Wispr body exceeds the bounded capture page limit")
+
+ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope:
+ """Keep each body attached to its exact native listing observation."""
+ if self.canonical_container_parents.get(record_type) != parent.identity.record_type:
+ raise ValueError("Unsupported Wispr child scope")
+ # Preserve existing provider-global body identity, including null container.
+ return CompletedScope(record_type=record_type, parent=parent.identity)
+
+ async def confirm_absent(self, record: SourceRecord) -> None:
+ """Wispr supplies no audited absence or revocation proof."""
+ raise ValueError("Wispr listing omission is not deletion or revocation evidence")
+
+ async def capture_page(
+ self,
+ scope: CompletedScope,
+ continuation: ScanContinuation,
+ *,
+ files: FileService,
+ parent: SourceRecord | None = None,
+ ) -> CapturePage:
+ """Commit listing pages and individual bodies through one durable frontier.
+
+ Completed bodies remain observed state in this cycle even if a later attempt
+ refreshes listing metadata. The next cycle refreshes them; this is no snapshot.
+ """
+ kind = scope.record_type
+ if (
+ kind in {"meeting_listing", "scratchpad_listing"}
+ and scope.parent is None
+ and parent is None
+ and scope.container_id is None
+ ):
+ return await self._listing_page(kind, continuation)
+ if (
+ kind not in {"meeting", "scratchpad_note"}
+ or parent is None
+ or self.canonical_container_parents.get(kind) != parent.identity.record_type
+ or scope.parent != parent.identity
+ or scope.container_id is not None
+ or continuation.value
+ ):
+ raise ValueError("Unsupported Wispr body scope or continuation")
+ identity = _Listing.model_validate(parent.payload).id
+ if identity != parent.identity.native_id:
+ raise ValueError("Wispr listing identity does not match its scope")
+ payload = await self._body(kind, identity)
+ payload["listing"] = parent.payload
+ self._archive_size(payload)
+ modified = _Version.model_validate(payload["responses"][0]["response"]).modified_at
+ return CapturePage(
+ records=(
+ CaptureRecord(
+ identity=RecordIdentity(record_type=kind, native_id=identity),
+ parent=parent.identity,
+ payload=payload,
+ completeness="partial",
+ source_updated_at=modified,
+ observed_at=datetime.now(timezone.utc),
+ ),
+ ),
+ continuation=ScanContinuation(),
+ final=True,
+ )
+
+ async def _listing_page(
+ self, kind: Literal["meeting_listing", "scratchpad_listing"], continuation: ScanContinuation
+ ) -> CapturePage:
+ progress = _Progress.model_validate(continuation.value)
+ arguments: dict[str, JsonValue] = {"limit": 200}
+ if progress.cursor:
+ arguments["cursor"] = progress.cursor
+ if progress.window.since:
+ arguments["since"] = progress.window.since.isoformat()
+ if progress.window.until:
+ arguments["until"] = progress.window.until.isoformat()
+ if kind == "meeting_listing":
+ page = MeetingPage.model_validate(
+ await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", arguments)
+ )
+ rows = page.meetings
+ dates = [_MeetingStart.model_validate(row).start for row in rows]
+ else:
+ page = ScratchpadPage.model_validate(
+ await self._execute("WISPR_FLOW_MCP_SEARCH_SCRATCHPAD_NOTES", arguments)
+ )
+ rows = page.notes
+ dates = [_Version.model_validate(row).modified_at for row in rows]
+ next_progress, final = self._advance_listing(progress, page, dates)
+ identities = [_Listing.model_validate(row).id for row in rows]
+ if len(identities) != len(set(identities)):
+ raise ValueError("Wispr listing page repeats a native identity")
+ records = tuple(
+ CaptureRecord(
+ identity=RecordIdentity(record_type=kind, native_id=identity),
+ payload=row,
+ completeness="metadata_only",
+ observed_at=datetime.now(timezone.utc),
+ )
+ for identity, row in zip(identities, rows, strict=True)
+ )
+ return CapturePage(
+ records=records,
+ continuation=ScanContinuation(value=next_progress.model_dump(mode="json")),
+ final=final,
+ )
+
+ def _advance_listing(
+ self,
+ progress: _Progress,
+ page: MeetingPage | ScratchpadPage,
+ dates: list[datetime | None],
+ ) -> tuple[_Progress, bool]:
+ """Advance native cursors and capped date windows without a completeness claim."""
+ count = progress.count + len(dates)
+ following = page.next_cursor if page.has_more else None
+ dates_complete = progress.dates_complete and all(value is not None for value in dates)
+ known_dates = [value for value in dates if value is not None] + [
+ d for d in (progress.earliest, progress.latest) if d is not None
+ ]
+ earliest = min(known_dates) if known_dates else None
+ latest = max(known_dates) if known_dates else None
+ capped = page.truncated or count > 1000 or (page.has_more and count >= 1000)
+ if capped:
+ if not dates_complete:
+ raise ValueError("Wispr query cap cannot be partitioned safely")
+ next_progress = self._partition(progress, earliest, latest)
+ final = False
+ elif page.has_more:
+ following_digest = (
+ hashlib.sha256(following.encode()).hexdigest()[:32] if following else None
+ )
+ if not following or following == progress.cursor or following_digest in progress.seen:
+ raise ValueError("Wispr returned incomplete or repeated pagination")
+ next_progress = _Progress.model_validate(
+ {
+ **progress.model_dump(),
+ "dates_complete": dates_complete,
+ "cursor": following,
+ "count": count,
+ "earliest": earliest,
+ "latest": latest,
+ "seen": [*progress.seen, following_digest],
+ }
+ )
+ final = False
+ elif progress.pending:
+ next_progress = _Progress(window=progress.pending[-1], pending=progress.pending[:-1])
+ final = False
+ else:
+ next_progress = _Progress()
+ final = True
+ return next_progress, final
+
+ @staticmethod
+ def _partition(
+ progress: _Progress, earliest: datetime | None, latest: datetime | None
+ ) -> _Progress:
+ if earliest is None or latest is None or earliest == latest or progress.window.depth >= 32:
+ raise ValueError("Wispr query cap cannot be partitioned safely")
+ split = earliest + (latest - earliest) / 2
+ window = progress.window
+ if (window.since is not None and split <= window.since) or (
+ window.until is not None and split >= window.until
+ ):
+ raise ValueError("Wispr date partition made no progress")
+ return _Progress(
+ window=_Window(since=window.since, until=split, depth=window.depth + 1),
+ pending=[
+ *progress.pending,
+ _Window(since=split, until=window.until, depth=window.depth + 1),
+ ],
+ )
diff --git a/backend/airweave/platform/sources/word.py b/backend/airweave/platform/sources/word.py
index 8a4178b66d..db7abc96ed 100644
--- a/backend/airweave/platform/sources/word.py
+++ b/backend/airweave/platform/sources/word.py
@@ -29,12 +29,12 @@
from airweave.platform.entities._base import BaseEntity
from airweave.platform.entities.word import WORD_EXTENSIONS, WordDocumentEntity
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/zendesk.py b/backend/airweave/platform/sources/zendesk.py
index 2f9ffafd58..329836bd75 100644
--- a/backend/airweave/platform/sources/zendesk.py
+++ b/backend/airweave/platform/sources/zendesk.py
@@ -25,12 +25,12 @@
_parse_dt,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/platform/sources/zoho_crm.py b/backend/airweave/platform/sources/zoho_crm.py
index 70376e0e0b..6c6982d4ca 100644
--- a/backend/airweave/platform/sources/zoho_crm.py
+++ b/backend/airweave/platform/sources/zoho_crm.py
@@ -37,12 +37,12 @@
ZohoCRMSalesOrderEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
ZOHO_API_LIMIT = 200
diff --git a/backend/airweave/platform/sources/zoom.py b/backend/airweave/platform/sources/zoom.py
index fa2004f3cc..6fb7fef55f 100644
--- a/backend/airweave/platform/sources/zoom.py
+++ b/backend/airweave/platform/sources/zoom.py
@@ -37,12 +37,12 @@
ZoomTranscriptEntity,
)
from airweave.platform.http_client.airweave_client import AirweaveHttpClient
-from airweave.platform.sources._base import BaseSource
-from airweave.platform.sources.http_helpers import raise_for_status
-from airweave.platform.sources.retry_helpers import (
+from airweave.platform.http_client.retry_helpers import (
retry_if_rate_limit_or_timeout,
wait_rate_limit_with_backoff,
)
+from airweave.platform.sources._base import BaseSource
+from airweave.platform.sources.http_helpers import raise_for_status
from airweave.schemas.source_connection import AuthenticationMethod, OAuthType
diff --git a/backend/airweave/schemas/entity.py b/backend/airweave/schemas/entity.py
index 1ec1e6d0bf..82e8809fa2 100644
--- a/backend/airweave/schemas/entity.py
+++ b/backend/airweave/schemas/entity.py
@@ -38,6 +38,7 @@ class EntityUpdate(BaseModel):
class EntityInDBBase(EntityBase):
"""Base schema for Entity stored in DB."""
+ sync_job_id: UUID | None
id: UUID
organization_id: UUID
created_at: datetime
diff --git a/backend/airweave/schemas/search_v2.py b/backend/airweave/schemas/search_v2.py
index 9eaa2d9537..8691bd48f5 100644
--- a/backend/airweave/schemas/search_v2.py
+++ b/backend/airweave/schemas/search_v2.py
@@ -42,7 +42,7 @@ class InstantSearchRequest(BaseModel):
filter: Optional[list[FilterGroup]] = Field(
default=None, description="Filter groups (combined with OR)."
)
- limit: int = Field(default=100, ge=1, le=1000, description="Max results to return.")
+ limit: int = Field(default=100, ge=1, le=200, description="Max results to return.")
offset: int = Field(default=0, ge=0, description="Number of results to skip.")
@field_validator("query")
@@ -70,7 +70,7 @@ class ClassicSearchRequest(BaseModel):
filter: Optional[list[FilterGroup]] = Field(
default=None, description="Filter groups (combined with OR)."
)
- limit: int = Field(default=100, ge=1, le=1000, description="Max results to return.")
+ limit: int = Field(default=100, ge=1, le=200, description="Max results to return.")
offset: int = Field(default=0, ge=0, description="Number of results to skip.")
@field_validator("query")
@@ -185,7 +185,10 @@ class BrowseResponse(BaseModel):
results: list[SearchResult] = Field(
default_factory=list, description="Rows on this page (one per source entity)."
)
- total: int = Field(..., description="Total entity count matching the filter.")
+ total: int | None = Field(
+ ..., description="Exact visible count, or null when the index cannot prove it."
+ )
+ retrieval_incomplete: bool = False
limit: int = Field(..., description="Limit echoed back from the request.")
offset: int = Field(..., description="Offset echoed back from the request.")
@@ -286,3 +289,11 @@ class SearchV2Response(BaseModel):
results: list[SearchResult] = Field(
default_factory=list, description="Search results ordered by relevance."
)
+
+ engine_partial: bool = False
+ engine_coverage_percent: float | None = None
+
+ retrieval_incomplete: bool = Field(
+ default=False, description="Unavailable index candidates may leave this page incomplete."
+ )
+ excluded_candidates: int = Field(default=0, ge=0)
diff --git a/backend/airweave/schemas/sync.py b/backend/airweave/schemas/sync.py
index 8738aa362d..ec8ec82149 100644
--- a/backend/airweave/schemas/sync.py
+++ b/backend/airweave/schemas/sync.py
@@ -126,6 +126,8 @@ class ScheduleResponse(BaseModel):
class SyncInDBBase(SyncBase):
"""Base schema for Sync stored in DB."""
+ provisioning_generation: int = Field(default=0, ge=0)
+ provisioning_ready_generation: int = Field(default=0, ge=0)
id: UUID
organization_id: UUID
created_at: datetime
diff --git a/backend/airweave/schemas/sync_job.py b/backend/airweave/schemas/sync_job.py
index 1d5b3c16a4..fd632031dd 100644
--- a/backend/airweave/schemas/sync_job.py
+++ b/backend/airweave/schemas/sync_job.py
@@ -68,6 +68,7 @@ class SyncJobUpdate(BaseModel):
class SyncJobInDBBase(SyncJobBase):
"""Base schema for SyncJob stored in DB."""
+ provisioning_generation: int = Field(default=0, ge=0)
id: UUID
organization_id: UUID
created_by_email: Optional[EmailStr] = None
diff --git a/backend/airweave/search/operations/embed_query.py b/backend/airweave/search/operations/embed_query.py
index 15b2b663e5..a3a1d76c9d 100644
--- a/backend/airweave/search/operations/embed_query.py
+++ b/backend/airweave/search/operations/embed_query.py
@@ -110,7 +110,7 @@ async def _generate_dense_embeddings(
f"[EmbedQuery] Generating {self.dense_embedder.dimensions}-dim embeddings "
f"for {len(queries)} queries"
)
- results = await self.dense_embedder.embed_many(queries)
+ results = await self.dense_embedder.embed_many(queries, purpose="query")
# Validate we got embeddings for all queries
if len(results) != len(queries):
diff --git a/backend/alembic/versions/0001_canonical_records.py b/backend/alembic/versions/0001_canonical_records.py
new file mode 100644
index 0000000000..c439d2ff7b
--- /dev/null
+++ b/backend/alembic/versions/0001_canonical_records.py
@@ -0,0 +1,107 @@
+"""Canonical source records, fenced writers and observed change snapshots.
+
+Revision ID: 0001
+Revises: 0000
+"""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0001"
+down_revision = "0000"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Legacy metadata remains revision zero until authoritative recapture."""
+ op.drop_constraint("fk_entity_sync_job_id", "entity", type_="foreignkey")
+ op.alter_column("entity", "sync_job_id", nullable=True)
+ op.create_foreign_key(
+ "fk_entity_sync_job_id",
+ "entity",
+ "sync_job",
+ ["sync_job_id"],
+ ["id"],
+ ondelete="SET NULL",
+ )
+ for column in (
+ sa.Column("native_id", sa.String(), nullable=True),
+ sa.Column("container_id", sa.String(), nullable=True),
+ sa.Column("parent_record_type", sa.String(), nullable=True),
+ sa.Column("parent_native_id", sa.String(), nullable=True),
+ sa.Column("parent_container_id", sa.String(), nullable=True),
+ sa.Column("source_payload", postgresql.JSONB(), nullable=True),
+ sa.Column("payload_schema_version", sa.Integer(), nullable=False, server_default="1"),
+ sa.Column("record_revision", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("capture_hash", sa.String(), nullable=True),
+ sa.Column("content_hash", sa.String(), nullable=True),
+ sa.Column("source_created_at", sa.DateTime(timezone=True), nullable=True),
+ sa.Column("source_updated_at", sa.DateTime(timezone=True), nullable=True),
+ sa.Column("observed_at", sa.DateTime(timezone=True), nullable=True),
+ sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
+ sa.Column("removal_reason", sa.String(), nullable=True),
+ sa.Column("completeness", sa.String(), nullable=True),
+ sa.Column("blob_references", postgresql.JSONB(), nullable=True),
+ sa.Column("last_seen_run_id", sa.UUID(), nullable=True),
+ sa.Column("indexed_revision", sa.BigInteger(), nullable=True),
+ sa.Column("indexed_pipeline_version", sa.BigInteger(), nullable=True),
+ sa.Column("projection_error", sa.String(), nullable=True),
+ ):
+ op.add_column("entity", column)
+ op.create_index(
+ "idx_entity_canonical_scope",
+ "entity",
+ ["sync_id", "entity_definition_short_name", "container_id"],
+ )
+ op.create_index(
+ "idx_entity_pending_revision",
+ "entity",
+ ["sync_id", "id"],
+ postgresql_where=sa.text(
+ "record_revision > 0 AND indexed_revision IS DISTINCT FROM record_revision"
+ ),
+ )
+ for column in (
+ sa.Column("observed_change_sequence", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("writer_epoch", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("writer_job_id", sa.UUID(), nullable=True),
+ sa.Column("writer_attempt_id", sa.UUID(), nullable=True),
+ sa.Column("writer_attempt_number", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("index_pipeline_version", sa.BigInteger(), nullable=False, server_default="1"),
+ ):
+ op.add_column("sync", column)
+ op.create_table(
+ "entity_change",
+ sa.Column("id", sa.UUID(), nullable=False, primary_key=True),
+ sa.Column("created_at", sa.DateTime(), nullable=False),
+ sa.Column("modified_at", sa.DateTime(), nullable=False),
+ sa.Column(
+ "organization_id",
+ sa.UUID(),
+ sa.ForeignKey("organization.id", ondelete="CASCADE"),
+ nullable=False,
+ ),
+ sa.Column(
+ "sync_id", sa.UUID(), sa.ForeignKey("sync.id", ondelete="CASCADE"), nullable=False
+ ),
+ sa.Column(
+ "entity_record_id",
+ sa.UUID(),
+ sa.ForeignKey("entity.id", ondelete="CASCADE"),
+ nullable=False,
+ ),
+ sa.Column("sequence", sa.BigInteger(), nullable=False),
+ sa.Column("record_revision", sa.BigInteger(), nullable=False),
+ sa.Column("kind", sa.String(), nullable=False),
+ sa.Column("snapshot", postgresql.JSONB(), nullable=False),
+ sa.UniqueConstraint("sync_id", "sequence", name="uq_entity_change_sequence"),
+ )
+ op.create_index("ix_entity_change_organization_id", "entity_change", ["organization_id"])
+
+
+def downgrade():
+ """Refuse to destroy canonical source state through an accidental downgrade."""
+ raise RuntimeError("Canonical record migration is forward-only; restore a verified backup")
diff --git a/backend/alembic/versions/0002_projection_publication.py b/backend/alembic/versions/0002_projection_publication.py
new file mode 100644
index 0000000000..44025b9664
--- /dev/null
+++ b/backend/alembic/versions/0002_projection_publication.py
@@ -0,0 +1,26 @@
+"""Revision-safe canonical index publication pointer.
+
+Revision ID: 0002
+Revises: 0001
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "0002"
+down_revision = "0001"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Add a publication pointer; old index state remains untrusted until rebuilt."""
+ op.add_column("entity", sa.Column("indexed_generation", sa.Uuid(), nullable=True))
+ op.add_column("entity", sa.Column("indexed_chunk_count", sa.BigInteger(), nullable=True))
+
+
+def downgrade():
+ """Remove only derived publication state, preserving source records."""
+ op.drop_column("entity", "indexed_chunk_count")
+ op.drop_column("entity", "indexed_generation")
diff --git a/backend/alembic/versions/0003_mail_thread_index.py b/backend/alembic/versions/0003_mail_thread_index.py
new file mode 100644
index 0000000000..09eb966ba6
--- /dev/null
+++ b/backend/alembic/versions/0003_mail_thread_index.py
@@ -0,0 +1,24 @@
+"""Index account-scoped Gmail thread traversal over original provider JSON."""
+
+from alembic import op
+
+revision = "0003"
+down_revision = "0002"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ """Keep thread lookup proportional to one account/thread, not its entire mailbox."""
+ op.execute("""
+ CREATE INDEX ix_entity_mail_thread ON entity
+ (organization_id, sync_id, (source_payload ->> 'threadId'),
+ source_created_at ASC NULLS LAST, id)
+ WHERE entity_definition_short_name = 'message'
+ AND record_revision > 0 AND deleted_at IS NULL
+ """)
+
+
+def downgrade() -> None:
+ """Drop only the derived lookup index; canonical payloads remain unchanged."""
+ op.drop_index("ix_entity_mail_thread", table_name="entity")
diff --git a/backend/alembic/versions/0004_projection_generation.py b/backend/alembic/versions/0004_projection_generation.py
new file mode 100644
index 0000000000..af1d80c850
--- /dev/null
+++ b/backend/alembic/versions/0004_projection_generation.py
@@ -0,0 +1,47 @@
+"""Durable prepared feed manifests and repeatable retired-generation cleanup.
+
+Revision ID: 0004
+Revises: 0003
+"""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0004"
+down_revision = "0003"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Keep body-free ownership IDs without cascading FKs so deletion survives source removal."""
+ op.create_table(
+ "projection_generation",
+ sa.Column("id", sa.Uuid(), primary_key=True),
+ sa.Column("created_at", sa.DateTime(), nullable=False),
+ sa.Column("modified_at", sa.DateTime(), nullable=False),
+ *(
+ sa.Column(name, sa.Uuid(), nullable=False)
+ for name in ("organization_id", "sync_id", "collection_id", "record_id")
+ ),
+ sa.Column("revision", sa.BigInteger(), nullable=False),
+ sa.Column("pipeline_version", sa.BigInteger(), nullable=False),
+ sa.Column("documents", postgresql.JSONB(), nullable=False),
+ sa.Column("retired_at", sa.DateTime(timezone=True)),
+ sa.Column("next_gc_at", sa.DateTime(timezone=True), nullable=False),
+ sa.Column("last_gc_at", sa.DateTime(timezone=True)),
+ sa.Column("delete_cursor", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("gc_passes", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column("gc_attempt", sa.Uuid()),
+ sa.Column("gc_error", sa.String(200)),
+ )
+ for name in ("organization_id", "sync_id", "record_id", "next_gc_at"):
+ op.create_index(f"ix_projection_generation_{name}", "projection_generation", [name])
+ op.create_index("idx_projection_gc_due", "projection_generation", ["next_gc_at", "id"])
+
+
+def downgrade():
+ """Do not silently discard cleanup obligations."""
+ raise RuntimeError("Projection generation manifests require explicit cleanup before removal")
diff --git a/backend/alembic/versions/0005_capture_scan.py b/backend/alembic/versions/0005_capture_scan.py
new file mode 100644
index 0000000000..d87af234f1
--- /dev/null
+++ b/backend/alembic/versions/0005_capture_scan.py
@@ -0,0 +1,56 @@
+"""Durable whole-scope capture progress.
+
+Revision ID: 0005
+Revises: 0004
+"""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0005"
+down_revision = "0004"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Keep one bounded continuation per source scope, not per job."""
+ op.create_table(
+ "capture_scan",
+ sa.Column("id", sa.Uuid(), primary_key=True),
+ sa.Column("created_at", sa.DateTime(), nullable=False),
+ sa.Column("modified_at", sa.DateTime(), nullable=False),
+ sa.Column(
+ "organization_id",
+ sa.Uuid(),
+ sa.ForeignKey("organization.id", ondelete="CASCADE"),
+ nullable=False,
+ ),
+ sa.Column(
+ "sync_id", sa.Uuid(), sa.ForeignKey("sync.id", ondelete="CASCADE"), nullable=False
+ ),
+ sa.Column("scope_key", sa.String(), nullable=False),
+ sa.Column("record_type", sa.String(), nullable=False),
+ sa.Column("container_id", sa.String()),
+ sa.Column("cycle_id", sa.Uuid(), nullable=False),
+ sa.Column("sweep_id", sa.Uuid(), nullable=False),
+ sa.Column("revision", sa.BigInteger(), nullable=False),
+ sa.Column("phase", sa.String(), nullable=False),
+ sa.Column("fingerprint", sa.String(64), nullable=False),
+ sa.Column("continuation", postgresql.JSONB(), nullable=False),
+ sa.Column("started_at", sa.DateTime(timezone=True), nullable=False),
+ sa.Column("completed_at", sa.DateTime(timezone=True)),
+ sa.UniqueConstraint("sync_id", "scope_key", name="uq_capture_scan_scope"),
+ sa.CheckConstraint(
+ "phase IN ('collecting','reconciling','complete')", name="ck_capture_scan_phase"
+ ),
+ sa.CheckConstraint("revision >= 1", name="ck_capture_scan_revision"),
+ )
+ op.create_index("ix_capture_scan_organization_id", "capture_scan", ["organization_id"])
+
+
+def downgrade():
+ """Active recovery obligations require an explicit operator decision."""
+ raise RuntimeError("Capture scans require explicit completion or abandonment before removal")
diff --git a/backend/alembic/versions/0006_record_visibility.py b/backend/alembic/versions/0006_record_visibility.py
new file mode 100644
index 0000000000..9399238520
--- /dev/null
+++ b/backend/alembic/versions/0006_record_visibility.py
@@ -0,0 +1,58 @@
+"""Epoch-attested canonical record-parent visibility.
+
+Revision ID: 0006
+Revises: 0005
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "0006"
+down_revision = "0005"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Attest only existing visible flat links; never revive withdrawn content."""
+ op.add_column(
+ "entity", sa.Column("visibility_epoch", sa.BigInteger(), nullable=False, server_default="1")
+ )
+ op.add_column("entity", sa.Column("parent_visibility_epoch", sa.BigInteger()))
+ op.create_check_constraint(
+ "ck_entity_visibility_epoch",
+ "entity",
+ "visibility_epoch >= 1 AND "
+ "(parent_visibility_epoch IS NULL OR parent_visibility_epoch >= 1)",
+ )
+ op.create_index(
+ "idx_entity_canonical_identity",
+ "entity",
+ ["sync_id", "entity_definition_short_name", "native_id", "container_id"],
+ )
+ op.create_index(
+ "idx_entity_canonical_parent",
+ "entity",
+ ["sync_id", "parent_record_type", "parent_native_id", "parent_container_id"],
+ )
+ op.execute("""
+ UPDATE entity child SET parent_visibility_epoch = parent.visibility_epoch
+ FROM entity parent
+ WHERE child.organization_id = parent.organization_id
+ AND child.sync_id = parent.sync_id
+ AND child.parent_record_type = parent.entity_definition_short_name
+ AND child.parent_native_id = parent.native_id
+ AND child.parent_container_id IS NOT DISTINCT FROM parent.container_id
+ AND child.record_revision > 0 AND parent.record_revision > 0
+ AND parent.parent_record_type IS NULL AND parent.deleted_at IS NULL
+ AND (parent.removal_reason IS NULL
+ OR parent.removal_reason NOT IN ('scope_removed', 'access_revoked'))
+ AND (child.removal_reason IS NULL
+ OR child.removal_reason NOT IN ('scope_removed', 'access_revoked'))
+ """)
+
+
+def downgrade():
+ """Removing visibility attestations requires an explicit operator decision."""
+ raise RuntimeError("Record visibility cannot safely downgrade to flat parent checks")
diff --git a/backend/alembic/versions/0007_scan_scope_owner.py b/backend/alembic/versions/0007_scan_scope_owner.py
new file mode 100644
index 0000000000..6c2b772e23
--- /dev/null
+++ b/backend/alembic/versions/0007_scan_scope_owner.py
@@ -0,0 +1,171 @@
+"""Exact parent ownership for nested capture scopes.
+
+Revision ID: 0007
+Revises: 0006
+"""
+
+import json
+from typing import Annotated, Literal
+from uuid import UUID
+
+import sqlalchemy as sa
+from pydantic import (
+ BaseModel,
+ ConfigDict,
+ Field,
+ StrictInt,
+ StrictStr,
+ ValidationError,
+ model_validator,
+)
+
+from alembic import op
+
+revision = "0007"
+down_revision = "0006"
+branch_labels = None
+depends_on = None
+
+Kind = Annotated[StrictStr, Field(min_length=1, max_length=200)]
+
+
+class LegacyConfiguration(BaseModel):
+ """Frozen schema1 migration contract, independent of future runtime models."""
+
+ model_config = ConfigDict(extra="forbid")
+ fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$")
+ root_record_type: Kind
+ child_record_types: list[Kind] = Field(default_factory=list, max_length=20)
+
+ @model_validator(mode="after")
+ def distinct(self):
+ """A flat root cannot be its own child or have duplicate child declarations."""
+ if self.root_record_type in self.child_record_types or len(
+ set(self.child_record_types)
+ ) != len(self.child_record_types):
+ raise ValueError("Invalid legacy flat topology")
+ return self
+
+
+class LegacyVersion(BaseModel):
+ """Exact existing checkpoint boundary."""
+
+ model_config = ConfigDict(extra="forbid")
+ cycle_id: UUID
+ revision: StrictInt = Field(ge=1)
+
+
+class LegacyCycle(BaseModel):
+ """Only active schema1 snapshots authorize ownership backfill."""
+
+ model_config = ConfigDict(extra="forbid")
+ schema_version: StrictInt = Field(ge=1, le=1)
+ version: LegacyVersion
+ configuration: LegacyConfiguration
+ phase: Literal["active"]
+ root_writer_attempt_id: UUID | None = None
+ completed_job_id: UUID | None = None
+
+
+def bind_legacy_scans(connection):
+ """Unresolved or malformed rows retain their original bytes for explicit recovery."""
+ cursors = connection.execute(
+ sa.text("SELECT organization_id,sync_id,cursor_data FROM sync_cursor")
+ ).mappings()
+ for cursor in cursors:
+ value = cursor["cursor_data"]
+ if not isinstance(value, dict):
+ continue
+ try:
+ cycle = LegacyCycle.model_validate(value.get("canonical_cycle"))
+ except ValidationError:
+ continue
+ scope = {
+ "org": cursor["organization_id"],
+ "sync": cursor["sync_id"],
+ "cycle": cycle.version.cycle_id,
+ }
+ rows = (
+ connection.execute(
+ sa.text("""
+ SELECT id,scope_key,record_type,container_id FROM capture_scan
+ WHERE organization_id=:org AND sync_id=:sync AND cycle_id=:cycle
+ """),
+ scope,
+ )
+ .mappings()
+ .all()
+ )
+ for row in rows:
+ expected = json.dumps(
+ [row["record_type"], row["container_id"]], ensure_ascii=False, separators=(",", ":")
+ )
+ if row["scope_key"] != expected:
+ continue
+ if (
+ row["container_id"] is None
+ and row["record_type"] == cycle.configuration.root_record_type
+ ):
+ connection.execute(
+ sa.text("UPDATE capture_scan SET membership_attempt_id=:attempt WHERE id=:id"),
+ {"attempt": cycle.root_writer_attempt_id, "id": row["id"]},
+ )
+ elif (
+ row["container_id"] is not None
+ and row["record_type"] in cycle.configuration.child_record_types
+ ):
+ parent = (
+ connection.execute(
+ sa.text("""
+ SELECT id,visibility_epoch FROM entity
+ WHERE organization_id=:org AND sync_id=:sync
+ AND entity_definition_short_name=:kind AND native_id=:native
+ AND container_id IS NULL AND record_revision>0
+ """),
+ {
+ "org": scope["org"],
+ "sync": scope["sync"],
+ "kind": cycle.configuration.root_record_type,
+ "native": row["container_id"],
+ },
+ )
+ .mappings()
+ .one_or_none()
+ )
+ if parent is not None:
+ connection.execute(
+ sa.text("""
+ UPDATE capture_scan SET parent_record_id=:parent,
+ parent_visibility_epoch=:epoch,
+ scope_key=:key WHERE id=:id
+ """),
+ {
+ "parent": parent["id"],
+ "epoch": parent["visibility_epoch"],
+ "key": expected + "|" + str(parent["id"]),
+ "id": row["id"],
+ },
+ )
+
+
+def upgrade():
+ """Bind old flat scans only to validated, declared tenant-scoped root records."""
+ op.add_column(
+ "capture_scan",
+ sa.Column("parent_record_id", sa.Uuid(), sa.ForeignKey("entity.id", ondelete="CASCADE")),
+ )
+ op.add_column("capture_scan", sa.Column("parent_visibility_epoch", sa.BigInteger()))
+ op.add_column("capture_scan", sa.Column("membership_attempt_id", sa.Uuid()))
+ bind_legacy_scans(op.get_bind())
+ op.create_index(
+ "idx_capture_scan_parent",
+ "capture_scan",
+ ["sync_id", "parent_record_id", "record_type"],
+ unique=True,
+ postgresql_where=sa.text("parent_record_id IS NOT NULL"),
+ )
+
+
+def downgrade():
+ """Nested recovery cannot safely be interpreted by the flat engine."""
+ raise RuntimeError("Nested capture scopes require explicit retirement before downgrade")
diff --git a/backend/alembic/versions/0008_scope_execution.py b/backend/alembic/versions/0008_scope_execution.py
new file mode 100644
index 0000000000..c6e9483c97
--- /dev/null
+++ b/backend/alembic/versions/0008_scope_execution.py
@@ -0,0 +1,21 @@
+"""Retain completed scope evidence independently of current page progress."""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0008"
+down_revision = "0007"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Keep existing rows explicitly without attested mixed-scope history."""
+ op.add_column("capture_scan", sa.Column("execution_state", postgresql.JSONB(), nullable=True))
+
+
+def downgrade():
+ """Remove mixed execution evidence only on an explicit schema downgrade."""
+ op.drop_column("capture_scan", "execution_state")
diff --git a/backend/alembic/versions/0009_extraction_coverage.py b/backend/alembic/versions/0009_extraction_coverage.py
new file mode 100644
index 0000000000..4763ae5845
--- /dev/null
+++ b/backend/alembic/versions/0009_extraction_coverage.py
@@ -0,0 +1,23 @@
+"""Generation-owned extraction coverage; legacy publications remain unknown."""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0009"
+down_revision = "0008"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Old generations remain explicitly unknown until reprojection."""
+ op.add_column(
+ "projection_generation", sa.Column("extraction_coverage", postgresql.JSONB(), nullable=True)
+ )
+
+
+def downgrade():
+ """Remove derived evidence only on explicit downgrade."""
+ op.drop_column("projection_generation", "extraction_coverage")
diff --git a/backend/alembic/versions/0010_owned_provisioning.py b/backend/alembic/versions/0010_owned_provisioning.py
new file mode 100644
index 0000000000..086b489fae
--- /dev/null
+++ b/backend/alembic/versions/0010_owned_provisioning.py
@@ -0,0 +1,73 @@
+"""Durable account provisioning and generation-fenced capture admission."""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0010"
+down_revision = "0009"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Existing sources/jobs remain unmanaged generation zero."""
+ for table, column in (
+ ("sync", "provisioning_generation"),
+ ("sync", "provisioning_ready_generation"),
+ ("sync_job", "provisioning_generation"),
+ ):
+ op.add_column(table, sa.Column(column, sa.BigInteger(), nullable=False, server_default="0"))
+ op.create_table(
+ "owned_provisioning",
+ sa.Column("id", sa.UUID(), primary_key=True),
+ sa.Column("created_at", sa.DateTime(), nullable=False),
+ sa.Column("modified_at", sa.DateTime(), nullable=False),
+ sa.Column(
+ "organization_id",
+ sa.UUID(),
+ sa.ForeignKey("organization.id", ondelete="CASCADE"),
+ nullable=False,
+ ),
+ sa.Column("client_namespace", sa.String(32), nullable=False),
+ sa.Column("account_id", sa.UUID(), nullable=False),
+ sa.Column("generation", sa.BigInteger(), nullable=False),
+ sa.Column("request_hash", sa.String(64), nullable=False),
+ sa.Column("request_payload", postgresql.JSONB(), nullable=False),
+ sa.Column("desired_state", sa.String(20), nullable=False),
+ sa.Column("observed_generation", sa.BigInteger(), nullable=False, server_default="0"),
+ sa.Column(
+ "source_connection_id",
+ sa.UUID(),
+ sa.ForeignKey("source_connection.id", ondelete="RESTRICT"),
+ ),
+ sa.Column("sync_id", sa.UUID(), sa.ForeignKey("sync.id", ondelete="RESTRICT")),
+ sa.Column("initial_job_id", sa.UUID(), sa.ForeignKey("sync_job.id", ondelete="SET NULL")),
+ sa.Column("cancellation_job_ids", postgresql.JSONB(), nullable=False, server_default="[]"),
+ sa.Column("verified_at", sa.DateTime()),
+ sa.UniqueConstraint(
+ "organization_id",
+ "client_namespace",
+ "account_id",
+ name="uq_owned_provisioning_account",
+ ),
+ sa.UniqueConstraint("sync_id", name="uq_owned_provisioning_sync"),
+ sa.CheckConstraint(
+ "generation > 0 AND observed_generation >= 0 AND observed_generation <= generation",
+ name="ck_owned_provisioning_generation",
+ ),
+ sa.CheckConstraint("client_namespace = 'almanac'", name="ck_owned_provisioning_namespace"),
+ sa.CheckConstraint(
+ "desired_state IN ('active', 'paused', 'disconnected')",
+ name="ck_owned_provisioning_state",
+ ),
+ )
+
+
+def downgrade():
+ """Only explicit rollback removes generation protections."""
+ op.drop_table("owned_provisioning")
+ op.drop_column("sync_job", "provisioning_generation")
+ op.drop_column("sync", "provisioning_ready_generation")
+ op.drop_column("sync", "provisioning_generation")
diff --git a/backend/alembic/versions/0011_retained_projection_text.py b/backend/alembic/versions/0011_retained_projection_text.py
new file mode 100644
index 0000000000..a8d32c1144
--- /dev/null
+++ b/backend/alembic/versions/0011_retained_projection_text.py
@@ -0,0 +1,28 @@
+"""Retain derived text descriptors with their projection publication generation."""
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+revision = "0011"
+down_revision = "0010"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Old generations have unknown representation coverage until reprojected."""
+ op.add_column(
+ "projection_generation",
+ sa.Column(
+ "text_representations",
+ postgresql.JSONB(none_as_null=True),
+ nullable=True,
+ ),
+ )
+
+
+def downgrade():
+ """Remove descriptors; no source originals are affected."""
+ op.drop_column("projection_generation", "text_representations")
diff --git a/backend/alembic/versions/0012_scan_parent_verification.py b/backend/alembic/versions/0012_scan_parent_verification.py
new file mode 100644
index 0000000000..46e5253ef0
--- /dev/null
+++ b/backend/alembic/versions/0012_scan_parent_verification.py
@@ -0,0 +1,28 @@
+"""Exact owner verification on the existing scope scan.
+
+Revision ID: 0012
+Revises: 0011
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "0012"
+down_revision = "0011"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Old scans have no verification; never backfill unproven permission."""
+ op.add_column("capture_scan", sa.Column("parent_verified_attempt_id", sa.Uuid(), nullable=True))
+ op.add_column(
+ "capture_scan", sa.Column("parent_verified_revision", sa.BigInteger(), nullable=True)
+ )
+
+
+def downgrade():
+ """Remove only admission receipts, not captured originals."""
+ op.drop_column("capture_scan", "parent_verified_revision")
+ op.drop_column("capture_scan", "parent_verified_attempt_id")
diff --git a/backend/entrypoint.dev.sh b/backend/entrypoint.dev.sh
index 6e018bdbd5..80e49f5d4c 100644
--- a/backend/entrypoint.dev.sh
+++ b/backend/entrypoint.dev.sh
@@ -21,11 +21,10 @@ def check_db():
check_db()
"
-# Run migrations using our existing Alembic setup
-echo "Running database migrations..."
-cd /app && poetry run alembic upgrade head
-# cd /app #dev time to not rerun alembic on conflicting branches
+echo "Starting against the provisioned database..."
+cd /app
+# Migrations and bootstrap are explicit deployment jobs, never replica startup.
# Start application with hot reloading enabled
echo "Starting application with hot reloading..."
-poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001 --reload --reload-dir /app/airweave
+poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001-dir /app/airweave
diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh
index 00457d17bc..5284e3fca5 100644
--- a/backend/entrypoint.sh
+++ b/backend/entrypoint.sh
@@ -1,56 +1,7 @@
-#!/bin/bash
-set -e
+#!/bin/sh
+set -eu
-# Wait for database to be ready
-echo "Waiting for database..."
-echo "POSTGRES_HOST: $POSTGRES_HOST"
-echo "POSTGRES_USER: $POSTGRES_USER"
-echo "POSTGRES_DB: $POSTGRES_DB"
-echo "POSTGRES_PORT: ${POSTGRES_PORT:-5432}"
-echo "POSTGRES_SSLMODE: ${POSTGRES_SSLMODE:-prefer}"
-
-poetry run python -c "
-from tenacity import retry, stop_after_attempt, wait_fixed
-from sqlalchemy import create_engine
-import os
-import sys
-
-@retry(stop=stop_after_attempt(15), wait=wait_fixed(5))
-def check_db():
- try:
- user = os.environ['POSTGRES_USER']
- password = os.environ['POSTGRES_PASSWORD']
- host = os.getenv('POSTGRES_HOST', 'db')
- db = os.getenv('POSTGRES_DB', 'airweave')
- port = os.getenv('POSTGRES_PORT', '5432')
- sslmode = os.getenv('POSTGRES_SSLMODE', 'prefer')
-
- print(f'Attempting to connect to: postgresql://{user}:***@{host}:{port}/{db}?sslmode={sslmode}')
-
- # Use SSL mode from environment variable
- connection_string = f'postgresql://{user}:{password}@{host}:{port}/{db}?sslmode={sslmode}'
- engine = create_engine(connection_string, connect_args={'connect_timeout': 10})
- connection = engine.connect()
- connection.close()
- print('Database connection successful!')
-
- except Exception as e:
- print(f'Database connection failed: {str(e)}')
- print(f'Error type: {type(e).__name__}')
- raise e
-
-try:
- check_db()
-except Exception as e:
- print(f'Final database connection error: {e}')
- sys.exit(1)
-"
-
-# Run migrations using our existing Alembic setup
-echo "Running database migrations..."
-cd /app && poetry run alembic upgrade head
-# cd /app #dev time to not rerun alembic on conflicting branches
-
-# Start application
-echo "Starting application..."
-poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001 --reload
+# Database migrations and bootstrap are explicit deployment jobs. Startup checks
+# the provisioned schema; the process supervisor handles unavailable dependencies.
+# exec delivers termination signals directly to Uvicorn for graceful shutdown.
+exec python -m uvicorn airweave.main:app --host 0.0.0.0 --port 8001
diff --git a/backend/pytest.ini b/backend/pytest.ini
index 9e11fe9871..1d23b72188 100644
--- a/backend/pytest.ini
+++ b/backend/pytest.ini
@@ -14,7 +14,7 @@ log_cli = true
log_cli_level = INFO
# Set PYTHONPATH
-pythonpath = .
+pythonpath = . ..
# Configure asyncio for async tests
asyncio_mode = auto
diff --git a/backend/scripts/REPROJECTION.md b/backend/scripts/REPROJECTION.md
new file mode 100644
index 0000000000..e16513b180
--- /dev/null
+++ b/backend/scripts/REPROJECTION.md
@@ -0,0 +1,49 @@
+# Canonical search metadata upgrade
+
+This is an operator action against an explicitly selected private source store,
+not a product CLI or a capture command. It performs no provider reads/writes.
+
+Deploy the expanded Vespa application schema **and only compatible API/workers**
+before changing versions. An old worker can publish a new numeric version without
+the new metadata; the number alone cannot prove compatible binaries. Stop/drain
+old workers through the normal release procedure, never mid-capture ad hoc.
+
+Dry-run (default), from `backend/` with the intended staging database configuration:
+
+```
+python -m scripts.reproject_canonical --organization ORGANIZATION_UUID \
+ --sync SYNC_UUID --expect 1 --target 2
+```
+
+Inspect the exact authenticated source binding, current version, captured count
+(including tombstones), and target-pending count. These counts describe captured
+rows, not search hits or active visible rows. Add `--apply` only after review.
+No command was applied while preparing this feature.
+
+Apply locks the existing Sync row, verifies organization/source/collection scope,
+and changes its existing `index_pipeline_version`. That immediately invalidates
+older publications through the existing SQL visibility fence. Existing pending
+projection selection now finds the records; capture checkpoints remain unchanged.
+The script starts the existing `ProjectCanonicalRecordsWorkflow` from its first
+page. Its normal page/retry/publication behavior remains authoritative.
+
+The database commit and Temporal launch are **not atomic**. If launch fails,
+pending records remain durable. Retry the identical expected/target arguments:
+current version must equal exactly expected or target, never an arbitrary later
+version. An already running deterministic workflow ID is reused; a completed or
+failed run can be restarted to drain remaining pending rows without another bump.
+An accepted workflow is not proof that all content became indexed. Unsupported
+conversions may remain pending and must be reported separately.
+
+Date/type-filtered searches reject versions below 2 with `reindex_required` before
+retrieval. Following the bump, pending publication coverage remains explicit;
+older generations cannot pass SQL validation. Relevance search is still bounded,
+not exhaustive, and this change introduces no chronological ordering or cursor.
+
+New canonical sources initialize their projection version to the current contract
+inside the existing sync-creation transaction, before creating a pending job.
+Capability resolution uses the same CanonicalSource/CanonicalPageSource protocols
+as the capture factory. Legacy sources retain their existing initial version.
+The version is an internal repository argument, not a client-controlled SyncCreate
+field; existing syncs are never upgraded implicitly. Sync service tests verify
+legacy, Gmail and Slack creation paths; this does not prove hosted provisioning.
diff --git a/backend/scripts/reproject_canonical.py b/backend/scripts/reproject_canonical.py
new file mode 100644
index 0000000000..6fc056f37f
--- /dev/null
+++ b/backend/scripts/reproject_canonical.py
@@ -0,0 +1,68 @@
+"""Dry-run by default; scoped re-projection without provider capture or writes."""
+
+import argparse
+import asyncio
+from uuid import UUID
+
+from temporalio.common import WorkflowIDReusePolicy
+from temporalio.exceptions import WorkflowAlreadyStartedError
+
+from airweave.core.config import settings
+from airweave.db.session import AsyncSessionLocal
+from airweave.domains.entities.canonical.reprojection import plan_reprojection
+from airweave.domains.temporal.client import get_client
+from airweave.domains.temporal.workflows.project_canonical_records import (
+ ProjectCanonicalRecordsWorkflow,
+)
+
+
+async def run(args: argparse.Namespace) -> None:
+ """Commit durable pending state, then start/resume the existing Temporal consumer."""
+ async with AsyncSessionLocal() as db:
+ plan = await plan_reprojection(
+ db,
+ args.organization,
+ args.sync,
+ expected_version=args.expect,
+ target_version=args.target,
+ apply=args.apply,
+ )
+ if args.apply:
+ await db.commit()
+ else:
+ await db.rollback()
+ print(plan.model_dump_json())
+ if not args.apply:
+ return
+ # Non-atomic by design: a launch failure leaves pending rows durable. Repeat
+ # the exact expect/target command; don't choose another version to retry.
+ client = await get_client()
+ try:
+ await client.start_workflow(
+ ProjectCanonicalRecordsWorkflow.run,
+ args=[str(args.organization), str(args.sync)],
+ id=plan.workflow_id,
+ task_queue=settings.TEMPORAL_TASK_QUEUE,
+ id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+ )
+ except WorkflowAlreadyStartedError as exc:
+ raise SystemExit(
+ "Projection is already running. The requested version remains durable; "
+ "retry this exact command after it finishes to include previously failed records."
+ ) from exc
+ print("Projection workflow accepted; this does not mean indexing is complete.")
+
+
+def main() -> None:
+ """Parse explicit scope and version pair; dry-run unless --apply is supplied."""
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--organization", type=UUID, required=True)
+ parser.add_argument("--sync", type=UUID, required=True)
+ parser.add_argument("--expect", type=int, required=True)
+ parser.add_argument("--target", type=int, required=True)
+ parser.add_argument("--apply", action="store_true")
+ asyncio.run(run(parser.parse_args()))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/backend/scripts/smoke_local_ocr.py b/backend/scripts/smoke_local_ocr.py
new file mode 100644
index 0000000000..33414fba6a
--- /dev/null
+++ b/backend/scripts/smoke_local_ocr.py
@@ -0,0 +1,32 @@
+"""Qualify packaged OCR weights and the real isolated worker without external inference."""
+
+import argparse
+import asyncio
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pymupdf
+
+from airweave.domains.ocr.local import LocalOcrProvider
+
+
+async def main(tessdata: str) -> None:
+ """A raster image ensures this checks OCR rather than a PDF text-layer shortcut."""
+ with TemporaryDirectory(prefix="ocr-smoke-") as directory:
+ path = Path(directory) / "scan.png"
+ with pymupdf.open() as document:
+ page = document.new_page(width=420, height=160)
+ page.insert_text((25, 80), "ALMANAC 7429", fontsize=30)
+ page.get_pixmap(dpi=150).save(path)
+ provider = LocalOcrProvider(tessdata_path=tessdata, languages=("eng", "hin"))
+ results = await provider.convert_batch([str(path)])
+ result = results[str(path)]
+ if result is None or "ALMANAC" not in result or "7429" not in result:
+ raise RuntimeError("Packaged local OCR did not recover the synthetic identifier")
+ print("Packaged local OCR passed; no external inference calls")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--tessdata", required=True)
+ asyncio.run(main(parser.parse_args().tessdata))
diff --git a/backend/tests/integration/test_canonical_search_prefilters.py b/backend/tests/integration/test_canonical_search_prefilters.py
new file mode 100644
index 0000000000..40c47faa95
--- /dev/null
+++ b/backend/tests/integration/test_canonical_search_prefilters.py
@@ -0,0 +1,101 @@
+"""Opt-in real Vespa: canonical filtering happens before a bounded result window."""
+
+import asyncio
+import json
+import os
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import httpx
+import pytest
+from vespa.application import Vespa
+
+from airweave.core.logging import logger
+from airweave.domains.embedders.types import DenseEmbedding, SparseEmbedding
+from airweave.domains.entities.canonical.search_metadata import epoch_microseconds
+from airweave.domains.entities.canonical.tests.vespa_helpers import deploy_schema
+from airweave.domains.search.adapters.vector_db.filter_translator import FilterTranslator
+from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB
+from airweave.domains.search.owned import OwnedSearchService
+from airweave.domains.search.owned_models import OwnedSearchRequest
+from airweave.domains.search.types import FilterGroup, SearchPlan, SearchQuery
+from airweave.domains.search.types.embeddings import QueryEmbeddings
+
+pytestmark = pytest.mark.skipif(
+ os.environ.get("OWNED_VESPA_TEST") != "1", reason="requires disposable real Vespa"
+)
+
+
+async def test_prefilter_excludes_over_two_hundred_wrong_dates_before_top_matches():
+ async with httpx.AsyncClient(timeout=120) as http:
+ await deploy_schema(http)
+ collection, sync = str(uuid4()), uuid4()
+ vector = [1.0] + [0.0] * 383
+ instant = datetime(2026, 1, 1, microsecond=1, tzinfo=timezone.utc)
+ urls = []
+ expected = None
+ try:
+ for number in range(204):
+ identity = f"canonical-prefilter-{uuid4()}"
+ url = f"http://localhost:8081/document/v1/airweave/base_entity/docid/{identity}"
+ urls.append(url)
+ correct = number == 203
+ known = number != 201
+ fields = {
+ "entity_id": identity,
+ "name": "fundraising",
+ "textual_representation": "fundraising fixture",
+ "payload": json.dumps({}),
+ "airweave_system_metadata_collection_id": collection,
+ "airweave_system_metadata_sync_id": str(sync),
+ "airweave_system_metadata_original_entity_id": identity,
+ "airweave_system_metadata_source_name": "google_calendar",
+ "airweave_system_metadata_entity_type": "GoogleCalendarEventEntity",
+ "airweave_system_metadata_canonical_record_type": "message"
+ if number == 202
+ else "event",
+ "airweave_system_metadata_source_created_known": int(known),
+ "airweave_system_metadata_source_created_us": epoch_microseconds(instant)
+ if number >= 201
+ else 0,
+ "dense_embedding": {"values": vector},
+ "sparse_embedding": {"cells": {"1": 1.0}},
+ }
+ response = await http.post(url, json={"fields": fields})
+ assert response.status_code == 200, response.text
+ if correct:
+ expected = identity
+ context = logger.with_context(request_id="canonical-prefilter-test")
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=context,
+ filter_translator=FilterTranslator(context),
+ )
+ request = OwnedSearchRequest(
+ query="fundraising",
+ sync_ids=(sync,),
+ mode="keyword",
+ record_types=("event",),
+ created_after=instant,
+ )
+ compiled = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary="fundraising"),
+ retrieval_strategy="keyword",
+ limit=200,
+ offset=0,
+ filter_groups=[
+ FilterGroup(conditions=OwnedSearchService._prefilters(request, [sync]))
+ ],
+ ),
+ QueryEmbeddings(
+ dense_embeddings=[DenseEmbedding(vector=vector)],
+ sparse_embedding=SparseEmbedding(indices=[1], values=[1.0]),
+ ),
+ collection,
+ )
+ result = await engine.execute_query(compiled)
+ assert not result.engine_partial
+ assert [hit.entity_id for hit in result.results] == [expected]
+ finally:
+ await asyncio.gather(*(http.delete(url) for url in urls))
diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py
new file mode 100644
index 0000000000..4097054ec8
--- /dev/null
+++ b/backend/tests/integration/test_owned_vespa.py
@@ -0,0 +1,444 @@
+"""Real engine smoke using only synthetic records in a disposable Vespa instance.
+
+Explicit opt-in: OWNED_VESPA_TEST=1. This deploys the repository schema package
+and is unsuitable for an existing application or customer index.
+"""
+
+import json
+import math
+import os
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import httpx
+import pytest
+from vespa.application import Vespa
+
+from airweave.core.logging import logger
+from airweave.domains.embedders.types import DenseEmbedding, SparseEmbedding
+from airweave.domains.entities.canonical.tests.vespa_helpers import deploy_schema
+from airweave.domains.search.adapters.vector_db.filter_translator import FilterTranslator
+from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB
+from airweave.domains.search.adapters.vector_db.vespa_config import ALL_VESPA_SCHEMAS
+from airweave.domains.search.types.embeddings import QueryEmbeddings
+from airweave.domains.search.types.filters import FilterCondition, FilterGroup
+from airweave.domains.search.types.plan import RetrievalStrategy, SearchPlan, SearchQuery
+
+pytestmark = pytest.mark.skipif(
+ os.environ.get("OWNED_VESPA_TEST") != "1", reason="requires disposable real Vespa"
+)
+
+
+@pytest.mark.asyncio
+async def test_real_schema_retrieval_collection_isolation_and_delete():
+ """Exercise real query compilation/rank profiles; vectors are fixed fixtures."""
+ async with httpx.AsyncClient(timeout=120) as http:
+ await deploy_schema(http)
+ collection, other = str(uuid4()), str(uuid4())
+ sync = str(uuid4())
+ vector = [1.0] + [0.0] * 383
+ urls = []
+ try:
+ for number, scope in enumerate((collection, other)):
+ identity = f"owned-smoke-{uuid4()}"
+ url = f"http://localhost:8081/document/v1/airweave/base_entity/docid/{identity}"
+ urls.append(url)
+ fields = {
+ "entity_id": identity,
+ "created_at": int(datetime(2026, 9, 30, tzinfo=timezone.utc).timestamp()),
+ "name": "Synthetic fundraising discussion",
+ "textual_representation": "fundraising discussion fixture",
+ "payload": json.dumps({"web_url": "https://example.com/fixture"}),
+ "airweave_system_metadata_collection_id": scope,
+ "airweave_system_metadata_sync_id": sync,
+ "airweave_system_metadata_source_name": "gmail",
+ "airweave_system_metadata_entity_type": "GmailMessageEntity",
+ "airweave_system_metadata_original_entity_id": identity,
+ "dense_embedding": {"values": vector},
+ "sparse_embedding": {"cells": {"1": 1.0}},
+ }
+ fed = await http.post(url, json={"fields": fields})
+ assert fed.status_code == 200, fed.text
+ if number == 0:
+ expected_id = identity
+ contextual = logger.with_context(request_id="synthetic-vespa-smoke")
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=contextual,
+ filter_translator=FilterTranslator(logger=contextual),
+ )
+ embeddings = QueryEmbeddings(
+ dense_embeddings=[DenseEmbedding(vector=vector)],
+ sparse_embedding=SparseEmbedding(indices=[1], values=[1.0]),
+ )
+ for mode in RetrievalStrategy:
+ compiled = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary="fundraising"),
+ retrieval_strategy=mode,
+ limit=20,
+ offset=0,
+ ),
+ embeddings,
+ collection,
+ )
+ result = await engine.execute_query(compiled)
+ assert not result.engine_partial
+ assert [hit.entity_id for hit in result.results] == [expected_id]
+ assert result.results[0].textual_representation == "fundraising discussion fixture"
+ for operator, expected_ids in (
+ ("greater_than_or_equal", []),
+ ("less_than", [expected_id]),
+ ):
+ filtered = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary="fundraising"),
+ retrieval_strategy=RetrievalStrategy.KEYWORD,
+ filter_groups=[
+ FilterGroup(
+ conditions=[
+ FilterCondition(
+ field="created_at",
+ operator=operator,
+ value="2026-09-30T00:00:00.500Z",
+ )
+ ]
+ )
+ ],
+ limit=20,
+ offset=0,
+ ),
+ embeddings,
+ collection,
+ )
+ result = await engine.execute_query(filtered)
+ assert not result.engine_partial
+ assert [hit.entity_id for hit in result.results] == expected_ids
+ deleted = await http.delete(urls[0])
+ assert deleted.status_code == 200
+ result = await engine.execute_query(compiled)
+ assert result.results == []
+ finally:
+ for url in urls:
+ response = await http.delete(url)
+ assert response.status_code in (200, 404)
+
+
+@pytest.mark.asyncio
+async def test_all_schemas_retrieve_varied_vectors_with_collection_and_sync_scope():
+ """Small all-schema recall regression; exact ID sets, not relevance-quality claims."""
+ collection, foreign_collection, sync, foreign_sync = (str(uuid4()) for _ in range(4))
+ filters = [
+ FilterGroup(
+ conditions=[
+ FilterCondition(
+ field="airweave_system_metadata.sync_id",
+ operator="in",
+ value=[sync],
+ )
+ ]
+ )
+ ]
+ contextual = logger.with_context(request_id="synthetic-all-schema-retrieval")
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=contextual,
+ filter_translator=FilterTranslator(logger=contextual),
+ )
+ urls, expected = [], set()
+ async with httpx.AsyncClient(timeout=120) as http:
+ await deploy_schema(http)
+ try:
+ for index, schema in enumerate(ALL_VESPA_SCHEMAS):
+ # Distinct finite, nonzero vectors exercise angular ANN across schemas.
+ vector = [math.sin((index + 1) * (dimension + 1)) for dimension in range(384)]
+ for scope, source in (
+ (collection, sync),
+ (foreign_collection, sync),
+ (collection, foreign_sync),
+ ):
+ identity = f"all-schema-{uuid4()}"
+ url = f"http://localhost:8081/document/v1/airweave/{schema}/docid/{identity}"
+ urls.append(url)
+ response = await http.post(
+ url,
+ json={
+ "fields": {
+ "entity_id": identity,
+ "name": "Synthetic fundraising discussion",
+ "textual_representation": "fundraising all schema fixture",
+ "payload": "{}",
+ "airweave_system_metadata_collection_id": scope,
+ "airweave_system_metadata_sync_id": source,
+ "airweave_system_metadata_source_name": "gmail",
+ "airweave_system_metadata_entity_type": "SyntheticEntity",
+ "airweave_system_metadata_original_entity_id": identity,
+ "dense_embedding": {"values": vector},
+ "sparse_embedding": {"cells": {"1": 1.0}},
+ }
+ },
+ )
+ assert response.status_code == 200, response.text
+ if scope == collection and source == sync:
+ expected.add(identity)
+ vectors = (
+ [1.0] * 384,
+ [-1.0] * 384,
+ [math.cos(dimension + 1) for dimension in range(384)],
+ )
+ for mode in RetrievalStrategy:
+ for vector in vectors if mode != RetrievalStrategy.KEYWORD else vectors[:1]:
+ compiled = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary="fundraising"),
+ retrieval_strategy=mode,
+ filter_groups=filters,
+ limit=20,
+ offset=0,
+ ),
+ QueryEmbeddings(
+ dense_embeddings=[DenseEmbedding(vector=vector)],
+ sparse_embedding=SparseEmbedding(indices=[1], values=[1.0]),
+ ),
+ collection,
+ )
+ params = compiled.raw["params"]
+ if mode == RetrievalStrategy.KEYWORD:
+ assert "ranking.matching.filterFirstThreshold" not in params
+ else:
+ assert params["ranking.matching.filterFirstThreshold"] == 0.0
+ # Force this fixture through ANN even when its random scope is tiny.
+ # Mutate the actual executable dictionary: raw is PrivateAttr-backed.
+ params["ranking.matching.approximateThreshold"] = 0.0
+ result = await engine.execute_query(compiled)
+ assert not result.engine_partial, (mode, result.engine_coverage_percent)
+ assert {item.entity_id for item in result.results} == expected, mode
+ assert len(result.results) == len(ALL_VESPA_SCHEMAS)
+ browsed = await engine.filter_search(filters, collection, limit=20)
+ assert {item.entity_id for item in browsed} == expected
+ assert await engine.count(filters, collection) == len(expected)
+ finally:
+ # Delete and read back only this run's random synthetic document IDs.
+ for url in urls:
+ deleted = await http.delete(url)
+ assert deleted.status_code in (200, 404), deleted.text
+ assert (await http.get(url)).status_code == 404
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+ os.environ.get("OWNED_MINILM_TEST") != "1", reason="requires disposable pinned MiniLM service"
+)
+async def test_real_minilm_paraphrase_retrieval_and_exact_keyword():
+ """Small synthetic retrieval baseline, not a general search-quality benchmark."""
+ import asyncio
+
+ from airweave.domains.embedders.dense.local import LocalDenseEmbedder
+ from airweave.domains.embedders.sparse.fastembed import FastEmbedSparseEmbedder
+
+ corpus = [
+ (
+ "funding",
+ "Investors and venture capital",
+ "We are raising a seed round to finance the startup. "
+ "Venture investors will purchase equity in the company.",
+ ),
+ (
+ "travel",
+ "Summer holiday",
+ "Book a seaside hotel for our family vacation. We will swim and relax on the beach.",
+ ),
+ (
+ "incident",
+ "Database incident ZXQ4829",
+ "The database connection pool exhausted its limit. "
+ "Restarting the server restored normal query latency.",
+ ),
+ (
+ "food",
+ "Dinner recipe",
+ "Roast potatoes with olive oil and garlic. Serve the vegetables with fresh bread.",
+ ),
+ ]
+ collection, foreign = str(uuid4()), str(uuid4())
+ embedder = LocalDenseEmbedder(inference_url="http://localhost:8080", dimensions=384)
+ sparse_embedder = FastEmbedSparseEmbedder(model="Qdrant/bm25")
+ urls = []
+ async with httpx.AsyncClient(timeout=120) as http:
+ try:
+ for _attempt in range(60):
+ try:
+ health = await http.get("http://localhost:8080/.well-known/ready", timeout=5)
+ if health.status_code == 204:
+ break
+ except httpx.HTTPError:
+ pass
+ await asyncio.sleep(2)
+ else:
+ pytest.fail("Pinned MiniLM service did not become healthy")
+ await deploy_schema(http)
+ vectors = await embedder.embed_many([title + "\n" + text for _, title, text in corpus])
+ sparse_vectors = await sparse_embedder.embed_many(
+ [title + "\n" + text for _, title, text in corpus]
+ )
+ identities = {}
+ for (key, title, text), vector, sparse in zip(
+ corpus, vectors, sparse_vectors, strict=True
+ ):
+ assert len(vector.vector) == 384 and any(vector.vector)
+ identity = f"quality-{uuid4()}"
+ identities[key] = identity
+ # A duplicate in another collection must not become a result.
+ for scope, suffix in ((collection, ""), (foreign, "-foreign")):
+ document_id = identity + suffix
+ url = f"http://localhost:8081/document/v1/airweave/base_entity/docid/{document_id}"
+ urls.append(url)
+ response = await http.post(
+ url,
+ json={
+ "fields": {
+ "entity_id": document_id,
+ "name": title,
+ "textual_representation": text,
+ "payload": json.dumps({"web_url": "https://example.com/" + key}),
+ "airweave_system_metadata_collection_id": scope,
+ "airweave_system_metadata_sync_id": str(uuid4()),
+ "airweave_system_metadata_source_name": "gmail",
+ "airweave_system_metadata_entity_type": "GmailMessageEntity",
+ "airweave_system_metadata_original_entity_id": document_id,
+ "dense_embedding": {"values": vector.vector},
+ "sparse_embedding": {
+ "cells": {
+ str(index): value
+ for index, value in zip(
+ sparse.indices, sparse.values, strict=True
+ )
+ }
+ },
+ }
+ },
+ )
+ assert response.status_code == 200, response.text
+ contextual = logger.with_context(request_id="synthetic-minilm-baseline")
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=contextual,
+ filter_translator=FilterTranslator(logger=contextual),
+ )
+ for query, expected, mode in (
+ (
+ "How are we obtaining money from equity investors?",
+ "funding",
+ RetrievalStrategy.SEMANTIC,
+ ),
+ ("How was the database outage resolved?", "incident", RetrievalStrategy.SEMANTIC),
+ ("ZXQ4829", "incident", RetrievalStrategy.KEYWORD),
+ ):
+ embeddings = QueryEmbeddings(
+ dense_embeddings=[await embedder.embed(query)]
+ if mode == RetrievalStrategy.SEMANTIC
+ else None,
+ sparse_embedding=await sparse_embedder.embed(query)
+ if mode == RetrievalStrategy.KEYWORD
+ else None,
+ )
+ compiled = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary=query), retrieval_strategy=mode, limit=4, offset=0
+ ),
+ embeddings,
+ collection,
+ )
+ results = await engine.execute_query(compiled)
+ assert not results.engine_partial
+ assert results.results, (query, "no matches")
+ assert results.results[0].entity_id == identities[expected], (
+ query,
+ [hit.name for hit in results.results],
+ )
+ assert all(hit.entity_id in identities.values() for hit in results.results)
+ finally:
+ await embedder.close()
+ for url in urls:
+ response = await http.delete(url)
+ assert response.status_code in (200, 404)
+
+
+@pytest.mark.asyncio
+async def test_dynamic_summary_preserves_full_chunks_and_inherited_fields():
+ """Actual default summaries must expose lexical fragments without replacing originals."""
+ collection, sync = str(uuid4()), str(uuid4())
+ vector = [1.0] + [0.0] * 383
+ # The match is beyond the owned service's 2,000-character prefix fallback.
+ original = "Ordinary introductory context. " * 200 + (
+ "Uniquelexicalneedle नमस्ते 原文 concludes this record."
+ )
+ contextual = logger.with_context(request_id="synthetic-dynamic-summary")
+ engine = VespaVectorDB(
+ app=Vespa(url="http://localhost", port=8081),
+ logger=contextual,
+ filter_translator=FilterTranslator(logger=contextual),
+ )
+ embeddings = QueryEmbeddings(
+ dense_embeddings=[DenseEmbedding(vector=vector)],
+ sparse_embedding=SparseEmbedding(indices=[1], values=[1.0]),
+ )
+ urls, expected = [], set()
+ async with httpx.AsyncClient(timeout=120) as http:
+ await deploy_schema(http)
+ try:
+ for schema in ("file_entity", "email_entity"):
+ identity = f"snippet-{uuid4()}"
+ expected.add(identity)
+ url = f"http://localhost:8081/document/v1/airweave/{schema}/docid/{identity}"
+ urls.append(url)
+ fed = await http.post(
+ url,
+ json={
+ "fields": {
+ "entity_id": identity,
+ "name": "Synthetic original",
+ "payload": "{}",
+ "url": "https://example.test/original",
+ "textual_representation": original,
+ "airweave_system_metadata_collection_id": collection,
+ "airweave_system_metadata_sync_id": sync,
+ "airweave_system_metadata_source_name": "gmail",
+ "airweave_system_metadata_entity_type": "GmailMessageEntity",
+ "airweave_system_metadata_original_entity_id": identity,
+ "dense_embedding": {"values": vector},
+ "sparse_embedding": {"cells": {"1": 1.0}},
+ }
+ },
+ )
+ assert fed.status_code == 200, fed.text
+ for mode in RetrievalStrategy:
+ compiled = await engine.compile_query(
+ SearchPlan(
+ query=SearchQuery(primary="Uniquelexicalneedle"),
+ retrieval_strategy=mode,
+ limit=10,
+ offset=0,
+ ),
+ embeddings,
+ collection,
+ )
+ assert "presentation.summary" not in compiled.raw["params"]
+ results = await engine.execute_query(compiled)
+ assert {hit.entity_id for hit in results.results} == expected
+ for hit in results.results:
+ assert hit.textual_representation == original
+ assert hit.airweave_system_metadata.original_entity_id == hit.entity_id
+ assert hit.airweave_system_metadata.sync_id == sync
+ assert hit.url == "https://example.test/original"
+ if mode != RetrievalStrategy.SEMANTIC:
+ assert hit.query_snippet is not None
+ assert "uniquelexicalneedle" in hit.query_snippet.lower()
+ assert len(hit.query_snippet) < len(original)
+ assert "नमस्ते" in hit.query_snippet and "原文" in hit.query_snippet
+ # Semantic-only summary contents are engine-dependent; owned
+ # HTTP tests prove they never replace the semantic fallback.
+ finally:
+ for url in urls:
+ assert (await http.delete(url)).status_code in (200, 404)
+ assert (await http.get(url)).status_code == 404
diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md
new file mode 100644
index 0000000000..fa881b8374
--- /dev/null
+++ b/backend/tests/live/README.md
@@ -0,0 +1,495 @@
+# Durable live capture proof
+
+`canonical_capture.py` is opt-in and makes read-only Gmail/Drive/Calendar provider requests.
+It commits real captured records and blobs into disposable local storage, then
+checks reads over fresh database connections, paginated lists/change feeds, exact
+payload equality, blob SHA256/size, and an identical replay with zero new changes.
+It does **not** assert a completed account sync or save an upstream checkpoint.
+
+Requirements:
+
+- A separately initialized PostgreSQL database `sync_tests`, user `sync_test`,
+ reachable only through an explicitly supplied private Unix-socket directory.
+- `CANONICAL_TEST_DATABASE_URL` using `postgresql+asyncpg`, that database/user,
+ no network hostname, and `?host=/absolute/private/socket&port=...`.
+- `COMPOSIO_API_KEY`, `LIVE_GMAIL_ACCOUNT_ID`, `LIVE_DRIVE_ACCOUNT_ID`, and
+ `LIVE_EXPECTED_EMAIL` in the subprocess environment. The two accounts must
+ already be connected. The harness checks their actual provider email identity.
+
+Optionally set `LIVE_PROVIDERS=gmail` or `LIVE_PROVIDERS=google_drive` to verify
+only one provider; default is both. `LIVE_PROVIDERS=google_calendar` selects Calendar
+and requires `LIVE_CALENDAR_ACCOUNT_ID`. It verifies the primary CalendarList ID
+against `LIVE_EXPECTED_EMAIL`. Only selected provider account variables are required.
+
+Run from `backend`, using your approved secret injector:
+
+```sh
+PYTHONPATH=. .venv/bin/python tests/live/canonical_capture.py
+```
+
+The source sample stops after at least three records and one blob, or 25 records.
+Gmail uses `newer_than:7d smaller:5M`; capture/download limits are reduced to 10 MiB
+for this probe. A Gmail/Drive sample without a blob fails. Calendar samples up to
+25 original records and requires an event; its JSON does not require blobs, so
+blob verification is null when none exist. Listing includes tombstones so sparse
+cancellations are verified alongside ordinary originals. No scope-completion observations
+are persisted. The application database hostname is overridden to an unusable
+value; only the explicit test URL reaches a database engine.
+
+A random `canonical_live_*` schema receives the real migrations. Files live under
+one temporary directory with mode `0700`, and the process uses umask `077`.
+Schema and files are removed in `finally`, including on failure. Database WAL may
+retain deleted pages until the disposable PostgreSQL cluster is destroyed; use a
+private test cluster and remove it after the larger verification session. Do not
+run this against a production database or retain/export provider payloads.
+
+Output contains only counts, booleans, safe exception class names, stages, HTTP
+status/host and allowlisted provider error codes. The live
+script never prints credentials, record contents, account identifiers, filenames,
+or SQL parameters. Automated mock/store tests remain separate from this proof.
+
+## Almanac HTTP consumer handoff
+
+For Gmail, optionally set `LIVE_ALMANAC_ROOT` to the Almanac checkout and
+`LIVE_ALMANAC_PYTHON` to its Python 3.12 **venv executable path** (do not resolve
+its symlink to the base interpreter). Run the same `canonical_capture.py` command
+in Airweave's Python 3.13 environment. The callback serves the committed sample
+through the actual `api_router` on an ephemeral loopback port, then invokes
+Almanac's `backend/tests/live/read_owned_gmail.py`.
+
+The child receives a mode0600 manifest inside the existing mode0700 directory;
+its environment contains no provider credentials. It reads an actual thread and
+record/blob endpoints using `OwnedMailThreadReader` and `OwnedSourceRecords`,
+compares native header/body/reference and payload digests, verifies blob SHA/size,
+and checks that an unauthenticated probe request is denied. The server uses a
+random-key-guarded **test context override**, not hosted WorkOS or real persisted
+API-key authentication. It creates no production Almanac account binding. The
+child and server are stopped before the enclosing harness removes schema/files.
+
+Verified 2026-09-30: 7 Gmail originals, 7 journal changes, 2 blobs totaling
+1,383,758 bytes, exact replay with 0 new changes and fresh-connection readback.
+The real HTTP → Almanac reader returned 1 message with complete body and verified
+native fields; both blobs passed HTTP SHA/size verification. That selected body
+was inline (`external_body_blobs=0`): external-body retrieval is covered by
+synthetic tests, not claimed as live-covered by this sample. After the initial
+venv-launch failure and successful rerun, independent checks found 0 retained
+`canonical_live_*` schemas and 0 private live directories. No complete-source
+coverage, checkpoint, hosted authentication, or deployed product proof is claimed.
+
+## Historical Slack and Wispr samples
+
+Both samples below used the former observation adapters. Slack and Wispr now
+use durable pages; `canonical_capture.py` rejects these sources and directs callers
+to the shared `provider_lifecycle.py`. Do not use `LIVE_PROVIDERS=wispr` or
+`LIVE_PROVIDERS=slack` with the old sampler.
+
+The current lifecycle selects `LIVE_LIFECYCLE_PROVIDER=wispr`, with
+`LIVE_WISPR_ACCOUNT_ID` and `LIVE_WISPR_USER_ID` for the existing connected account.
+It verifies exact ACTIVE Composio account ID, toolkit and principal. This is weaker
+than a provider mailbox/profile check and is not an Almanac user binding. Wispr
+requires no email variable. Slack verifies `auth.test` and matching `users.info`
+email against `LIVE_EXPECTED_EMAIL`.
+
+The historical samples stopped after three distinct records or25 observations.
+Exact native JSON survived durable readback and identical replay added no journal
+changes. Source markers/checkpoints remained deliberately uncommitted. These are
+historical sample results, not verification of the new page adapters.
+
+Verified 2026-09-30: Slack produced four observations and three distinct records
+(channel/message), with three journal changes, exact fresh-connection readback
+and zero replay changes. Its provider email matched. No file-bearing message was
+in that sample, so attachment bytes/partial file coverage were not live-tested.
+Wispr produced three meeting records and three journal changes; all three remain
+explicitly partial because the provider omits raw editor data/deletion history.
+Its replay added zero changes. Neither sample contained blobs or proved complete
+source coverage. Both private schema and files were removed afterward.
+
+## Current Gmail page lifecycle and bounded recovery
+
+The current Gmail source uses the production durable page/checkpoint protocol.
+The existing `provider_lifecycle.py` remains the runner; no live execution of this
+new harness is claimed here. The earlier lifecycle results below describe the
+previous observation adapter.
+
+With `LIVE_LIFECYCLE_PROVIDER=gmail`, the default still freezes one seven-day
+query and runs two complete configured-scope passes in fresh processes. Validation
+now reads `canonical_cycle` and the fenced checkpoint; legacy marker counts and
+`history_id` are not evidence for the page source. Payload digest differences are
+reported honestly rather than requiring a stationary mailbox.
+
+For an explicitly selected unfiltered **partial recovery** probe, additionally set
+`LIVE_GMAIL_UNFILTERED=1 LIVE_GMAIL_RESUME=1`. This clears the source's default
+label/category filters rather than treating an empty custom query as unfiltered.
+The first child exits with code75 only after SQL proves a message page committed.
+A second child uses the same running job and a new attempt, resumes the identical
+cycle/sweep/version, and exits with code76 after another acknowledged page.
+Parent SQL readback independently verifies advancement. No payloads or provider
+cursors enter printed evidence. A mailbox completing before interruption reports
+`recovery_not_exercised`; it is not successful recovery proof.
+
+The existing600 provider-request,250 observation,600-second and256MiB blob-write
+budgets are shared across the interruption/resume pair; each MIME blob remains
+limited to10MiB. Exact known-message refreshes consume request, observation and
+storage budgets too. Normal identity verification still reads Gmail's profile in
+each process. Separate counters require **zero capture-boundary profile reads**
+in the resumed process; identity verification is not counted as a boundary reset.
+
+`LIVE_GMAIL_UNFILTERED=1` without the resume flag selects two bounded full/changes
+processes. A changes claim requires the first baseline actually to finish, the
+next process to load its promoted checkpoint and issue history requests using that
+exact boundary, and the second checkpoint to commit. Budget exhaustion remains
+incomplete; no sampled records or interrupted baseline establish mailbox completion.
+This option does not increase any existing download or process limit.
+
+## Historical configured Gmail reconciliation lifecycle
+
+`PYTHONPATH=. .venv/bin/python tests/live/provider_lifecycle.py` uses the same private
+PostgreSQL, Composio key, `LIVE_GMAIL_ACCOUNT_ID` and `LIVE_EXPECTED_EMAIL` inputs.
+It freezes a seven-day `after: