From ccefe7650c0a9e91379ae31b3cbec0340c39b338 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 01:48:07 -0700 Subject: [PATCH 001/266] fix(sync): fail runs when checkpoint persistence fails --- .../domains/sync_pipeline/orchestrator.py | 2 + .../sync_pipeline/tests/test_orchestrator.py | 28 +++++++++++++ .../airweave/domains/syncs/cursors/service.py | 16 ++++---- .../cursors/tests/test_failure_semantics.py | 39 +++++++++++++++++++ 4 files changed, 77 insertions(+), 8 deletions(-) create mode 100644 backend/airweave/domains/syncs/cursors/tests/test_failure_semantics.py diff --git a/backend/airweave/domains/sync_pipeline/orchestrator.py b/backend/airweave/domains/sync_pipeline/orchestrator.py index d8e30cc928..e5f12e5e72 100644 --- a/backend/airweave/domains/sync_pipeline/orchestrator.py +++ b/backend/airweave/domains/sync_pipeline/orchestrator.py @@ -664,6 +664,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 diff --git a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py index cec2075275..f2446ac29e 100644 --- a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py +++ b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py @@ -468,3 +468,31 @@ 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.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/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()) == {} From 77f90e51a8a4282b1fdcbdb49bbb4503d42f1d0a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 01:57:23 -0700 Subject: [PATCH 002/266] feat(store): add fenced canonical records and authenticated exact reads --- .env.example | 14 +- backend/SERVICE_AUTH.md | 55 +++ backend/airweave/adapters/identity/auth0.py | 4 +- backend/airweave/api/auth.py | 25 +- backend/airweave/api/context_resolver.py | 33 +- backend/airweave/api/deps.py | 14 +- .../api/tests/test_context_resolver_auth.py | 58 ++- .../airweave/api/tests/test_service_auth.py | 177 ++++++++ .../api/tests/test_sync_authorization.py | 39 ++ backend/airweave/api/v1/api.py | 2 + backend/airweave/api/v1/endpoints/records.py | 79 ++++ backend/airweave/api/v1/endpoints/sync.py | 19 +- backend/airweave/core/config/__init__.py | 3 +- backend/airweave/core/config/enums.py | 8 + backend/airweave/core/config/settings.py | 62 ++- .../config/tests/test_settings_validators.py | 3 + backend/airweave/core/container/factory.py | 2 +- backend/airweave/db/bootstrap.py | 41 ++ backend/airweave/domains/embedders/config.py | 14 +- .../domains/entities/canonical/models.py | 68 ++++ .../domains/entities/canonical/query.py | 135 +++++++ .../entities/canonical/query_models.py | 62 +++ .../domains/entities/canonical/query_store.py | 50 +++ .../domains/entities/canonical/requests.py | 99 +++++ .../domains/entities/canonical/service.py | 57 +++ .../domains/entities/canonical/store.py | 377 ++++++++++++++++++ .../entities/canonical/tests/conftest.py | 132 ++++++ .../entities/canonical/tests/helpers.py | 23 ++ .../entities/canonical/tests/test_http.py | 66 +++ .../entities/canonical/tests/test_postgres.py | 344 ++++++++++++++++ .../entities/canonical/tests/test_query.py | 122 ++++++ backend/airweave/main.py | 22 +- backend/airweave/models/__init__.py | 2 + backend/airweave/models/entity.py | 41 +- backend/airweave/models/entity_change.py | 24 ++ backend/airweave/models/sync.py | 9 +- backend/airweave/models/sync_job.py | 1 - backend/airweave/schemas/entity.py | 1 + .../versions/0001_canonical_records.py | 104 +++++ backend/entrypoint.dev.sh | 9 +- backend/entrypoint.sh | 9 +- 41 files changed, 2291 insertions(+), 118 deletions(-) create mode 100644 backend/SERVICE_AUTH.md create mode 100644 backend/airweave/api/tests/test_service_auth.py create mode 100644 backend/airweave/api/tests/test_sync_authorization.py create mode 100644 backend/airweave/api/v1/endpoints/records.py create mode 100644 backend/airweave/db/bootstrap.py create mode 100644 backend/airweave/domains/entities/canonical/models.py create mode 100644 backend/airweave/domains/entities/canonical/query.py create mode 100644 backend/airweave/domains/entities/canonical/query_models.py create mode 100644 backend/airweave/domains/entities/canonical/query_store.py create mode 100644 backend/airweave/domains/entities/canonical/requests.py create mode 100644 backend/airweave/domains/entities/canonical/service.py create mode 100644 backend/airweave/domains/entities/canonical/store.py create mode 100644 backend/airweave/domains/entities/canonical/tests/conftest.py create mode 100644 backend/airweave/domains/entities/canonical/tests/helpers.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_http.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_postgres.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_query.py create mode 100644 backend/airweave/models/entity_change.py create mode 100644 backend/alembic/versions/0001_canonical_records.py diff --git a/.env.example b/.env.example index bac0af33c9..aa157cb7fa 100644 --- a/.env.example +++ b/.env.example @@ -14,15 +14,15 @@ REDIS_HOST=localhost REDIS_PORT=6379 # Authentication -# Generated by start.sh if empty -FIRST_SUPERUSER= -# Generated by start.sh if empty -FIRST_SUPERUSER_PASSWORD= +# Local demo bootstrap only; omit these in api_key mode. +# FIRST_SUPERUSER= +# FIRST_SUPERUSER_PASSWORD= ENCRYPTION_KEY= -AUTH_ENABLED=false +AUTH_MODE=api_key +# Opt into AUTH_MODE=local only for isolated local development. STATE_SECRET= -# Auth0 Configuration (required if AUTH_ENABLED=true) +# Auth0 Configuration (required if AUTH_MODE=auth0) AUTH0_DOMAIN= AUTH0_AUDIENCE= AUTH0_RULE_NAMESPACE= @@ -123,8 +123,6 @@ ADDITIONAL_CORS_ORIGINS= # Other Settings PROJECT_NAME=Airweave LOG_LEVEL=INFO -RUN_ALEMBIC_MIGRATIONS=true -RUN_DB_SYNC=true CODE_SUMMARIZER_ENABLED=false SVIX_URL=http://localhost:8071 diff --git a/backend/SERVICE_AUTH.md b/backend/SERVICE_AUTH.md new file mode 100644 index 0000000000..72b3eb9fc8 --- /dev/null +++ b/backend/SERVICE_AUTH.md @@ -0,0 +1,55 @@ +# Standalone service authentication and database setup + +Use `AUTH_MODE=api_key` for the owned service. Every ordinary API request needs +an existing organization-scoped `X-API-Key`. The default is `api_key`: omitting +credentials does not select a demo user. This mode makes no Auth0 network calls +and does not create users from service credentials. User-only routes still +require user authentication; existing Connect sessions retain their own token +and scope checks. Auth0 mode preserves the existing user-first authentication +precedence when both user and API-key credentials are supplied. + +`AUTH_MODE=auth0` retains hosted Auth0 user authentication and API-key access. +The existing Auth0 configuration remains required in that mode. +`AUTH_MODE=local` explicitly enables insecure demo authentication and is allowed +only for `ENVIRONMENT=local` or `test`. Never deploy local mode. The deprecated +`AUTH_ENABLED=true/false` is accepted as a migration input mapping to auth0/local; +contradicting it with AUTH_MODE is an error. Remove AUTH_ENABLED when adopting +AUTH_MODE. The environment label is operator-controlled, not an isolation boundary. + +Keys are checked against persisted credentials on every request; expiry and +revocation take effect without waiting for a key-to-organization cache. All +record and job authorization still uses the key's organization. Authenticated +sync SSE subscriptions authorize job ownership before subscribing. + +## Explicit database setup + +Neither API startup nor the container entrypoint migrates the database or seeds +users, organizations, or keys. `RUN_ALEMBIC_MIGRATIONS` is no longer a runtime +switch. Run setup as a separately authorized deployment job with reviewed, +isolated database credentials: + +```sh +# From backend/, against the intended isolated target only: +poetry run alembic upgrade head +poetry run python -m airweave.db.bootstrap +``` + +The second command installs native destination definitions and the initial embedding +metadata. Ordinary startup validates that metadata without inserting it. Service +mode does not require FIRST_SUPERUSER or FIRST_SUPERUSER_PASSWORD; these are +local bootstrap inputs only. Provision the +intended service organization and key through the existing model/CRUD operator +workflow, store the key in secret management, then start the API and workers. +Do not print keys in logs. Runtime replicas should not need migration privileges. +Temporal database setup is separate and must also target isolated resources. + +For a disposable local demo only: + +```sh +AUTH_MODE=local poetry run python -m airweave.db.bootstrap --local-superuser +``` + +This opts into the legacy demo user/org/key provisioning. It is refused in +api_key or auth0 mode. The production entrypoint starts Uvicorn without reload. +Almanac/WorkOS remains end-user identity authority; this service does not +introduce an independent customer login system. diff --git a/backend/airweave/adapters/identity/auth0.py b/backend/airweave/adapters/identity/auth0.py index f271f44614..3fa63d8064 100644 --- a/backend/airweave/adapters/identity/auth0.py +++ b/backend/airweave/adapters/identity/auth0.py @@ -17,7 +17,7 @@ wait_exponential, ) -from airweave.core.config import settings +from airweave.core.config import AuthMode, settings from airweave.core.logging import logger from airweave.core.protocols.identity import ( IdentityProvider, @@ -416,7 +416,7 @@ async def add_enabled_connection_to_organization( # --------------------------------------------------------------------------- auth0_management_client: Optional[Auth0ManagementClient] = None -if settings.AUTH_ENABLED: +if settings.AUTH_MODE == AuthMode.AUTH0: auth0_management_client = Auth0ManagementClient() diff --git a/backend/airweave/api/auth.py b/backend/airweave/api/auth.py index d98b35eaed..b4b3ab9518 100644 --- a/backend/airweave/api/auth.py +++ b/backend/airweave/api/auth.py @@ -5,7 +5,7 @@ from fastapi_auth0 import Auth0, Auth0User from jose import jwt -from airweave.core.config import settings +from airweave.core.config import AuthMode, settings # Add a method to auth0 instance to verify tokens directly @@ -19,9 +19,12 @@ async def get_user_from_token(token: str): Auth0User if token is valid, None otherwise """ # If auth is disabled, just return a mock user - if not settings.AUTH_ENABLED: + if settings.AUTH_MODE == AuthMode.LOCAL: return Auth0User(sub="mock-user-id", email=settings.FIRST_SUPERUSER) + if settings.AUTH_MODE != AuthMode.AUTH0: + return None + try: if not token: return None @@ -63,7 +66,7 @@ async def get_user_from_token(token: str): # Initialize Auth0 only if authentication is enabled -if settings.AUTH_ENABLED: +if settings.AUTH_MODE == AuthMode.AUTH0: auth0 = Auth0( domain=settings.AUTH0_DOMAIN, api_audience=settings.AUTH0_AUDIENCE, @@ -71,11 +74,11 @@ async def get_user_from_token(token: str): ) else: # Create a mock Auth0 instance that doesn't make network calls - class MockAuth0: - """A mock Auth0 class that doesn't make network calls for testing/development.""" + class InactiveAuth0: + """Dependency for modes without Auth0; only explicit local mode returns a demo user.""" def __init__(self): - """Initialize the mock Auth0 instance.""" + """Initialize a network-free Auth0 dependency.""" self.domain = "mock-domain.auth0.com" self.audience = "https://mock-api/" self.algorithms = ["RS256"] @@ -83,9 +86,11 @@ def __init__(self): self.auth0_user_model = Auth0User async def get_user(self): - """Always return a mock user in development mode.""" + """Return a demo user only when insecure local mode is explicit.""" # For development and testing - return Auth0User(sub="mock-user-id", email=settings.FIRST_SUPERUSER) + if settings.AUTH_MODE == AuthMode.LOCAL: + return Auth0User(sub="mock-user-id", email=settings.FIRST_SUPERUSER) + return None - auth0 = MockAuth0() - logging.info("Using mock Auth0 instance because AUTH_ENABLED=False") + auth0 = InactiveAuth0() + logging.info("Auth0 disabled for authentication mode %s", settings.AUTH_MODE.value) diff --git a/backend/airweave/api/context_resolver.py b/backend/airweave/api/context_resolver.py index d1d5128700..edf954eba2 100644 --- a/backend/airweave/api/context_resolver.py +++ b/backend/airweave/api/context_resolver.py @@ -18,8 +18,12 @@ from airweave import schemas from airweave.analytics.service import analytics from airweave.api.context import ApiContext, RequestHeaders -from airweave.core.config import settings -from airweave.core.exceptions import NotFoundException, RateLimitExceededException +from airweave.core.config import AuthMode, settings +from airweave.core.exceptions import ( + NotFoundException, + PermissionException, + RateLimitExceededException, +) from airweave.core.logging import logger from airweave.core.protocols.cache import ContextCache from airweave.core.protocols.rate_limiter import RateLimiter @@ -105,14 +109,14 @@ async def authenticate_user_only( Used by ``deps.get_user``. """ - if not settings.AUTH_ENABLED: + if settings.AUTH_MODE == AuthMode.LOCAL: user = await self._fetch_system_user(db) if user: return user raise HTTPException(status_code=401, detail="User not found") - if not auth0_user: - raise HTTPException(status_code=401, detail="User email not found in Auth0") + if settings.AUTH_MODE != AuthMode.AUTH0 or not auth0_user: + raise HTTPException(status_code=401, detail="User authentication is required") user = await self._fetch_auth0_user(db, auth0_user) if not user: @@ -130,9 +134,9 @@ async def _authenticate( x_api_key: Optional[str], request: Request, ) -> AuthResult: - if not settings.AUTH_ENABLED: + if settings.AUTH_MODE == AuthMode.LOCAL: return await self._authenticate_system(db) - if auth0_user: + if settings.AUTH_MODE == AuthMode.AUTH0 and auth0_user: return await self._authenticate_auth0(db, auth0_user) if x_api_key: return await self._authenticate_api_key(db, x_api_key, request) @@ -165,14 +169,6 @@ async def _authenticate_api_key( self, db: AsyncSession, api_key: str, request: Request ) -> AuthResult: try: - org_id = await self._cache.get_api_key_org_id(api_key) - if org_id: - return AuthResult( - method=AuthMethod.API_KEY, - metadata={"api_key_id": "cached", "created_by": None}, - api_key_org_id=str(org_id), - ) - api_key_obj = await self._api_keys.get_by_key(db, key=api_key) org_id = api_key_obj.organization_id @@ -183,8 +179,6 @@ async def _authenticate_api_key( f"endpoint={request.url.path} created_by={api_key_obj.created_by_email}" ) - await self._cache.set_api_key_org_id(api_key, org_id) - return AuthResult( method=AuthMethod.API_KEY, metadata={ @@ -194,7 +188,7 @@ async def _authenticate_api_key( api_key_org_id=str(org_id), ) - except (ValueError, NotFoundException) as e: + except (ValueError, NotFoundException, PermissionException) as e: logger.error(f"API key validation failed: {e}") if "expired" in str(e): raise HTTPException(status_code=403, detail="API key has expired") from e @@ -291,8 +285,7 @@ async def _validate_organization_access( ) elif auth.method == AuthMethod.API_KEY and x_api_key: - api_key_obj = await self._api_keys.get_by_key(db, key=x_api_key) - if str(api_key_obj.organization_id) != organization_id: + if auth.api_key_org_id != organization_id: raise HTTPException( status_code=403, detail=f"API key does not have access to organization {organization_id}", diff --git a/backend/airweave/api/deps.py b/backend/airweave/api/deps.py index 1e0b305910..92211f7edf 100644 --- a/backend/airweave/api/deps.py +++ b/backend/airweave/api/deps.py @@ -17,13 +17,16 @@ from airweave.api.context_resolver import ContextResolver from airweave.api.inject import Inject # noqa: F401 — re-exported for backward compat from airweave.core import container as container_mod -from airweave.core.config import settings +from airweave.core.config import AuthMode, settings from airweave.core.container import Container from airweave.core.logging import ContextualLogger from airweave.core.protocols.cache import ContextCache from airweave.core.protocols.rate_limiter import RateLimiter from airweave.core.shared_models import AuthMethod 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.store import CanonicalRecordStore from airweave.domains.organizations.repository import ApiKeyRepository, OrganizationRepository from airweave.domains.users.repository import UserRepository @@ -139,7 +142,7 @@ async def get_user_from_token(token: str, db: AsyncSession) -> Optional[schemas. if token.startswith("Bearer "): token = token[7:] - if not settings.AUTH_ENABLED: + if settings.AUTH_MODE == AuthMode.LOCAL: user = await crud.user.get_by_email(db, email=settings.FIRST_SUPERUSER) if user: return schemas.User.model_validate(user) @@ -230,3 +233,10 @@ async def get_connect_session( ) except (KeyError, ValueError) as e: raise HTTPException(status_code=401, detail="Invalid session token payload") from e + + +def get_canonical_query_service() -> CanonicalQueryService: + """Wire stateless exact-record reads with the deployment's cursor signing key.""" + return CanonicalQueryService( + CanonicalRecordStore(), CanonicalQueryStore(), settings.STATE_SECRET + ) diff --git a/backend/airweave/api/tests/test_context_resolver_auth.py b/backend/airweave/api/tests/test_context_resolver_auth.py index cae054848d..0d633b9082 100644 --- a/backend/airweave/api/tests/test_context_resolver_auth.py +++ b/backend/airweave/api/tests/test_context_resolver_auth.py @@ -4,7 +4,7 @@ - Auth method dispatch (system / Auth0 / API key / none) - Auth0 cache hit returns cached user without DB call - Auth0 cache miss populates cache after DB call -- API key cache hit returns org_id without DB call +- Cached key mappings never bypass credential validation - No auth raises 401 """ @@ -38,14 +38,17 @@ class TestAuthDispatch: @pytest.mark.asyncio @patch("airweave.api.context_resolver.settings") async def test_auth_disabled_uses_system(self, mock_settings): - mock_settings.AUTH_ENABLED = False + mock_settings.AUTH_MODE = "local" mock_settings.FIRST_SUPERUSER = "admin@test.com" resolver = _make_resolver() with patch.object(resolver, "_authenticate_system", new_callable=AsyncMock) as mock_sys: mock_sys.return_value = AuthResult(method=AuthMethod.SYSTEM) result = await resolver._authenticate( - db=AsyncMock(), auth0_user=None, x_api_key=None, request=MagicMock(), + db=AsyncMock(), + auth0_user=None, + x_api_key=None, + request=MagicMock(), ) mock_sys.assert_called_once() assert result.method == AuthMethod.SYSTEM @@ -53,7 +56,7 @@ async def test_auth_disabled_uses_system(self, mock_settings): @pytest.mark.asyncio @patch("airweave.api.context_resolver.settings") async def test_auth0_user_takes_priority_over_api_key(self, mock_settings): - mock_settings.AUTH_ENABLED = True + mock_settings.AUTH_MODE = "auth0" resolver = _make_resolver() auth0_user = MagicMock(email="user@test.com", id="auth0|123") @@ -61,7 +64,10 @@ async def test_auth0_user_takes_priority_over_api_key(self, mock_settings): with patch.object(resolver, "_authenticate_auth0", new_callable=AsyncMock) as mock_auth0: mock_auth0.return_value = AuthResult(method=AuthMethod.AUTH0) result = await resolver._authenticate( - db=AsyncMock(), auth0_user=auth0_user, x_api_key="some-key", request=MagicMock(), + db=AsyncMock(), + auth0_user=auth0_user, + x_api_key="some-key", + request=MagicMock(), ) mock_auth0.assert_called_once() assert result.method == AuthMethod.AUTH0 @@ -69,15 +75,19 @@ async def test_auth0_user_takes_priority_over_api_key(self, mock_settings): @pytest.mark.asyncio @patch("airweave.api.context_resolver.settings") async def test_api_key_used_when_no_auth0_user(self, mock_settings): - mock_settings.AUTH_ENABLED = True + mock_settings.AUTH_MODE = "auth0" resolver = _make_resolver() with patch.object(resolver, "_authenticate_api_key", new_callable=AsyncMock) as mock_api: mock_api.return_value = AuthResult( - method=AuthMethod.API_KEY, api_key_org_id=str(ORG_ID), + method=AuthMethod.API_KEY, + api_key_org_id=str(ORG_ID), ) result = await resolver._authenticate( - db=AsyncMock(), auth0_user=None, x_api_key="key-123", request=MagicMock(), + db=AsyncMock(), + auth0_user=None, + x_api_key="key-123", + request=MagicMock(), ) mock_api.assert_called_once() assert result.method == AuthMethod.API_KEY @@ -85,12 +95,15 @@ async def test_api_key_used_when_no_auth0_user(self, mock_settings): @pytest.mark.asyncio @patch("airweave.api.context_resolver.settings") async def test_no_auth_raises_401(self, mock_settings): - mock_settings.AUTH_ENABLED = True + mock_settings.AUTH_MODE = "auth0" resolver = _make_resolver() with pytest.raises(HTTPException) as exc: await resolver._authenticate( - db=AsyncMock(), auth0_user=None, x_api_key=None, request=MagicMock(), + db=AsyncMock(), + auth0_user=None, + x_api_key=None, + request=MagicMock(), ) assert exc.value.status_code == 401 @@ -151,19 +164,22 @@ async def test_cache_miss_db_returns_none(self): assert "ghost@test.com" not in cache._users -class TestApiKeyCacheIntegration: - """Verify API key auth uses the cache correctly.""" +class TestApiKeyValidation: + """A cached org mapping is not authorization after key revocation.""" @pytest.mark.asyncio - async def test_cache_hit_returns_org_id(self): + async def test_stale_cache_does_not_authorize_revoked_key(self): + from airweave.core.exceptions import NotFoundException + cache = FakeContextCache() cache._api_keys["my-secret-key"] = ORG_ID - resolver = _make_resolver(cache=cache) - result = await resolver._authenticate_api_key( - db=AsyncMock(), api_key="my-secret-key", request=MagicMock(), - ) - - assert result.method == AuthMethod.API_KEY - assert result.api_key_org_id == str(ORG_ID) - assert result.metadata["api_key_id"] == "cached" + resolver._api_keys.get_by_key = AsyncMock(side_effect=NotFoundException("revoked")) + with pytest.raises(HTTPException) as error: + await resolver._authenticate_api_key( + db=AsyncMock(), + api_key="my-secret-key", + request=MagicMock(), + ) + assert error.value.status_code == 403 + resolver._api_keys.get_by_key.assert_awaited_once() diff --git a/backend/airweave/api/tests/test_service_auth.py b/backend/airweave/api/tests/test_service_auth.py new file mode 100644 index 0000000000..850bbfccee --- /dev/null +++ b/backend/airweave/api/tests/test_service_auth.py @@ -0,0 +1,177 @@ +"""Service-key mode never falls back to mock users or stale credential caches.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from airweave.api.context_resolver import AuthResult, ContextResolver +from airweave.core.config import AuthMode, settings +from airweave.core.config.settings import Settings +from airweave.core.exceptions import NotFoundException, PermissionException +from airweave.core.shared_models import AuthMethod + + +def resolver(): + return ContextResolver( + cache=MagicMock(), + rate_limiter=MagicMock(), + user_repo=MagicMock(), + api_key_repo=SimpleNamespace(get_by_key=AsyncMock()), + org_repo=MagicMock(), + ) + + +def test_auth_mode_defaults_and_legacy_migration(): + assert Settings.model_fields["AUTH_MODE"].default == AuthMode.API_KEY + assert Settings.migrate_auth_mode({"AUTH_ENABLED": True})["AUTH_MODE"] == "auth0" + with pytest.raises(ValueError, match="conflicts"): + Settings.migrate_auth_mode({"AUTH_ENABLED": False, "AUTH_MODE": "api_key"}) + + +def test_local_mode_cannot_run_in_production(): + configuration = settings.model_dump() + configuration.update(AUTH_MODE="local", AUTH_ENABLED=None, ENVIRONMENT="prd") + with pytest.raises(ValueError, match="only allowed"): + Settings(**configuration) + + +@pytest.mark.asyncio +async def test_api_key_mode_rejects_missing_key_and_user_auth(monkeypatch): + monkeypatch.setattr(settings, "AUTH_MODE", AuthMode.API_KEY) + service = resolver() + with pytest.raises(HTTPException) as error: + await service._authenticate( + MagicMock(), SimpleNamespace(email="fake@example.com"), None, MagicMock() + ) + assert error.value.status_code == 401 + with pytest.raises(HTTPException) as error: + await service.authenticate_user_only(MagicMock(), SimpleNamespace(email="fake@example.com")) + assert error.value.status_code == 401 + service._users.get_by_email.assert_not_called() + + +@pytest.mark.asyncio +async def test_key_validated_each_request_and_real_identity_preserved(): + service = resolver() + org, key = uuid4(), uuid4() + service._api_keys.get_by_key.return_value = SimpleNamespace( + organization_id=org, id=key, created_by_email="operator@example.com" + ) + request = MagicMock() + request.url.path = "/search" + first = await service._authenticate_api_key(MagicMock(), "opaque-key", request) + assert first.api_key_org_id == str(org) + assert first.metadata["api_key_id"] == str(key) + service._api_keys.get_by_key.side_effect = NotFoundException("revoked") + with pytest.raises(HTTPException) as error: + await service._authenticate_api_key(MagicMock(), "opaque-key", request) + assert error.value.status_code == 403 + assert service._api_keys.get_by_key.await_count == 2 + service._cache.get_api_key_org_id.assert_not_called() + + +@pytest.mark.asyncio +async def test_expired_key_maps_to_auth_error(): + service = resolver() + service._api_keys.get_by_key.side_effect = PermissionException("API key has expired") + with pytest.raises(HTTPException) as error: + await service._authenticate_api_key(MagicMock(), "opaque-key", MagicMock()) + assert error.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_key_cannot_override_organization(): + service = resolver() + auth = AuthResult(method=AuthMethod.API_KEY, api_key_org_id=str(uuid4())) + with pytest.raises(HTTPException) as error: + await service._validate_organization_access(MagicMock(), str(uuid4()), auth, "opaque-key") + assert error.value.status_code == 403 + service._api_keys.get_by_key.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_api_key_mode_user_dependencies_never_supply_demo_identity(monkeypatch): + from airweave.api import auth, deps + + monkeypatch.setattr(settings, "AUTH_MODE", AuthMode.API_KEY) + assert await auth.get_user_from_token("anything") is None + assert await auth.auth0.get_user() is None + assert await deps.get_user_from_token("Bearer anything", MagicMock()) is None + + +@pytest.mark.asyncio +async def test_service_bootstrap_refuses_demo_superuser(monkeypatch): + from airweave.db import bootstrap + + monkeypatch.setattr(settings, "AUTH_MODE", AuthMode.API_KEY) + session_factory = MagicMock() + monkeypatch.setattr(bootstrap, "AsyncSessionLocal", session_factory) + with pytest.raises(ValueError, match="requires AUTH_MODE=local"): + await bootstrap.bootstrap(local_superuser=True) + session_factory.assert_not_called() + + +def test_api_key_configuration_does_not_require_auth0(): + configuration = settings.model_dump() + configuration.update(AUTH_MODE="api_key", AUTH_ENABLED=None) + for field in configuration: + if field.startswith("AUTH0_"): + configuration[field] = None + assert Settings(**configuration).AUTH_MODE == AuthMode.API_KEY + configuration["AUTH_MODE"] = "auth0" + with pytest.raises(ValueError, match="must be set when AUTH_MODE=auth0"): + Settings(**configuration) + + +@pytest.mark.parametrize("value", [True, "true", "yes", "on", "1"]) +def test_legacy_truthy_auth_never_becomes_local(value): + assert Settings.migrate_auth_mode({"AUTH_ENABLED": value})["AUTH_MODE"] == AuthMode.AUTH0 + + +@pytest.mark.parametrize("value", ["typo", "", "sometimes"]) +def test_invalid_legacy_auth_rejected(value): + with pytest.raises(ValueError): + Settings.migrate_auth_mode({"AUTH_ENABLED": value}) + + +def test_service_settings_require_no_demo_credentials(): + configuration = settings.model_dump() + configuration.update( + AUTH_MODE="api_key", + AUTH_ENABLED=None, + FIRST_SUPERUSER=None, + FIRST_SUPERUSER_PASSWORD=None, + ) + result = Settings(**configuration) + assert result.FIRST_SUPERUSER is None + assert result.FIRST_SUPERUSER_PASSWORD is None + + +@pytest.mark.asyncio +async def test_startup_missing_embedding_metadata_never_seeds(): + from airweave.domains.embedders import config + + result = MagicMock() + result.scalar_one_or_none.return_value = None + db = MagicMock(execute=AsyncMock(return_value=result), commit=AsyncMock()) + with pytest.raises(config.EmbeddingConfigError, match="bootstrap"): + await config.validate_embedding_config(db) + db.add.assert_not_called() + db.commit.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_explicit_bootstrap_initializes_embedding_metadata(): + from airweave.domains.embedders import config + + result = MagicMock() + result.scalar_one_or_none.return_value = None + db = MagicMock(execute=AsyncMock(return_value=result), commit=AsyncMock()) + await config.initialize_embedding_config(db) + row = db.add.call_args.args[0] + assert row.dense_embedder == config.DENSE_EMBEDDER + assert row.embedding_dimensions == config.EMBEDDING_DIMENSIONS + db.commit.assert_awaited_once() diff --git a/backend/airweave/api/tests/test_sync_authorization.py b/backend/airweave/api/tests/test_sync_authorization.py new file mode 100644 index 0000000000..0463411295 --- /dev/null +++ b/backend/airweave/api/tests/test_sync_authorization.py @@ -0,0 +1,39 @@ +"""An authenticated key cannot subscribe to another organization's job topics.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from airweave.api.v1.endpoints import sync + + +@pytest.mark.asyncio +@pytest.mark.parametrize("handler", [sync.subscribe_sync_job, sync.subscribe_entity_state]) +async def test_foreign_job_never_subscribes(monkeypatch, handler): + owner, foreign_owner, job_id = uuid4(), uuid4(), uuid4() + db = object() + ctx = SimpleNamespace(organization=SimpleNamespace(id=owner)) + + async def scoped_get(session, *, id, ctx): + assert session is db + assert id == job_id + return SimpleNamespace(id=id) if ctx.organization.id == foreign_owner else None + + monkeypatch.setattr(sync.crud.sync_job, "get", scoped_get) + pubsub = SimpleNamespace(subscribe=AsyncMock()) + with pytest.raises(HTTPException) as error: + await handler(job_id=job_id, ctx=ctx, db=db, pubsub=pubsub) + assert error.value.status_code == 404 + pubsub.subscribe.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authorized_job_passes_context_to_scoped_repository(monkeypatch): + get = AsyncMock(return_value=object()) + monkeypatch.setattr(sync.crud.sync_job, "get", get) + db, ctx, job_id = object(), object(), uuid4() + await sync.authorize_job(db, ctx, job_id) + get.assert_awaited_once_with(db, id=job_id, ctx=ctx) diff --git a/backend/airweave/api/v1/api.py b/backend/airweave/api/v1/api.py index 70b86c8d8c..0d1f5ae7eb 100644 --- a/backend/airweave/api/v1/api.py +++ b/backend/airweave/api/v1/api.py @@ -14,6 +14,7 @@ file_retrieval, health, organizations, + records, search, search_legacy, source_connections, @@ -47,6 +48,7 @@ source_rate_limits.router, prefix="/source-rate-limits", tags=["source-rate-limits"] ) api_router.include_router(sync.router, prefix="/sync", tags=["sync"]) +api_router.include_router(records.router, prefix="/sync", tags=["records"]) api_router.include_router(entities.router, prefix="/entities", tags=["entities"]) api_router.include_router(entity_counts.router, prefix="/entity-counts", tags=["entity-counts"]) api_router.include_router(file_retrieval.router, prefix="/files", tags=["files"]) diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py new file mode 100644 index 0000000000..a4b752e44a --- /dev/null +++ b/backend/airweave/api/v1/endpoints/records.py @@ -0,0 +1,79 @@ +"""Internal service API for committed source records; provider APIs are never read here.""" + +from typing import Literal +from uuid import UUID + +from fastapi import Depends, Query, Request +from fastapi.responses import JSONResponse +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api import deps +from airweave.api.context import ApiContext +from airweave.api.router import TrailingSlashRouter +from airweave.db.session import get_db +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.query import CanonicalQueryService +from airweave.domains.entities.canonical.query_models import ( + RecordChangePage, + RecordFilters, + RecordListQuery, + RecordPage, +) +from airweave.domains.entities.canonical.store import CanonicalStoreError + +router = TrailingSlashRouter() + + +async def record_error_response(request: Request, error: CanonicalStoreError) -> JSONResponse: + """Preserve machine-readable recovery codes without exposing SQL/provider diagnostics.""" + status = 404 if error.code in {"source_not_found", "record_not_found"} else 400 + return JSONResponse( + status_code=status, + content={"error": {"code": error.code, "message": str(error), "retryable": False}}, + ) + + +@router.get("/{sync_id}/records", response_model=RecordPage) +async def list_records( + sync_id: UUID, + record_type: str | None = None, + container_id: str | None = None, + state: Literal["active", "deleted", "all"] = "active", + cursor: str | None = None, + limit: int = Query(100, ge=1, le=500), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), +) -> RecordPage: + """List exact stored records in stable ID order; continuation is a live traversal.""" + query = RecordListQuery( + filters=RecordFilters(record_type=record_type, container_id=container_id, state=state), + cursor=cursor, + limit=limit, + ) + return await service.list_records(db, ctx.organization.id, sync_id, query) + + +@router.get("/{sync_id}/records/changes", response_model=RecordChangePage) +async def record_changes( + sync_id: UUID, + cursor: str | None = None, + limit: int = Query(100, ge=1, le=500), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), +) -> RecordChangePage: + """Read observed changes from the beginning, or resume a source-bound cursor.""" + return await service.changes(db, ctx.organization.id, sync_id, cursor=cursor, limit=limit) + + +@router.get("/{sync_id}/records/{record_id}", response_model=SourceRecord) +async def read_record( + sync_id: UUID, + record_id: UUID, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), +) -> SourceRecord: + """Read current committed provider state, including tombstones and completeness.""" + return await service.read(db, ctx.organization.id, sync_id, record_id) diff --git a/backend/airweave/api/v1/endpoints/sync.py b/backend/airweave/api/v1/endpoints/sync.py index 24add9764f..747c9b3ace 100644 --- a/backend/airweave/api/v1/endpoints/sync.py +++ b/backend/airweave/api/v1/endpoints/sync.py @@ -5,23 +5,33 @@ from typing import AsyncGenerator from uuid import UUID -from fastapi import Depends +from fastapi import Depends, HTTPException from fastapi.responses import StreamingResponse +from sqlalchemy.ext.asyncio import AsyncSession +from airweave import crud from airweave.api import deps from airweave.api.context import ApiContext from airweave.api.deps import Inject from airweave.api.router import TrailingSlashRouter from airweave.core.logging import logger from airweave.core.protocols import PubSub +from airweave.db.session import get_db router = TrailingSlashRouter() +async def authorize_job(db: AsyncSession, ctx: ApiContext, job_id: UUID) -> None: + """Authorize ownership before subscribing to an otherwise unscoped pubsub topic.""" + if await crud.sync_job.get(db, id=job_id, ctx=ctx) is None: + raise HTTPException(status_code=404, detail="Sync job not found") + + @router.get("/job/{job_id}/subscribe") async def subscribe_sync_job( job_id: UUID, ctx: ApiContext = Depends(deps.get_context), + db: AsyncSession = Depends(get_db), pubsub: PubSub = Inject(PubSub), ) -> StreamingResponse: """Server-Sent Events (SSE) endpoint to subscribe to a sync job's progress. @@ -30,12 +40,15 @@ async def subscribe_sync_job( ----- job_id: The ID of the job to subscribe to ctx: The API context + db: Database session for organization-scoped job authorization pubsub: PubSub adapter for event streaming Returns: -------- StreamingResponse: The streaming response """ + await authorize_job(db, ctx, job_id) + logger.info(f"SSE sync subscription authenticated for user: {ctx}, job: {job_id}") # Track active SSE connections @@ -91,6 +104,7 @@ async def event_stream() -> AsyncGenerator[str, None]: async def subscribe_entity_state( job_id: UUID, ctx: ApiContext = Depends(deps.get_context), + db: AsyncSession = Depends(get_db), pubsub: PubSub = Inject(PubSub), ) -> StreamingResponse: """SSE endpoint for total entity state updates during sync. @@ -103,12 +117,15 @@ async def subscribe_entity_state( ----- job_id: The ID of the job to subscribe to ctx: The API context + db: Database session for organization-scoped job authorization pubsub: PubSub adapter for event streaming Returns: -------- StreamingResponse: Server-sent events with entity state updates """ + await authorize_job(db, ctx, job_id) + logger.info(f"SSE entity state subscription for user: {ctx}, job: {job_id}") channel = f"sync_job_state:{job_id}" diff --git a/backend/airweave/core/config/__init__.py b/backend/airweave/core/config/__init__.py index f0b74a7163..8df18b9b71 100644 --- a/backend/airweave/core/config/__init__.py +++ b/backend/airweave/core/config/__init__.py @@ -14,11 +14,12 @@ ... """ -from airweave.core.config.enums import Environment, StorageBackendType +from airweave.core.config.enums import AuthMode, Environment, StorageBackendType from airweave.core.config.settings import MISTRAL_DEFAULT_BASE_URL, Settings __all__ = [ "Settings", + "AuthMode", "StorageBackendType", "Environment", "settings", diff --git a/backend/airweave/core/config/enums.py b/backend/airweave/core/config/enums.py index f2b4688f7d..2bc7e33bb4 100644 --- a/backend/airweave/core/config/enums.py +++ b/backend/airweave/core/config/enums.py @@ -30,3 +30,11 @@ class Environment(str, Enum): TEST = "test" DEV = "dev" PRD = "prd" + + +class AuthMode(str, Enum): + """Authentication authority; local explicitly opts into insecure development.""" + + API_KEY = "api_key" + AUTH0 = "auth0" + LOCAL = "local" diff --git a/backend/airweave/core/config/settings.py b/backend/airweave/core/config/settings.py index c4fbc8db35..080cbe5a2b 100644 --- a/backend/airweave/core/config/settings.py +++ b/backend/airweave/core/config/settings.py @@ -5,10 +5,10 @@ from typing import Optional -from pydantic import PostgresDsn, ValidationInfo, field_validator +from pydantic import PostgresDsn, TypeAdapter, ValidationInfo, field_validator, model_validator from pydantic_settings import BaseSettings -from airweave.core.config.enums import Environment, StorageBackendType +from airweave.core.config.enums import AuthMode, Environment, StorageBackendType _BANNED_PASSWORDS: frozenset[str] = frozenset( { @@ -63,9 +63,6 @@ class Settings(BaseSettings): POSTGRES_PASSWORD (str): The PostgreSQL password. SQLALCHEMY_ASYNC_DATABASE_URI (Optional[PostgresDsn]): The SQLAlchemy async database URI. LOCAL_NGROK_SERVER (Optional[str]): The local ngrok server URL. - RUN_ALEMBIC_MIGRATIONS (bool): Whether to run the alembic migrations. - RUN_DB_SYNC (bool): Whether to run the system sync to process sources, - destinations, and entity types. REDIS_HOST (str): The Redis server hostname (single-host mode). REDIS_PORT (int): The Redis server port (single-host mode). REDIS_PASSWORD (Optional[str]): The Redis password (if authentication is enabled). @@ -108,11 +105,12 @@ class Settings(BaseSettings): # Rate limiting DISABLE_RATE_LIMIT: bool = False # For testing purposes - disables rate limiting completely - FIRST_SUPERUSER: str - FIRST_SUPERUSER_PASSWORD: str + FIRST_SUPERUSER: Optional[str] = None + FIRST_SUPERUSER_PASSWORD: Optional[str] = None FIRST_SUPERUSER_NAME: str = "Admin" - AUTH_ENABLED: Optional[bool] = False + AUTH_MODE: AuthMode = AuthMode.API_KEY + AUTH_ENABLED: Optional[bool] = None # Deprecated migration input; use AUTH_MODE. AUTH0_DOMAIN: Optional[str] = None AUTH0_AUDIENCE: Optional[str] = None AUTH0_RULE_NAMESPACE: Optional[str] = None @@ -144,8 +142,6 @@ class Settings(BaseSettings): LOCAL_NGROK_SERVER: Optional[str] = None - RUN_ALEMBIC_MIGRATIONS: bool = True - RUN_DB_SYNC: bool = True ENABLE_INTERNAL_SOURCES: bool = False # Enable internal/testing sources (stub, snapshot) # Redis configuration. @@ -358,6 +354,32 @@ def parse_cors_origins(cls, v: Optional[str]) -> Optional[list[str]]: return [origin.strip() for origin in v.split(",") if origin.strip()] + @model_validator(mode="before") + @classmethod + def migrate_auth_mode(cls, values): + """Preserve explicit legacy configuration without retaining two authorities.""" + values = dict(values) + legacy = values.get("AUTH_ENABLED") + if legacy is not None: + enabled = TypeAdapter(bool).validate_python(legacy) + inferred = AuthMode.AUTH0 if enabled else AuthMode.LOCAL + if "AUTH_MODE" in values and values["AUTH_MODE"] != inferred: + raise ValueError("AUTH_ENABLED conflicts with AUTH_MODE; remove AUTH_ENABLED") + values["AUTH_MODE"] = inferred + return values + + @model_validator(mode="after") + def validate_auth_mode(self): + """Insecure authentication is never available in a deployed environment.""" + if self.AUTH_MODE == AuthMode.LOCAL and self.ENVIRONMENT not in { + Environment.LOCAL, + Environment.TEST, + }: + raise ValueError("AUTH_MODE=local is only allowed in local/test environments") + if self.AUTH_MODE == AuthMode.LOCAL and not self.FIRST_SUPERUSER: + raise ValueError("FIRST_SUPERUSER is required for AUTH_MODE=local") + return self + @field_validator( "AUTH0_DOMAIN", "AUTH0_AUDIENCE", @@ -368,7 +390,7 @@ def parse_cors_origins(cls, v: Optional[str]) -> Optional[list[str]]: mode="before", ) def validate_auth0_settings(cls, v: str, info: ValidationInfo) -> str: - """Validate Auth0 settings when AUTH_ENABLED is True. + """Validate Auth0 settings only for AUTH_MODE=auth0. Args: @@ -382,12 +404,12 @@ def validate_auth0_settings(cls, v: str, info: ValidationInfo) -> str: Raises: ------ - ValueError: If AUTH_ENABLED is True and the Auth0 setting is empty. + ValueError: If AUTH_MODE=auth0 and the Auth0 setting is empty. """ - auth_enabled = info.data.get("AUTH_ENABLED", False) + auth_enabled = info.data.get("AUTH_MODE") == AuthMode.AUTH0 if auth_enabled and not v: field_name = info.field_name - raise ValueError(f"{field_name} must be set when AUTH_ENABLED is True") + raise ValueError(f"{field_name} must be set when AUTH_MODE=auth0") return v @field_validator( @@ -458,8 +480,12 @@ def validate_svix_jwt_secret(cls, v: str) -> str: return v @field_validator("FIRST_SUPERUSER_PASSWORD", mode="before") - def validate_first_superuser_password(cls, v: str, info: ValidationInfo) -> str: - """Reject weak superuser passwords in non-local environments.""" + def validate_first_superuser_password( + cls, v: Optional[str], info: ValidationInfo + ) -> Optional[str]: + """Validate bootstrap passwords when supplied; service mode needs none.""" + if v is None: + return None environment = info.data.get("ENVIRONMENT", "local") env_str = environment.value if isinstance(environment, Environment) else environment if env_str in ("local", "test"): @@ -476,7 +502,9 @@ def validate_first_superuser_password(cls, v: str, info: ValidationInfo) -> str: return v @field_validator("FIRST_SUPERUSER", mode="before") - def validate_first_superuser_email(cls, v: str, info: ValidationInfo) -> str: + def validate_first_superuser_email( + cls, v: Optional[str], info: ValidationInfo + ) -> Optional[str]: """Reject well-known placeholder emails in non-local environments.""" environment = info.data.get("ENVIRONMENT", "local") env_str = environment.value if isinstance(environment, Environment) else environment diff --git a/backend/airweave/core/config/tests/test_settings_validators.py b/backend/airweave/core/config/tests/test_settings_validators.py index e3eee42504..399871b833 100644 --- a/backend/airweave/core/config/tests/test_settings_validators.py +++ b/backend/airweave/core/config/tests/test_settings_validators.py @@ -8,6 +8,7 @@ # can override *only* the field under test and trigger the expected error. _BASE_ENV: dict[str, str] = { "ENVIRONMENT": "local", + "AUTH_MODE": "api_key", "FIRST_SUPERUSER": "test@example.com", "FIRST_SUPERUSER_PASSWORD": "testpassword123", "ENCRYPTION_KEY": "SpgLrrEEgJ/7QdhSMSvagL1juEY5eoyCG0tZN7OSQV0=", @@ -21,6 +22,8 @@ def _build_settings(monkeypatch, overrides: dict[str, str]): """Set env vars from _BASE_ENV + overrides, then construct Settings.""" + # These field-validator tests must not inherit insecure local auth from conftest. + monkeypatch.delenv("AUTH_ENABLED", raising=False) for key, value in {**_BASE_ENV, **overrides}.items(): monkeypatch.setenv(key, value) diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index a1cc47bb54..94ce802393 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -1079,7 +1079,7 @@ def _create_usage_ledger(settings: Settings, billing_deps: dict) -> UsageLedgerP def _create_identity_provider(settings: Settings) -> IdentityProvider: """Create identity provider: Auth0 if enabled, otherwise null implementation.""" - if settings.AUTH_ENABLED: + if settings.AUTH_MODE == "auth0": from airweave.adapters.identity.auth0 import auth0_management_client if auth0_management_client: diff --git a/backend/airweave/db/bootstrap.py b/backend/airweave/db/bootstrap.py new file mode 100644 index 0000000000..15909e426b --- /dev/null +++ b/backend/airweave/db/bootstrap.py @@ -0,0 +1,41 @@ +"""Explicit database setup: run after reviewed Alembic migrations, never on API startup. + +Usage: python -m airweave.db.bootstrap +Local demo only: AUTH_MODE=local python -m airweave.db.bootstrap --local-superuser +Service organization/key provisioning remains an explicit operator action. +""" + +import argparse +import asyncio + +from airweave.core.config import AuthMode, settings +from airweave.db.init_db import init_db +from airweave.db.init_db_native import init_db_with_native_connections +from airweave.db.session import AsyncSessionLocal +from airweave.domains.embedders.config import initialize_embedding_config + + +async def bootstrap(*, local_superuser: bool = False) -> None: + """Install native definitions, optionally opting into local demo identity.""" + if local_superuser and settings.AUTH_MODE != AuthMode.LOCAL: + raise ValueError("--local-superuser requires AUTH_MODE=local") + if local_superuser and not (settings.FIRST_SUPERUSER and settings.FIRST_SUPERUSER_PASSWORD): + raise ValueError("Local bootstrap requires FIRST_SUPERUSER and FIRST_SUPERUSER_PASSWORD") + async with AsyncSessionLocal() as db: + await initialize_embedding_config(db) + if local_superuser: + await init_db(db) + else: + await init_db_with_native_connections(db) + + +def main() -> None: + """Run the explicitly requested bootstrap operation.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--local-superuser", action="store_true") + args = parser.parse_args() + asyncio.run(bootstrap(local_superuser=args.local_superuser)) + + +if __name__ == "__main__": + main() diff --git a/backend/airweave/domains/embedders/config.py b/backend/airweave/domains/embedders/config.py index f5170a3221..125c20f3a4 100644 --- a/backend/airweave/domains/embedders/config.py +++ b/backend/airweave/domains/embedders/config.py @@ -107,6 +107,13 @@ async def validate_embedding_config(db: AsyncSession) -> None: await _reconcile_db(db) +async def initialize_embedding_config(db: AsyncSession) -> None: + """Persist first-deployment metadata only from explicit database bootstrap.""" + if not DENSE_EMBEDDER or not EMBEDDING_DIMENSIONS or not SPARSE_EMBEDDER: + raise EmbeddingConfigError("Set dense/sparse embedder and dimensions before bootstrap") + await _reconcile_db(db, initialize=True) + + def _validate_dimensions(dense_spec: DenseEmbedderEntry) -> None: """Validate EMBEDDING_DIMENSIONS against the dense embedder spec.""" if dense_spec.supports_matryoshka: @@ -175,12 +182,17 @@ def _validate_local_reachability(dense_spec: DenseEmbedderEntry) -> None: ) from exc -async def _reconcile_db(db: AsyncSession) -> None: +async def _reconcile_db(db: AsyncSession, *, initialize: bool = False) -> None: """Reconcile code config against the vector_db_deployment_metadata table.""" result = await db.execute(select(VectorDbDeploymentMetadata).limit(1)) row = result.scalar_one_or_none() if row is None: + if not initialize: + raise EmbeddingConfigError( + "Embedding metadata is missing; run python -m airweave.db.bootstrap " + "against the intended database before starting the service." + ) row = VectorDbDeploymentMetadata( dense_embedder=DENSE_EMBEDDER, embedding_dimensions=EMBEDDING_DIMENSIONS, diff --git a/backend/airweave/domains/entities/canonical/models.py b/backend/airweave/domains/entities/canonical/models.py new file mode 100644 index 0000000000..dc87978b4a --- /dev/null +++ b/backend/airweave/domains/entities/canonical/models.py @@ -0,0 +1,68 @@ +"""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.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 + revision: int + payload: dict[str, JsonValue] + 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 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 diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py new file mode 100644 index 0000000000..3e649aa924 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/query.py @@ -0,0 +1,135 @@ +"""Exact reads and continuation contracts over the owned record store.""" + +from uuid import UUID + +from jose import JWTError, jwt +from pydantic import ValidationError +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.query_models import ( + RecordChangePage, + RecordCursor, + RecordListQuery, + RecordPage, +) +from airweave.domains.entities.canonical.query_store import CanonicalQueryStore +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError + + +class InvalidRecordCursor(CanonicalStoreError): + """A cursor belongs to another scope/query or cannot be verified.""" + + code = "invalid_cursor" + + +class RecordNotFound(CanonicalStoreError): + """No canonical record is visible at this identity.""" + + code = "record_not_found" + + +class CanonicalQueryService: + """Own read semantics; cursor signatures never replace tenant authorization.""" + + def __init__( + self, records: CanonicalRecordStore, queries: CanonicalQueryStore, signing_key: str + ) -> None: + """Bind persistence and the existing deployment secret used for signed cursors.""" + self.records = records + self.queries = queries + self.signing_key = signing_key + + def _encode(self, cursor: RecordCursor) -> str: + return jwt.encode(cursor.model_dump(mode="json"), self.signing_key, algorithm="HS256") + + def _decode(self, token: str, organization_id: UUID, sync_id: UUID, mode: str) -> RecordCursor: + try: + cursor = RecordCursor.model_validate( + jwt.decode(token, self.signing_key, algorithms=["HS256"]) + ) + except (JWTError, ValidationError, ValueError) as exc: + raise InvalidRecordCursor("Invalid continuation cursor; restart this query") from exc + if ( + cursor.organization_id != organization_id + or cursor.sync_id != sync_id + or cursor.mode != mode + ): + raise InvalidRecordCursor("Cursor belongs to a different source or operation") + return cursor + + async def read( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID, record_id: UUID + ) -> SourceRecord: + """Return exact current state, including an explicit tombstone if deleted.""" + record = await self.records.read(db, organization_id, sync_id, record_id) + if record is None: + raise RecordNotFound("Record not found in this source") + return record + + async def list_records( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID, query: RecordListQuery + ) -> RecordPage: + """Traverse current committed rows; callers must preserve filters on continuation.""" + after_id = None + if query.cursor is not None: + cursor = self._decode(query.cursor, organization_id, sync_id, "list") + if cursor.filters != query.filters: + raise InvalidRecordCursor("Preserve the original filters when continuing a list") + after_id = cursor.after_id + rows = await self.queries.list_records( + db, organization_id, sync_id, query.filters, after_id=after_id, limit=query.limit + ) + more = len(rows) > query.limit + page = rows[: query.limit] + next_cursor = None + if more: + next_cursor = self._encode( + RecordCursor( + mode="list", + organization_id=organization_id, + sync_id=sync_id, + filters=query.filters, + after_id=page[-1].id, + ) + ) + return RecordPage(records=page, next_cursor=next_cursor, has_more=more) + + async def changes( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + *, + cursor: str | None = None, + limit: int = 100, + ) -> RecordChangePage: + """Start at the beginning; finish each committed window before polling the next.""" + position = ( + self._decode(cursor, organization_id, sync_id, "changes") + if cursor + else RecordCursor(mode="changes", organization_id=organization_id, sync_id=sync_id) + ) + page = await self.records.changes( + db, + organization_id, + sync_id, + after=position.after_sequence, + limit=limit, + high_watermark=position.high_watermark, + ) + next_cursor = self._encode( + RecordCursor( + mode="changes", + organization_id=organization_id, + sync_id=sync_id, + after_sequence=page.next_sequence, + high_watermark=page.high_watermark if page.has_more else None, + ) + ) + return RecordChangePage( + changes=page.changes, + next_cursor=next_cursor, + has_more=page.has_more, + high_watermark=page.high_watermark, + ) diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py new file mode 100644 index 0000000000..20a289e3f6 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -0,0 +1,62 @@ +"""Typed exact-record queries; search ranking is a separate index concern.""" + +from typing import Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord + + +class RecordFilters(BaseModel): + """Supported exact filters, preserved in every continuation token.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + record_type: str | None = None + container_id: str | None = None + state: Literal["active", "deleted", "all"] = "active" + + +class RecordListQuery(BaseModel): + """Stable-ID ordered live traversal, explicitly not a frozen snapshot.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + filters: RecordFilters = Field(default_factory=RecordFilters) + limit: int = Field(default=100, ge=1, le=500) + cursor: str | None = None + + +class RecordPage(BaseModel): + """Committed records with truthful traversal semantics.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + records: tuple[SourceRecord, ...] + next_cursor: str | None + has_more: bool + consistency: Literal["live"] = "live" + order: Literal["id_asc"] = "id_asc" + + +class RecordChangePage(BaseModel): + """A bounded journal window plus a cursor for future polling.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + changes: tuple[ObservedChange, ...] + next_cursor: str + has_more: bool + high_watermark: int + + +class RecordCursor(BaseModel): + """Versioned signed cursor payload; no provider secrets or record content.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + purpose: Literal["canonical_records"] = "canonical_records" + version: Literal[1] = 1 + mode: Literal["list", "changes"] + organization_id: UUID + sync_id: UUID + filters: RecordFilters | None = None + after_id: UUID | None = None + after_sequence: int = Field(default=0, ge=0) + high_watermark: int | None = Field(default=None, ge=0) diff --git a/backend/airweave/domains/entities/canonical/query_store.py b/backend/airweave/domains/entities/canonical/query_store.py new file mode 100644 index 0000000000..ec1073a9e6 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/query_store.py @@ -0,0 +1,50 @@ +"""Read-only SQL for exact canonical record listing.""" + +from uuid import UUID + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.query_models import RecordFilters +from airweave.domains.entities.canonical.store import SourceNotFound, source_record +from airweave.models.entity import Entity +from airweave.models.sync import Sync + + +class CanonicalQueryStore: + """No provider or index calls; reads only committed canonical rows.""" + + async def list_records( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + filters: RecordFilters, + *, + after_id: UUID | None, + limit: int, + ) -> tuple[SourceRecord, ...]: + """Fetch one extra row so continuation does not require a count query.""" + scope = await db.scalar( + select(Sync.id).where(Sync.id == sync_id, Sync.organization_id == organization_id) + ) + if scope is None: + raise SourceNotFound("Source does not exist in this organization") + statement = select(Entity).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.record_revision > 0, + ) + if filters.record_type is not None: + statement = statement.where(Entity.entity_definition_short_name == filters.record_type) + if filters.container_id is not None: + statement = statement.where(Entity.container_id == filters.container_id) + if filters.state == "active": + statement = statement.where(Entity.deleted_at.is_(None)) + elif filters.state == "deleted": + statement = statement.where(Entity.deleted_at.is_not(None)) + if after_id is not None: + statement = statement.where(Entity.id > after_id) + rows = await db.scalars(statement.order_by(Entity.id).limit(limit + 1)) + return tuple(source_record(row) for row in rows) diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py new file mode 100644 index 0000000000..f977c9b2e9 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -0,0 +1,99 @@ +"""Validated source capture commands. Provider payloads remain opaque JSON.""" + +import json +from typing import Literal +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator + + +class RecordIdentity(BaseModel): + """Native identity inside one sync, including provider-required container.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + record_type: str = Field(min_length=1) + native_id: str = Field(min_length=1) + container_id: str | None = None + + @property + def entity_key(self) -> str: + """Unambiguous encoding; native IDs remain separately available for actions.""" + return json.dumps( + [self.container_id, self.native_id], ensure_ascii=False, separators=(",", ":") + ) + + +class BlobReference(BaseModel): + """An already-written immutable blob required by this record.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + key: str = Field(min_length=1) + sha256: str = Field(pattern=r"^[a-f0-9]{64}$") + size_bytes: int = Field(ge=0) + media_type: str | None = None + + +class CaptureRecord(BaseModel): + """One current provider observation; sparse deletions are valid records.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + identity: RecordIdentity + payload: dict[str, JsonValue] + payload_schema_version: int = Field(default=1, ge=1) + kind: Literal["upsert", "delete"] = "upsert" + removal_reason: ( + Literal["provider_deleted", "scope_removed", "access_revoked", "absent"] | None + ) = None + completeness: Literal["complete", "metadata_only", "partial"] = "complete" + content_hash: str | None = None + source_created_at: AwareDatetime | None = None + source_updated_at: AwareDatetime | None = None + observed_at: AwareDatetime + blobs: tuple[BlobReference, ...] = () + + @model_validator(mode="after") + def validate_removal(self) -> "CaptureRecord": + """Keep provider deletion distinct from inaccessible or deselected scope.""" + if self.kind == "delete" and self.removal_reason is None: + raise ValueError("delete observations require removal_reason") + if self.kind == "upsert" and self.removal_reason is not None: + raise ValueError("upsert observations cannot have removal_reason") + return self + + +class WriterFence(BaseModel): + """Only the current source run may persist records or its checkpoint.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + organization_id: UUID + sync_id: UUID + job_id: UUID + epoch: int = Field(ge=1) + attempt_id: UUID + attempt_number: int = Field(ge=1) + + +class CaptureBatch(BaseModel): + """A bounded atomic capture batch; order within the batch is preserved.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + records: tuple[CaptureRecord, ...] = Field(max_length=500) + + +class CompletedScope(BaseModel): + """Evidence that one exact type/container was exhaustively enumerated.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + record_type: str = Field(min_length=1) + container_id: str | None = None + + +class ReconcileScope(BaseModel): + """Only call after successful complete enumeration, never partial failure.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + scope: CompletedScope + observed_at: AwareDatetime + limit: int = Field(default=250, ge=1, le=500) diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py new file mode 100644 index 0000000000..5a4c44f20c --- /dev/null +++ b/backend/airweave/domains/entities/canonical/service.py @@ -0,0 +1,57 @@ +"""Transaction boundary for source capture; external I/O happens before entry.""" + +from uuid import UUID + +from pydantic import JsonValue +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.db.unit_of_work import UnitOfWork +from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult +from airweave.domains.entities.canonical.requests import CaptureBatch, ReconcileScope, WriterFence +from airweave.domains.entities.canonical.store import CanonicalRecordStore + + +class CanonicalCaptureService: + """Own capture transactions; callers supply a fresh session without pending writes.""" + + def __init__(self, store: CanonicalRecordStore): + """Inject the persistence owner without hidden global sessions.""" + self.store = store + + async def activate_writer( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + job_id: UUID, + *, + attempt_id: UUID, + attempt_number: int, + ) -> WriterFence: + """Activate one valid source run under the same lock used for capture.""" + async with UnitOfWork(db): + return await self.store.activate_writer( + db, + organization_id, + sync_id, + job_id, + attempt_id=attempt_id, + attempt_number=attempt_number, + ) + + async def capture(self, db: AsyncSession, batch: CaptureBatch) -> CaptureResult: + """Commit record state and observed changes atomically.""" + async with UnitOfWork(db): + return await self.store.capture(db, batch) + + async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> ReconcileResult: + """Commit one bounded exact-scope absence reconciliation batch.""" + async with UnitOfWork(db): + return await self.store.reconcile_scope(db, request) + + async def save_checkpoint( + self, db: AsyncSession, fence: WriterFence, cursor_data: dict[str, JsonValue] + ) -> None: + """Advance only after the caller's full capture/reconciliation completion barrier.""" + async with UnitOfWork(db): + await self.store.save_checkpoint(db, fence, cursor_data) diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py new file mode 100644 index 0000000000..91fd0e0a9d --- /dev/null +++ b/backend/airweave/domains/entities/canonical/store.py @@ -0,0 +1,377 @@ +"""SQL persistence for canonical records. Never commits a caller's transaction.""" + +import hashlib +import json +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.models import ( + CaptureResult, + ChangePage, + ObservedChange, + ReconcileResult, + SourceRecord, +) +from airweave.domains.entities.canonical.requests import ( + CaptureBatch, + CaptureRecord, + ReconcileScope, + RecordIdentity, + WriterFence, +) +from airweave.models.entity import Entity +from airweave.models.entity_change import EntityChange +from airweave.models.sync import Sync +from airweave.models.sync_cursor import SyncCursor +from airweave.models.sync_job import SyncJob + + +class CanonicalStoreError(Exception): + """Stable boundary error, translated by the API edge.""" + + code = "canonical_store_error" + + +class SourceNotFound(CanonicalStoreError): + """No source in the authenticated organization.""" + + code = "source_not_found" + + +class StaleWriter(CanonicalStoreError): + """A cancelled or superseded source activity cannot write.""" + + code = "stale_writer" + + +class WriterBusy(CanonicalStoreError): + """Another source run still owns the capture writer.""" + + code = "writer_busy" + + +def capture_fingerprint(record: CaptureRecord) -> str: + """Observation times never create changes; sparse tombstone payloads do.""" + material = record.model_dump(mode="json", exclude={"observed_at"}) + serialized = json.dumps( + material, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False + ) + return hashlib.sha256(serialized.encode()).hexdigest() + + +def source_record(entity: Entity) -> SourceRecord: + """Canonical rows only; legacy metadata must be recaptured before reading.""" + return SourceRecord( + id=entity.id, + sync_id=entity.sync_id, + identity=RecordIdentity( + record_type=entity.entity_definition_short_name, + native_id=entity.native_id, + container_id=entity.container_id, + ), + revision=entity.record_revision, + payload=entity.source_payload, + payload_schema_version=entity.payload_schema_version, + capture_hash=entity.capture_hash, + content_hash=entity.content_hash, + completeness=entity.completeness, + observed_at=entity.observed_at, + source_created_at=entity.source_created_at, + source_updated_at=entity.source_updated_at, + deleted_at=entity.deleted_at, + removal_reason=entity.removal_reason, + blobs=entity.blob_references or (), + indexed_revision=entity.indexed_revision, + indexed_pipeline_version=entity.indexed_pipeline_version, + ) + + +class CanonicalRecordStore: + """One short Sync row lock serializes commits and allocates journal positions.""" + + async def _sync( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID, *, lock: bool = False + ) -> Sync: + statement = select(Sync).where(Sync.id == sync_id, Sync.organization_id == organization_id) + if lock: + statement = statement.with_for_update().execution_options(populate_existing=True) + sync = (await db.execute(statement)).scalar_one_or_none() + if sync is None: + raise SourceNotFound("Source does not exist in this organization") + return sync + + async def activate_writer( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + job_id: UUID, + *, + attempt_id: UUID, + attempt_number: int, + ) -> WriterFence: + """Claim a pending/running job; an identical activation is idempotent.""" + if attempt_number < 1: + raise ValueError("attempt_number must be positive") + sync = await self._sync(db, organization_id, sync_id, lock=True) + job = ( + await db.execute( + select(SyncJob) + .where( + SyncJob.id == job_id, + SyncJob.sync_id == sync_id, + SyncJob.organization_id == organization_id, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + if job is None or job.status not in ("pending", "running"): + raise StaleWriter("Writer must be an active job of this source") + identical = ( + sync.writer_job_id == job_id + and sync.writer_attempt_id == attempt_id + and sync.writer_attempt_number == attempt_number + ) + if sync.writer_job_id == job_id and not identical: + if attempt_number <= sync.writer_attempt_number: + raise StaleWriter("Attempt has been superseded or conflicts with current attempt") + if sync.writer_job_id != job_id: + if sync.writer_job_id is not None: + status = await db.scalar( + select(SyncJob.status).where(SyncJob.id == sync.writer_job_id) + ) + if status in ("pending", "running", "cancelling"): + raise WriterBusy("Another source run still owns capture") + if not identical: + sync.writer_epoch += 1 + sync.writer_job_id = job_id + sync.writer_attempt_id = attempt_id + sync.writer_attempt_number = attempt_number + await db.flush() + return WriterFence( + organization_id=organization_id, + sync_id=sync_id, + job_id=job_id, + epoch=sync.writer_epoch, + attempt_id=attempt_id, + attempt_number=attempt_number, + ) + + async def _fenced_sync(self, db: AsyncSession, fence: WriterFence) -> Sync: + sync = await self._sync(db, fence.organization_id, fence.sync_id, lock=True) + if ( + sync.writer_epoch != fence.epoch + or sync.writer_job_id != fence.job_id + or sync.writer_attempt_id != fence.attempt_id + or sync.writer_attempt_number != fence.attempt_number + ): + raise StaleWriter("Source run has been superseded") + status = await db.scalar( + select(SyncJob.status).where(SyncJob.id == fence.job_id).with_for_update() + ) + if status not in ("pending", "running"): + raise StaleWriter("Source run is no longer active") + return sync + + async def capture(self, db: AsyncSession, batch: CaptureBatch) -> CaptureResult: + """Apply observations and append snapshots under the same writer lock.""" + sync = await self._fenced_sync(db, batch.fence) + return await self._capture_locked(db, sync, batch) + + async def _capture_locked( + self, db: AsyncSession, sync: Sync, batch: CaptureBatch + ) -> CaptureResult: + changes = [] + unchanged = 0 + for observation in batch.records: + identity = observation.identity + entity = ( + await db.execute( + select(Entity) + .where( + Entity.organization_id == batch.fence.organization_id, + Entity.sync_id == sync.id, + Entity.entity_id == identity.entity_key, + Entity.entity_definition_short_name == identity.record_type, + ) + .execution_options(populate_existing=True) + ) + ).scalar_one_or_none() + fingerprint = capture_fingerprint(observation) + if entity is not None and entity.capture_hash == fingerprint: + entity.last_seen_run_id = batch.fence.attempt_id + entity.observed_at = observation.observed_at + unchanged += 1 + continue + if entity is None: + entity = Entity( + id=uuid4(), + organization_id=batch.fence.organization_id, + sync_id=sync.id, + entity_id=identity.entity_key, + entity_definition_short_name=identity.record_type, + record_revision=0, + ) + db.add(entity) + entity.sync_job_id = batch.fence.job_id + entity.native_id = identity.native_id + entity.container_id = identity.container_id + entity.source_payload = observation.payload + entity.payload_schema_version = observation.payload_schema_version + entity.record_revision += 1 + entity.capture_hash = fingerprint + entity.content_hash = observation.content_hash + entity.hash = observation.content_hash or fingerprint + entity.source_created_at = observation.source_created_at + entity.source_updated_at = observation.source_updated_at + entity.observed_at = observation.observed_at + entity.deleted_at = observation.observed_at if observation.kind == "delete" else None + entity.removal_reason = observation.removal_reason + entity.completeness = observation.completeness + entity.blob_references = [blob.model_dump(mode="json") for blob in observation.blobs] + entity.last_seen_run_id = batch.fence.attempt_id + entity.projection_error = None + # Flush before reading default/index state and before another observation of this ID. + await db.flush() + sync.observed_change_sequence += 1 + record = source_record(entity) + change = ObservedChange( + sequence=sync.observed_change_sequence, kind=observation.kind, record=record + ) + db.add( + EntityChange( + organization_id=batch.fence.organization_id, + sync_id=sync.id, + entity_record_id=entity.id, + sequence=change.sequence, + record_revision=record.revision, + kind=observation.kind, + snapshot=record.model_dump(mode="json"), + ) + ) + changes.append(change) + await db.flush() + return CaptureResult( + changes=tuple(changes), sequence=sync.observed_change_sequence, unchanged=unchanged + ) + + async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> ReconcileResult: + """Tombstone a bounded batch absent from one successfully enumerated scope.""" + sync = await self._fenced_sync(db, request.fence) + statement = ( + select(Entity) + .where( + Entity.organization_id == request.fence.organization_id, + Entity.sync_id == sync.id, + Entity.entity_definition_short_name == request.scope.record_type, + Entity.container_id.is_not_distinct_from(request.scope.container_id), + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + Entity.last_seen_run_id.is_distinct_from(request.fence.attempt_id), + ) + .order_by(Entity.id) + .limit(request.limit + 1) + ) + entities = list((await db.scalars(statement)).all()) + observations = tuple( + CaptureRecord( + identity=source_record(entity).identity, + payload=entity.source_payload, + payload_schema_version=entity.payload_schema_version, + kind="delete", + removal_reason="absent", + completeness=entity.completeness, + content_hash=entity.content_hash, + source_created_at=entity.source_created_at, + source_updated_at=entity.source_updated_at, + observed_at=request.observed_at, + blobs=entity.blob_references or (), + ) + for entity in entities[: request.limit] + ) + result = await self._capture_locked( + db, sync, CaptureBatch(fence=request.fence, records=observations) + ) + return ReconcileResult(capture=result, has_more=len(entities) > request.limit) + + async def save_checkpoint( + self, db: AsyncSession, fence: WriterFence, cursor_data: dict + ) -> None: + """Call only after capture barrier and successful exact-scope reconciliation.""" + await self._fenced_sync(db, fence) + cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + if cursor is None: + cursor = SyncCursor( + organization_id=fence.organization_id, + sync_id=fence.sync_id, + ) + db.add(cursor) + cursor.cursor_data = cursor_data + cursor.last_updated = datetime.now(timezone.utc) + await db.flush() + + async def read( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID, record_id: UUID + ) -> SourceRecord | None: + """Return current state including tombstones; never return another tenant's row.""" + entity = await db.scalar( + select(Entity).where( + Entity.id == record_id, + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.record_revision > 0, + ) + ) + return source_record(entity) if entity is not None else None + + async def changes( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + *, + after: int = 0, + limit: int = 100, + high_watermark: int | None = None, + ) -> ChangePage: + """Read immutable snapshots, bounded by a committed per-sync watermark.""" + if after < 0 or not 1 <= limit <= 500: + raise ValueError("Invalid change page bounds") + sync = await self._sync(db, organization_id, sync_id) + upper = sync.observed_change_sequence if high_watermark is None else high_watermark + if upper < after or upper > sync.observed_change_sequence: + raise ValueError("Invalid change high watermark") + rows = list( + ( + await db.scalars( + select(EntityChange) + .where( + EntityChange.organization_id == organization_id, + EntityChange.sync_id == sync_id, + EntityChange.sequence > after, + EntityChange.sequence <= upper, + ) + .order_by(EntityChange.sequence) + .limit(limit + 1) + ) + ).all() + ) + page = rows[:limit] + more = len(rows) > limit + return ChangePage( + changes=tuple( + ObservedChange( + sequence=row.sequence, + kind=row.kind, + record=SourceRecord.model_validate(row.snapshot), + ) + for row in page + ), + next_sequence=page[-1].sequence if more else upper, + high_watermark=upper, + has_more=more, + ) diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py new file mode 100644 index 0000000000..5333bd9594 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -0,0 +1,132 @@ +"""Isolated real PostgreSQL fixtures shared by canonical store/query tests.""" + +import importlib.util +import os +from datetime import datetime, timezone +from pathlib import Path +from uuid import uuid4 + +import pytest +from sqlalchemy import text +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from airweave.domains.entities.canonical.service import CanonicalCaptureService +from airweave.domains.entities.canonical.store import CanonicalRecordStore +from airweave.models import Organization, Sync, SyncJob +from alembic.migration import MigrationContext +from alembic.operations import Operations + + +def migrate(connection, filename): + """Run the actual migration with this connection's isolated search_path.""" + path = Path(__file__).resolve().parents[5] / "alembic" / "versions" / filename + spec = importlib.util.spec_from_file_location("canonical_test_migration", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + with Operations.context(MigrationContext.configure(connection)): + module.upgrade() + + +def seed_legacy(connection): + """Insert an existing metadata row using the pre-migration schema.""" + from sqlalchemy import MetaData, Table + + metadata = MetaData() + tables = { + name: Table(name, metadata, autoload_with=connection) + for name in ("organization", "sync", "sync_job", "entity") + } + org_id, sync_id, job_id, entity_id = [uuid4() for _ in range(4)] + now = datetime.now(timezone.utc).replace(tzinfo=None) + times = {"created_at": now, "modified_at": now} + connection.execute( + tables["organization"].insert().values(id=org_id, name="Legacy synthetic", **times) + ) + connection.execute( + tables["sync"] + .insert() + .values( + id=sync_id, + organization_id=org_id, + name="Legacy sync", + status="ACTIVE", + sync_type="full", + **times, + ) + ) + connection.execute( + tables["sync_job"] + .insert() + .values( + id=job_id, + organization_id=org_id, + sync_id=sync_id, + status="completed", + entities_inserted=1, + entities_updated=0, + entities_deleted=0, + entities_kept=0, + entities_skipped=0, + scheduled=False, + **times, + ) + ) + connection.execute( + tables["entity"] + .insert() + .values( + id=entity_id, + organization_id=org_id, + sync_id=sync_id, + sync_job_id=job_id, + entity_id="legacy-native-key", + entity_definition_short_name="event", + hash="old-hash", + **times, + ) + ) + + +@pytest.fixture +async def database(request): + """Every test gets its own schema; no global application DSN is consulted.""" + url = os.environ.get("CANONICAL_TEST_DATABASE_URL") + if not url: + pytest.skip("Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database") + schema = "canonical_test_" + uuid4().hex + admin = create_async_engine(url) + async with admin.begin() as connection: + await connection.execute(text(f'CREATE SCHEMA "{schema}"')) + engine = create_async_engine(url, connect_args={"server_settings": {"search_path": schema}}) + try: + async with engine.begin() as connection: + await connection.run_sync(migrate, "0000_baseline.py") + if getattr(request, "param", None) == "legacy": + await connection.run_sync(seed_legacy) + await connection.run_sync(migrate, "0001_canonical_records.py") + yield async_sessionmaker(engine, expire_on_commit=False) + finally: + await engine.dispose() + async with admin.begin() as connection: + await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + await admin.dispose() + + +@pytest.fixture +async def source(database): + organization_id, sync_id, job_id = uuid4(), uuid4(), uuid4() + async with database() as db: + db.add(Organization(id=organization_id, name="Synthetic canonical tests")) + await db.flush() + db.add(Sync(id=sync_id, organization_id=organization_id, name="Synthetic Gmail")) + await db.flush() + db.add( + SyncJob(id=job_id, organization_id=organization_id, sync_id=sync_id, status="running") + ) + await db.commit() + service = CanonicalCaptureService(CanonicalRecordStore()) + async with database() as db: + fence = await service.activate_writer( + db, organization_id, sync_id, job_id, attempt_id=uuid4(), attempt_number=1 + ) + return service, fence diff --git a/backend/airweave/domains/entities/canonical/tests/helpers.py b/backend/airweave/domains/entities/canonical/tests/helpers.py new file mode 100644 index 0000000000..30a8a73d9f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/helpers.py @@ -0,0 +1,23 @@ +"""Synthetic capture inputs and transaction invocation for database tests.""" + +from datetime import datetime, timezone + +from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord, RecordIdentity + + +def observation(native_id="one", container_id=None, **changes): + values = { + "identity": RecordIdentity( + record_type="event", native_id=native_id, container_id=container_id + ), + "payload": {"id": native_id, "summary": "Synthetic event"}, + "observed_at": datetime.now(timezone.utc), + "content_hash": "same-body", + } + values.update(changes) + return CaptureRecord(**values) + + +async def capture(database, service, fence, *records): + async with database() as db: + return await service.capture(db, CaptureBatch(fence=fence, records=records)) 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..1eda304049 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_http.py @@ -0,0 +1,66 @@ +"""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 CaptureBatch +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.tests.helpers import observation + + +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"))) + ) + app = FastAPI() + app.include_router(router, prefix="/sync") + app.add_exception_handler(CanonicalStoreError, record_error_response) + owner = fence.organization_id + + 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] = lambda: CanonicalQueryService( + CanonicalRecordStore(), CanonicalQueryStore(), "synthetic-cursor-secret" + ) + 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 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_query.py b/backend/airweave/domains/entities/canonical/tests/test_query.py new file mode 100644 index 0000000000..4f02991db3 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_query.py @@ -0,0 +1,122 @@ +"""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] diff --git a/backend/airweave/main.py b/backend/airweave/main.py index 2702ea968b..271adebfa3 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) @@ -141,6 +126,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..b995014cad 100644 --- a/backend/airweave/models/__init__.py +++ b/backend/airweave/models/__init__.py @@ -8,6 +8,7 @@ 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 @@ -37,6 +38,7 @@ "BillingPeriod", "Collection", "Entity", + "EntityChange", "EntityCount", "Connection", "ConnectionInitSession", diff --git a/backend/airweave/models/entity.py b/backend/airweave/models/entity.py index a5aa81522d..a46fa9ad94 100644 --- a/backend/airweave/models/entity.py +++ b/backend/airweave/models/entity.py @@ -1,9 +1,11 @@ """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, 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 +27,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 +41,28 @@ 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) + 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_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 +75,17 @@ class Entity(OrganizationBase): ) __table_args__ = ( + 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/sync.py b/backend/airweave/models/sync.py index baa0ce4c49..918c98b67c 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,13 @@ class Sync(OrganizationBase, UserMixin): sync_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) sync_config: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True) + 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..8c0cdc6f64 100644 --- a/backend/airweave/models/sync_job.py +++ b/backend/airweave/models/sync_job.py @@ -50,7 +50,6 @@ class SyncJob(OrganizationBase, UserMixin): "Entity", back_populates="sync_job", lazy="noload", - cascade="all, delete-orphan", passive_deletes=True, ) 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/alembic/versions/0001_canonical_records.py b/backend/alembic/versions/0001_canonical_records.py new file mode 100644 index 0000000000..63f9b17d11 --- /dev/null +++ b/backend/alembic/versions/0001_canonical_records.py @@ -0,0 +1,104 @@ +"""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("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/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..1226cf135b 100644 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -46,11 +46,10 @@ except Exception as 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 +echo "Starting against the provisioned database..." +cd /app +# Migrations and bootstrap are explicit deployment jobs, never replica startup. # Start application echo "Starting application..." -poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001 --reload +poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001 From 9ff95ebadb94e25ac135f774f3ecb307f32210d1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 02:20:17 -0700 Subject: [PATCH 003/266] Capture original records and immutable blobs through managed source connections --- .../domains/auth_provider/auth_result.py | 13 +- .../auth_provider/providers/composio.py | 91 ++++- .../domains/entities/canonical/__init__.py | 1 + .../domains/entities/canonical/models.py | 2 + .../domains/entities/canonical/query_store.py | 15 +- .../domains/entities/canonical/requests.py | 15 + .../domains/entities/canonical/service.py | 36 +- .../domains/entities/canonical/source.py | 44 +++ .../domains/entities/canonical/store.py | 218 +++++++++- .../canonical/tests/test_capture_pipeline.py | 183 +++++++++ .../canonical/tests/test_parent_visibility.py | 84 ++++ .../domains/sources/exceptions/__init__.py | 1 + .../domains/sources/exceptions/_exceptions.py | 4 + .../domains/sources/exceptions/classifier.py | 9 + backend/airweave/domains/sources/lifecycle.py | 25 +- .../sources/token_providers/protocol.py | 40 +- backend/airweave/domains/sources/types.py | 5 + .../airweave/domains/storage/file_service.py | 66 +++- .../storage/tests/test_file_service.py | 37 +- .../sync_pipeline/canonical_capture.py | 248 ++++++++++++ .../domains/sync_pipeline/capture_attempt.py | 28 ++ .../domains/sync_pipeline/contexts/runtime.py | 2 + .../airweave/domains/sync_pipeline/factory.py | 72 +++- .../domains/sync_pipeline/orchestrator.py | 28 +- .../domains/sync_pipeline/protocols.py | 2 + .../airweave/domains/sync_pipeline/stream.py | 3 +- .../sync_pipeline/tests/test_factory.py | 34 +- .../sync_pipeline/tests/test_orchestrator.py | 3 +- .../tests/test_orchestrator_coverage.py | 1 + backend/airweave/models/entity.py | 3 + backend/airweave/platform/configs/config.py | 6 + .../platform/cursors/google_calendar.py | 11 + .../airweave/platform/cursors/google_drive.py | 2 + .../platform/entities/google_drive.py | 4 +- .../platform/http_client/airweave_client.py | 15 +- .../http_client/composio_transport.py | 262 +++++++++++++ backend/airweave/platform/sources/__init__.py | 2 + backend/airweave/platform/sources/gmail.py | 59 ++- .../platform/sources/gmail_capture.py | 213 ++++++++++ .../platform/sources/google_calendar.py | 35 +- .../airweave/platform/sources/google_drive.py | 41 +- .../airweave/platform/sources/http_helpers.py | 3 + .../sources/records/google_calendar.py | 164 ++++++++ .../platform/sources/records/google_drive.py | 157 ++++++++ .../sources/records/google_drive_content.py | 59 +++ backend/airweave/platform/sources/slack.py | 180 ++++++++- .../sources/tests/test_gmail_capture.py | 371 ++++++++++++++++++ backend/airweave/platform/sources/wispr.py | 245 ++++++++++++ .../versions/0001_canonical_records.py | 3 + .../http_client/test_composio_transport.py | 308 +++++++++++++++ .../sources/records/test_google_calendar.py | 122 ++++++ .../sources/records/test_google_drive.py | 108 +++++ .../records/test_google_drive_content.py | 74 ++++ .../platform/sources/test_slack_capture.py | 120 ++++++ .../platform/sources/test_wispr_capture.py | 105 +++++ 55 files changed, 3877 insertions(+), 105 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/__init__.py create mode 100644 backend/airweave/domains/entities/canonical/source.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_parent_visibility.py create mode 100644 backend/airweave/domains/sync_pipeline/canonical_capture.py create mode 100644 backend/airweave/domains/sync_pipeline/capture_attempt.py create mode 100644 backend/airweave/platform/cursors/google_calendar.py create mode 100644 backend/airweave/platform/http_client/composio_transport.py create mode 100644 backend/airweave/platform/sources/gmail_capture.py create mode 100644 backend/airweave/platform/sources/records/google_calendar.py create mode 100644 backend/airweave/platform/sources/records/google_drive.py create mode 100644 backend/airweave/platform/sources/records/google_drive_content.py create mode 100644 backend/airweave/platform/sources/tests/test_gmail_capture.py create mode 100644 backend/airweave/platform/sources/wispr.py create mode 100644 backend/tests/unit/platform/http_client/test_composio_transport.py create mode 100644 backend/tests/unit/platform/sources/records/test_google_calendar.py create mode 100644 backend/tests/unit/platform/sources/records/test_google_drive.py create mode 100644 backend/tests/unit/platform/sources/records/test_google_drive_content.py create mode 100644 backend/tests/unit/platform/sources/test_slack_capture.py create mode 100644 backend/tests/unit/platform/sources/test_wispr_capture.py diff --git a/backend/airweave/domains/auth_provider/auth_result.py b/backend/airweave/domains/auth_provider/auth_result.py index d2b50b0ace..6424bbc85f 100644 --- a/backend/airweave/domains/auth_provider/auth_result.py +++ b/backend/airweave/domains/auth_provider/auth_result.py @@ -1,20 +1,29 @@ """Auth result types for auth provider credential fetch.""" +from __future__ import annotations + from dataclasses import dataclass -from typing import Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional + +if TYPE_CHECKING: + from airweave.domains.sources.token_providers.protocol import ( + ManagedAuthProvider, + ManagedToolAuthProvider, + ) @dataclass class AuthResult: """Result of auth provider credential fetch. - Credentials are always used directly by the source (no proxy mode). + Direct credentials and managed request access are distinct alternatives. source_config carries non-secret config fields (e.g., instance_url) that the auth provider extracted alongside credentials. """ credentials: Optional[Dict[str, Any]] = None source_config: Optional[Dict[str, Any]] = None + managed_auth: ManagedAuthProvider | ManagedToolAuthProvider | None = None @classmethod def direct(cls, credentials: Dict[str, Any]) -> "AuthResult": diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 9d95daf5fe..d3daafc57d 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -1,6 +1,7 @@ """Composio Test Auth Provider - provides authentication services for other integrations.""" -from typing import Any, Dict, List, Optional, Set +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set +from urllib.parse import quote from uuid import UUID import httpx @@ -20,6 +21,9 @@ from airweave.platform.configs.config import ComposioConfig from airweave.platform.decorators import auth_provider +if TYPE_CHECKING: + from airweave.domains.auth_provider.auth_result import AuthResult + @auth_provider( name="Composio", @@ -53,6 +57,7 @@ class ComposioAuthProvider(BaseAuthProvider): # Key: Airweave short name, Value: Composio toolkit slug # Only include mappings where names differ between Airweave and Composio SLUG_NAME_MAPPING = { + "wispr": "wispr_flow_mcp", "google_drive": "googledrive", "google_calendar": "googlecalendar", "google_docs": "googledocs", @@ -86,9 +91,81 @@ async def create( instance.api_key = credentials["api_key"] instance.auth_config_id = config.get("auth_config_id") instance.account_id = config.get("account_id") + instance.user_id = config.get("user_id") instance._last_credential_blob = None return instance + async def get_auth_result( + self, + source_short_name: str, + source_auth_config_fields: List[str], + optional_fields: Optional[Set[str]] = None, + source_config_field_mappings: Optional[Dict[str, str]] = None, + source_connection_id: Optional[UUID] = None, + ) -> "AuthResult": + """Bind managed access without requesting or extracting OAuth tokens.""" + from airweave.domains.auth_provider.auth_result import AuthResult + from airweave.domains.auth_provider.exceptions import AuthProviderConfigError + from airweave.domains.sources.token_providers.protocol import ( + ManagedAuthProvider, + ManagedToolAuthProvider, + ) + + hosts = { + "gmail": frozenset({"gmail.googleapis.com"}), + "google_calendar": frozenset({"www.googleapis.com"}), + "google_drive": frozenset({"www.googleapis.com"}), + "slack": frozenset({"slack.com"}), + } + if source_short_name not in hosts and source_short_name != "wispr": + raise AuthProviderConfigError( + "Managed Composio requests are not supported for this source yet", + provider_name="composio", + ) + if not self.account_id: + raise AuthProviderConfigError( + "An explicit Composio account is required", provider_name="composio" + ) + async with httpx.AsyncClient() as client: + account = await self._get_with_auth( + client, + "https://backend.composio.dev/api/v3/connected_accounts/" + + quote(self.account_id, safe=""), + ) + if account.get("toolkit", {}).get("slug", "").lower() != self._get_composio_slug( + source_short_name + ): + raise AuthProviderConfigError( + "Connected account toolkit does not match source", provider_name="composio" + ) + if self.auth_config_id and account.get("auth_config", {}).get("id") != self.auth_config_id: + raise AuthProviderConfigError( + "Connected account auth config does not match", provider_name="composio" + ) + if account.get("status", "").upper() != "ACTIVE": + from airweave.domains.auth_provider.exceptions import AuthProviderAuthError + + raise AuthProviderAuthError( + "Connected account requires attention", provider_name="composio" + ) + if source_short_name == "wispr": + if not self.user_id: + raise AuthProviderConfigError( + "Wispr requires an explicit Composio user binding", provider_name="composio" + ) + return AuthResult( + managed_auth=ManagedToolAuthProvider( + api_key=self.api_key, connected_account_id=self.account_id, user_id=self.user_id + ) + ) + return AuthResult( + managed_auth=ManagedAuthProvider( + api_key=self.api_key, + connected_account_id=self.account_id, + allowed_hosts=hosts[source_short_name], + ) + ) + def _get_composio_slug(self, airweave_short_name: str) -> str: """Get the Composio toolkit slug for an Airweave source short name. @@ -143,13 +220,20 @@ async def _get_with_auth( except httpx.HTTPStatusError as e: status = e.response.status_code self.logger.error(f"HTTP error from Composio API: {status} for {url}") - if status == 401: + if status in (401, 403): raise AuthProviderAuthError( "Composio API key is invalid or revoked", provider_name="composio", ) from e + if status == 404: + raise AuthProviderAccountNotFoundError( + "Composio connected account was not found", provider_name="composio" + ) from None if status == 429: - retry_after = float(e.response.headers.get("retry-after", 30)) + try: + retry_after = float(e.response.headers.get("retry-after", 30)) + except ValueError: + retry_after = 30.0 raise AuthProviderRateLimitError( "Composio API rate-limited", provider_name="composio", @@ -220,6 +304,7 @@ async def get_creds_for_source( source_short_name: The short name of the source to get credentials for source_auth_config_fields: The fields required for the source auth config optional_fields: Fields that can be skipped if not available in Composio + source_connection_id: Source connection requesting credentials. Returns: Credentials dictionary for the source diff --git a/backend/airweave/domains/entities/canonical/__init__.py b/backend/airweave/domains/entities/canonical/__init__.py new file mode 100644 index 0000000000..9d725c01ad --- /dev/null +++ b/backend/airweave/domains/entities/canonical/__init__.py @@ -0,0 +1 @@ +"""Canonical source records and their committed observed changes.""" diff --git a/backend/airweave/domains/entities/canonical/models.py b/backend/airweave/domains/entities/canonical/models.py index dc87978b4a..9560c22cfe 100644 --- a/backend/airweave/domains/entities/canonical/models.py +++ b/backend/airweave/domains/entities/canonical/models.py @@ -16,8 +16,10 @@ class SourceRecord(BaseModel): 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 diff --git a/backend/airweave/domains/entities/canonical/query_store.py b/backend/airweave/domains/entities/canonical/query_store.py index ec1073a9e6..c2436cfd6c 100644 --- a/backend/airweave/domains/entities/canonical/query_store.py +++ b/backend/airweave/domains/entities/canonical/query_store.py @@ -7,7 +7,12 @@ from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import RecordFilters -from airweave.domains.entities.canonical.store import SourceNotFound, source_record +from airweave.domains.entities.canonical.store import ( + SourceNotFound, + parent_is_visible, + source_record, + with_content_access, +) from airweave.models.entity import Entity from airweave.models.sync import Sync @@ -35,6 +40,7 @@ async def list_records( Entity.organization_id == organization_id, Entity.sync_id == sync_id, Entity.record_revision > 0, + parent_is_visible(), ) if filters.record_type is not None: statement = statement.where(Entity.entity_definition_short_name == filters.record_type) @@ -47,4 +53,9 @@ async def list_records( if after_id is not None: statement = statement.where(Entity.id > after_id) rows = await db.scalars(statement.order_by(Entity.id).limit(limit + 1)) - return tuple(source_record(row) for row in rows) + return tuple( + with_content_access( + source_record(row), row.removal_reason not in ("access_revoked", "scope_removed") + ) + for row in rows + ) diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index f977c9b2e9..3002e1b509 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -31,6 +31,9 @@ class BlobReference(BaseModel): sha256: str = Field(pattern=r"^[a-f0-9]{64}$") size_bytes: int = Field(ge=0) media_type: str | None = None + source_path: str | None = Field( + default=None, description="RFC 6901 JSON Pointer into the unchanged native payload" + ) class CaptureRecord(BaseModel): @@ -38,6 +41,7 @@ class CaptureRecord(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) identity: RecordIdentity + parent: RecordIdentity | None = None payload: dict[str, JsonValue] payload_schema_version: int = Field(default=1, ge=1) kind: Literal["upsert", "delete"] = "upsert" @@ -97,3 +101,14 @@ class ReconcileScope(BaseModel): scope: CompletedScope observed_at: AwareDatetime limit: int = Field(default=250, ge=1, le=500) + + +class StartedScope(CompletedScope): + """Start/restart exact-scope full enumeration, discarding earlier attempt sightings.""" + + +class RemovedScope(CompletedScope): + """Provider-confirmed scope loss; applies to children even if seen in this run.""" + + removal_reason: Literal["scope_removed", "access_revoked"] + observed_at: AwareDatetime diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index 5a4c44f20c..e321131355 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -7,7 +7,13 @@ from airweave.db.unit_of_work import UnitOfWork from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult -from airweave.domains.entities.canonical.requests import CaptureBatch, ReconcileScope, WriterFence +from airweave.domains.entities.canonical.requests import ( + CaptureBatch, + ReconcileScope, + RemovedScope, + StartedScope, + WriterFence, +) from airweave.domains.entities.canonical.store import CanonicalRecordStore @@ -49,6 +55,34 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re async with UnitOfWork(db): return await self.store.reconcile_scope(db, request) + async def start_scope(self, db: AsyncSession, fence: WriterFence, scope: StartedScope) -> None: + """Commit a new exact-scope enumeration boundary.""" + async with UnitOfWork(db): + await self.store.start_scope(db, fence, scope) + + async def remove_scope( + self, + db: AsyncSession, + fence: WriterFence, + scope: RemovedScope, + *, + limit: int = 250, + ) -> ReconcileResult: + """Commit a bounded batch of confirmed inaccessible/deselected records.""" + async with UnitOfWork(db): + return await self.store.remove_scope(db, fence, scope, limit=limit) + + async def reconcile_parents( + self, + db: AsyncSession, + fence: WriterFence, + *, + limit: int = 250, + ) -> ReconcileResult: + """Commit bounded durable removal of children with inaccessible/missing parents.""" + async with UnitOfWork(db): + return await self.store.reconcile_parents(db, fence, limit=limit) + async def save_checkpoint( self, db: AsyncSession, fence: WriterFence, cursor_data: dict[str, JsonValue] ) -> None: diff --git a/backend/airweave/domains/entities/canonical/source.py b/backend/airweave/domains/entities/canonical/source.py new file mode 100644 index 0000000000..e3c646840d --- /dev/null +++ b/backend/airweave/domains/entities/canonical/source.py @@ -0,0 +1,44 @@ +"""Explicit source capability for original provider observations from one crawl.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RemovedScope, + StartedScope, +) + +if TYPE_CHECKING: + from airweave.domains.browse_tree.types import NodeSelectionData + from airweave.domains.storage.file_service import FileService + from airweave.domains.syncs.cursors.cursor import SyncCursor + +SourceObservation = CaptureRecord | CompletedScope | RemovedScope | StartedScope + + +@runtime_checkable +class CanonicalSource(Protocol): + """Opt-in contract; declared types are audited operational records, not search chunks.""" + + canonical_record_types: tuple[str, ...] + + def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[SourceObservation, None]: + """Yield native records and only successfully completed exact enumeration scopes.""" + ... + + +@runtime_checkable +class ContainerScopedSource(Protocol): + """Audited child record types whose visibility depends on a root container record.""" + + canonical_container_parents: dict[str, str] diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 91fd0e0a9d..6d913b9f53 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -5,8 +5,9 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import select +from sqlalchemy import and_, exists, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import aliased from airweave.domains.entities.canonical.models import ( CaptureResult, @@ -20,6 +21,8 @@ CaptureRecord, ReconcileScope, RecordIdentity, + RemovedScope, + StartedScope, WriterFence, ) from airweave.models.entity import Entity @@ -72,6 +75,15 @@ def source_record(entity: Entity) -> SourceRecord: native_id=entity.native_id, container_id=entity.container_id, ), + parent=( + RecordIdentity( + record_type=entity.parent_record_type, + native_id=entity.parent_native_id, + container_id=entity.parent_container_id, + ) + if entity.parent_record_type is not None + else None + ), revision=entity.record_revision, payload=entity.source_payload, payload_schema_version=entity.payload_schema_version, @@ -89,6 +101,46 @@ def source_record(entity: Entity) -> SourceRecord: ) +def active_parent_exists(): + """Root container visibility; caller still supplies tenant/sync authorization.""" + parent = aliased(Entity, name="canonical_parent") + return exists( + select(parent.id).where( + parent.organization_id == Entity.organization_id, + parent.sync_id == Entity.sync_id, + parent.entity_definition_short_name == Entity.parent_record_type, + parent.native_id == Entity.parent_native_id, + parent.container_id.is_not_distinct_from(Entity.parent_container_id), + parent.record_revision > 0, + parent.deleted_at.is_(None), + parent.parent_record_type.is_(None), + ) + ).correlate(Entity) + + +def parent_is_visible(): + """Deny children of missing/deleted containers immediately, before cleanup catches up.""" + return or_(Entity.parent_record_type.is_(None), active_parent_exists()) + + +def content_is_available(): + """Shared SQL gate for read, historical payload access and search postvalidation.""" + return and_( + parent_is_visible(), + or_( + Entity.removal_reason.is_(None), + Entity.removal_reason.not_in(("scope_removed", "access_revoked")), + ), + ) + + +def with_content_access(record: SourceRecord, available: bool) -> SourceRecord: + """Retain citation identity/version while explicitly withholding revoked content.""" + if available: + return record + return record.model_copy(update={"content_access": "unavailable", "payload": {}, "blobs": ()}) + + class CanonicalRecordStore: """One short Sync row lock serializes commits and allocates journal positions.""" @@ -220,6 +272,13 @@ async def _capture_locked( entity.sync_job_id = batch.fence.job_id entity.native_id = identity.native_id entity.container_id = identity.container_id + entity.parent_record_type = ( + observation.parent.record_type if observation.parent else None + ) + entity.parent_native_id = observation.parent.native_id if observation.parent else None + entity.parent_container_id = ( + observation.parent.container_id if observation.parent else None + ) entity.source_payload = observation.payload entity.payload_schema_version = observation.payload_schema_version entity.record_revision += 1 @@ -280,6 +339,7 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re observations = tuple( CaptureRecord( identity=source_record(entity).identity, + parent=source_record(entity).parent, payload=entity.source_payload, payload_schema_version=entity.payload_schema_version, kind="delete", @@ -298,6 +358,137 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re ) return ReconcileResult(capture=result, has_more=len(entities) > request.limit) + async def start_scope(self, db: AsyncSession, fence: WriterFence, scope: StartedScope) -> None: + """Restart a full-scan scope without inheriting partial incremental sightings.""" + await self._fenced_sync(db, fence) + await db.execute( + update(Entity) + .where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == scope.record_type, + Entity.container_id.is_not_distinct_from(scope.container_id), + Entity.record_revision > 0, + ) + .values(last_seen_run_id=None) + ) + await db.flush() + + async def remove_scope( + self, + db: AsyncSession, + fence: WriterFence, + scope: RemovedScope, + *, + limit: int = 250, + ) -> ReconcileResult: + """Tombstone a bounded exact scope regardless of this attempt's sightings.""" + if not 1 <= limit <= 500: + raise ValueError("Scope removal batch limit must be between 1 and 500") + sync = await self._fenced_sync(db, fence) + entities = list( + ( + await db.scalars( + select(Entity) + .where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == scope.record_type, + Entity.container_id.is_not_distinct_from(scope.container_id), + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + ) + .order_by(Entity.id) + .limit(limit + 1) + ) + ).all() + ) + observations = tuple( + CaptureRecord( + identity=source_record(entity).identity, + parent=source_record(entity).parent, + payload=entity.source_payload, + payload_schema_version=entity.payload_schema_version, + kind="delete", + removal_reason=scope.removal_reason, + completeness=entity.completeness, + content_hash=entity.content_hash, + source_created_at=entity.source_created_at, + source_updated_at=entity.source_updated_at, + observed_at=scope.observed_at, + blobs=entity.blob_references or (), + ) + for entity in entities[:limit] + ) + result = await self._capture_locked( + db, sync, CaptureBatch(fence=fence, records=observations) + ) + return ReconcileResult(capture=result, has_more=len(entities) > limit) + + async def reconcile_parents( + self, + db: AsyncSession, + fence: WriterFence, + *, + limit: int = 250, + ) -> ReconcileResult: + """Durable unfinished removal is derived from child rows, not transient events.""" + if not 1 <= limit <= 500: + raise ValueError("Parent reconciliation limit must be between 1 and 500") + sync = await self._fenced_sync(db, fence) + rows = list( + ( + await db.scalars( + select(Entity) + .where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + Entity.parent_record_type.is_not(None), + ~active_parent_exists(), + ) + .order_by(Entity.id) + .limit(limit + 1) + ) + ).all() + ) + observations = [] + for entity in rows[:limit]: + record = source_record(entity) + reason = await db.scalar( + select(Entity.removal_reason).where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == record.parent.record_type, + Entity.native_id == record.parent.native_id, + Entity.container_id.is_not_distinct_from(record.parent.container_id), + Entity.record_revision > 0, + ) + ) + observations.append( + CaptureRecord( + identity=record.identity, + parent=record.parent, + payload=record.payload, + payload_schema_version=record.payload_schema_version, + kind="delete", + removal_reason="access_revoked" + if reason == "access_revoked" + else "scope_removed", + completeness=record.completeness, + content_hash=record.content_hash, + source_created_at=record.source_created_at, + source_updated_at=record.source_updated_at, + observed_at=datetime.now(timezone.utc), + blobs=record.blobs, + ) + ) + result = await self._capture_locked( + db, sync, CaptureBatch(fence=fence, records=tuple(observations)) + ) + return ReconcileResult(capture=result, has_more=len(rows) > limit) + async def save_checkpoint( self, db: AsyncSession, fence: WriterFence, cursor_data: dict ) -> None: @@ -326,7 +517,10 @@ async def read( Entity.record_revision > 0, ) ) - return source_record(entity) if entity is not None else None + if entity is None: + return None + available = await db.scalar(select(content_is_available()).where(Entity.id == entity.id)) + return with_content_access(source_record(entity), bool(available)) async def changes( self, @@ -362,12 +556,30 @@ async def changes( ) page = rows[:limit] more = len(rows) > limit + availability = ( + dict( + ( + await db.execute( + select(Entity.id, content_is_available()).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.id.in_([row.entity_record_id for row in page]), + ) + ) + ).all() + ) + if page + else {} + ) return ChangePage( changes=tuple( ObservedChange( sequence=row.sequence, kind=row.kind, - record=SourceRecord.model_validate(row.snapshot), + record=with_content_access( + SourceRecord.model_validate(row.snapshot), + bool(availability.get(row.entity_record_id, False)), + ), ) for row in page ), diff --git a/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py b/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py new file mode 100644 index 0000000000..a58672771d --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py @@ -0,0 +1,183 @@ +"""Single-crawl capture path against PostgreSQL, independent of index services.""" + +import logging +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select +from temporalio.testing import ActivityEnvironment + +from airweave.adapters.event_bus.fake import FakeEventBus +from airweave.domains.entities.canonical.requests import CompletedScope, RemovedScope, StartedScope +from airweave.domains.entities.canonical.tests.helpers import observation +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.exceptions import SyncFailureError +from airweave.domains.sync_pipeline.orchestrator import SyncOrchestrator +from airweave.domains.sync_pipeline.pipeline.entity_tracker import EntityTracker +from airweave.domains.sync_pipeline.stream import AsyncSourceStream +from airweave.domains.sync_pipeline.worker_pool import AsyncWorkerPool +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.models import Entity +from airweave.models import SyncCursor as StoredCursor + +pytestmark = pytest.mark.integration + + +def components(database, source): + service, fence = source + logger = logging.getLogger("canonical-capture-test") + config = SyncConfig() + config.behavior.skip_guardrails = True + ctx = SimpleNamespace( + organization_id=fence.organization_id, + sync_id=fence.sync_id, + sync_job_id=fence.job_id, + sync=SimpleNamespace(id=fence.sync_id), + sync_job=SimpleNamespace(id=fence.job_id), + organization=SimpleNamespace(id=fence.organization_id), + collection_id=uuid4(), + source_connection_id=uuid4(), + source_short_name="test", + execution_config=config, + logger=logger, + batch_size=1, + max_batch_latency_ms=0, + should_batch=True, + ) + bus = FakeEventBus() + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + ("event",), + CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), + ) + runtime = SyncRuntime( + source=SimpleNamespace(source_name="synthetic", supports_access_control=False), + entity_tracker=EntityTracker(fence.job_id, fence.sync_id, logger), + cursor=SyncCursor(fence.sync_id), + canonical_capture=pipeline, + ) + return ctx, pipeline, runtime, bus + + +def orchestrator(ctx, pipeline, runtime, generator, bus): + return SyncOrchestrator( + entity_pipeline=pipeline, + worker_pool=AsyncWorkerPool(logger=ctx.logger), + stream=AsyncSourceStream(generator, logger=ctx.logger), + sync_context=ctx, + runtime=runtime, + access_control_pipeline=AsyncMock(), + event_bus=bus, + usage_checker=AsyncMock(), + usage_ledger=AsyncMock(), + sync_cursor_service=AsyncMock(), + state_machine=AsyncMock(), + lifecycle_data=SimpleNamespace(), + sync_state_machine=AsyncMock(), + ) + + +async def test_original_payload_commits_without_index_and_preserves_stream_order(database, source): + ctx, pipeline, runtime, bus = components(database, source) + + async def generate(): + yield observation(payload={"native-only": {"nested": "retained"}, "version": 1}) + yield observation(payload={"native-only": {"nested": "retained"}, "version": 2}) + yield CompletedScope(record_type="event") + runtime.cursor.update(history_id="committed") + + runner = orchestrator(ctx, pipeline, runtime, generate(), bus) + await runner._start_sync() + await runner._process_entities() + await runner._cleanup_orphaned_entities_if_needed() + await runner._save_cursor_data() + async with database() as db: + row = (await db.scalars(select(Entity))).one() + assert row.source_payload == {"native-only": {"nested": "retained"}, "version": 2} + assert row.record_revision == 2 and row.indexed_revision is None + assert (await db.scalar(select(StoredCursor.cursor_data)))["history_id"] == "committed" + assert runtime.destinations == [] + assert len(bus.events) == 2 + + +async def test_failed_enumeration_keeps_partial_capture_but_no_cursor_or_orphan_delete( + database, source +): + ctx, pipeline, runtime, bus = components(database, source) + await pipeline.start(ctx) + await pipeline.process([observation("prior", "a")], ctx, runtime) + + async def generate(): + yield StartedScope(record_type="event", container_id="a") + yield observation("partial", "a") + runtime.cursor.update(history_id="must-not-save") + raise ConnectionError("provider page failed") + + runner = orchestrator(ctx, pipeline, runtime, generate(), bus) + await runner._start_sync() + with pytest.raises(ConnectionError, match="provider page failed"): + await runner._process_entities() + async with database() as db: + rows = list((await db.scalars(select(Entity))).all()) + assert "prior" in {row.native_id for row in rows} + assert all(row.deleted_at is None for row in rows) + assert await db.scalar(select(StoredCursor.id)) is None + + +async def test_same_attempt_rebaseline_and_confirmed_scope_loss(database, source): + ctx, pipeline, runtime, _ = components(database, source) + await pipeline.start(ctx) + await pipeline.process( + [ + observation("incremental-before-410", "a"), + observation("unrelated", "b"), + StartedScope(record_type="event", container_id="a"), + observation("full-after-reset", "a"), + CompletedScope(record_type="event", container_id="a"), + ], + ctx, + runtime, + ) + await pipeline.cleanup_orphaned_entities(ctx, runtime) + async with database() as db: + rows = {row.native_id: row for row in (await db.scalars(select(Entity))).all()} + assert rows["incremental-before-410"].removal_reason == "absent" + assert rows["full-after-reset"].deleted_at is None + assert rows["unrelated"].deleted_at is None + await pipeline.process( + [ + RemovedScope( + record_type="event", + container_id="a", + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) + ], + ctx, + runtime, + ) + async with database() as db: + rows = {row.native_id: row for row in (await db.scalars(select(Entity))).all()} + assert rows["full-after-reset"].removal_reason == "access_revoked" + assert rows["unrelated"].deleted_at is None + with pytest.raises(SyncFailureError, match="after confirmed scope loss"): + await pipeline.process([observation("late", "a")], ctx, runtime) + + +def test_attempt_identity_uses_temporal_metadata_or_explicit_local_identity(): + with pytest.raises(ValueError, match="explicit CaptureAttempt"): + resolve_capture_attempt() + explicit = CaptureAttempt(id=uuid4(), number=3) + assert resolve_capture_attempt(explicit) == explicit + environment = ActivityEnvironment() + first = environment.run(resolve_capture_attempt) + second = environment.run(resolve_capture_attempt) + assert first == second and first.number == environment.info.attempt 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/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..b7b45d37d6 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,6 +163,7 @@ 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 @@ -398,6 +403,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 +623,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 +640,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 +717,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/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/storage/file_service.py b/backend/airweave/domains/storage/file_service.py index 945ebf40d6..50a431a473 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,7 +21,12 @@ 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.paths import paths from airweave.domains.storage.protocols import StorageBackend @@ -43,13 +49,57 @@ 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, + ) -> 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) + except httpx.HTTPStatusError as exc: + if exc.response.status_code != 401 or not auth.supports_refresh: + raise + headers = await authorization_headers(auth, refresh=True) + await self._stream_download(client, url, headers, temp_path, logger) + 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 +111,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 @@ -190,8 +239,7 @@ async def download_from_url( except httpx.HTTPStatusError as first_error: if first_error.response.status_code == 401 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 +249,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/tests/test_file_service.py b/backend/airweave/domains/storage/tests/test_file_service.py index 65637d0309..a2fc54b294 100644 --- a/backend/airweave/domains/storage/tests/test_file_service.py +++ b/backend/airweave/domains/storage/tests/test_file_service.py @@ -19,9 +19,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 +126,36 @@ 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): + 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") 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..c471866741 --- /dev/null +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -0,0 +1,248 @@ +"""Capture pipeline: native records commit before independent search projection.""" + +from collections.abc import 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.models import CaptureResult +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.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] | 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 {}) + for child, parent in self._container_parents.items(): + if child not in self._record_types 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("Only root-container visibility relationships are supported") + 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, + ) + + 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, + 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 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 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/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/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index cdf225b02f..e25c87c509 100644 --- a/backend/airweave/domains/sync_pipeline/factory.py +++ b/backend/airweave/domains/sync_pipeline/factory.py @@ -22,6 +22,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 +32,9 @@ 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.service import CanonicalCaptureService +from airweave.domains.entities.canonical.source import CanonicalSource, ContainerScopedSource +from airweave.domains.entities.canonical.store import CanonicalRecordStore from airweave.domains.entities.protocols import ( EntityCountRepositoryProtocol, EntityRepositoryProtocol, @@ -43,6 +47,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 @@ -145,6 +151,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() @@ -185,13 +192,20 @@ async def create_orchestrator( access_token=access_token, ) 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 if isinstance(source_result.source, CanonicalSource) else None + ) + 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 +238,24 @@ 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), + 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) @@ -316,12 +345,20 @@ def _build_stream( sync_context: SyncContext, ) -> AsyncSourceStream: """Build the async source stream from the source generator.""" - return AsyncSourceStream( - source_generator=runtime.source.generate_entities( + 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, ) @@ -344,6 +381,11 @@ async def _build_source( ) -> SourceBuildResult: """Build source instance, cursor, file service, and node selections.""" if execution_config and execution_config.behavior.replay_from_arf: + source_class = self._source_registry.get(source_connection.short_name).source_class_ref + if isinstance(source_class, CanonicalSource): + 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 +399,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, diff --git a/backend/airweave/domains/sync_pipeline/orchestrator.py b/backend/airweave/domains/sync_pipeline/orchestrator.py index e5f12e5e72..f9371e5faf 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,7 +54,7 @@ class SyncOrchestrator: def __init__( self, - entity_pipeline: EntityPipeline, + entity_pipeline: EntityPipeline | CanonicalCapturePipeline, worker_pool: AsyncWorkerPool, stream: AsyncSourceStream, sync_context: SyncContext, @@ -196,14 +197,15 @@ 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) + await self.stream.start() async def _process_entities(self) -> None: # noqa: C901 """Process entities using micro-batching with bounded inner concurrency.""" @@ -306,6 +308,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 +460,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 +525,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 +649,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 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..f7b3a1cfa0 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 @@ -132,7 +134,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]) +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() @@ -177,6 +180,9 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(): ): mock_source = MagicMock() mock_source.generate_entities = MagicMock(return_value=AsyncMock()) + if 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 +200,19 @@ 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 + if canonical: + assert isinstance(orchestrator.entity_pipeline, CanonicalCapturePipeline) + mock_build_destinations.assert_not_called() + mock_disp_builder.assert_not_called() + 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 # --------------------------------------------------------------------------- @@ -397,10 +411,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 +543,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 +559,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 diff --git a/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py b/backend/airweave/domains/sync_pipeline/tests/test_orchestrator.py index f2446ac29e..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, @@ -479,6 +479,7 @@ async def test_checkpoint_failure_prevents_completed_transition(): 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, 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/models/entity.py b/backend/airweave/models/entity.py index a46fa9ad94..f6c6bf0359 100644 --- a/backend/airweave/models/entity.py +++ b/backend/airweave/models/entity.py @@ -44,6 +44,9 @@ class Entity(OrganizationBase): # 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) 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") diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 77d8d004c7..6b3438d427 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -1183,6 +1183,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..45e1e104ba --- /dev/null +++ b/backend/airweave/platform/cursors/google_calendar.py @@ -0,0 +1,11 @@ +"""Independent continuation tokens for each calendar's unexpanded event stream.""" + +from pydantic import Field + +from airweave.platform.cursors._base import BaseCursor + + +class GoogleCalendarCursor(BaseCursor): + """Token parameters are fixed by the canonical generator, not user date ranges.""" + + calendar_tokens: dict[str, str] = 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/entities/google_drive.py b/backend/airweave/platform/entities/google_drive.py index c4ff91a50b..801f75c736 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 = ( 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/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/__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/gmail.py b/backend/airweave/platform/sources/gmail.py index 57819cc3bc..a739f224c0 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -11,6 +11,8 @@ import asyncio import base64 +import json +from contextlib import aclosing from datetime import datetime from typing import Any, AsyncGenerator, Dict, List, Optional, Set @@ -20,8 +22,12 @@ 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.source import SourceObservation 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, + authorization_headers, +) from airweave.domains.storage import FileSkippedException from airweave.domains.storage.file_service import FileService from airweave.domains.syncs.cursors.cursor import SyncCursor @@ -37,6 +43,7 @@ ) from airweave.platform.http_client.airweave_client import AirweaveHttpClient from airweave.platform.sources._base import BaseSource +from airweave.platform.sources.gmail_capture import GmailCapture from airweave.platform.sources.http_helpers import raise_for_status from airweave.platform.sources.retry_helpers import ( wait_rate_limit_with_backoff, @@ -83,6 +90,24 @@ class GmailSource(BaseSource): It supports syncing email threads, individual messages, and file attachments. """ + canonical_record_types = ("message",) + + async def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[SourceObservation, None]: + """Capture native messages once; thread/search projections are downstream.""" + if node_selections: + raise ValueError("Gmail capture does not support selected node scopes") + capture = GmailCapture( + self._get, self._build_gmail_query(), files=files, attachment_get=self._get_mime_json + ) + async for observation in capture.generate(cursor): + yield observation + # ----------------------- # Construction / Config # ----------------------- @@ -90,7 +115,7 @@ class GmailSource(BaseSource): async def create( cls, *, - auth: TokenProviderProtocol, + auth: SourceAuthProvider, logger: ContextualLogger, http_client: AirweaveHttpClient, config: GmailConfig, @@ -230,13 +255,11 @@ 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), @@ -260,8 +283,7 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: if response.status_code == 429: self.logger.warning( - f"Got 429 Rate Limited from Gmail API. Headers: {response.headers}. " - f"Body: {response.text}." + "Gmail rate limit reached" ) raise_for_status( @@ -274,6 +296,27 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: self.logger.debug(f"Response data keys: {list(data.keys())}") return data + async def _get_mime_json(self, url: str, *, max_bytes: int) -> dict: + """Read bounded attachment JSON through the managed source HTTP transport.""" + headers = await self._authed_headers() + async with self.http_client.stream("GET", url, headers=headers) as response: + raise_for_status( + response, source_short_name=self.short_name, + token_provider_kind=self.auth.provider_kind, + ) + body = bytearray() + async with aclosing(response.aiter_bytes()) as chunks: + async for chunk in chunks: + if len(body) + len(chunk) > max_bytes: + raise FileSkippedException( + "Encoded MIME body exceeds size limit", "MIME body" + ) + body.extend(chunk) + value = json.loads(body) + if not isinstance(value, dict): + raise ValueError("Gmail MIME response must be an object") + return value + # ----------------------- # Cursor helper # ----------------------- diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py new file mode 100644 index 0000000000..ef5bb93e7f --- /dev/null +++ b/backend/airweave/platform/sources/gmail_capture.py @@ -0,0 +1,213 @@ +"""Original Gmail message capture, independent of derived thread/search entities.""" + +import base64 +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator +from datetime import datetime, timezone + +import httpx + +from airweave.domains.entities.canonical.requests import ( + BlobReference, + CaptureRecord, + CompletedScope, + RecordIdentity, +) +from airweave.domains.entities.canonical.source import SourceObservation +from airweave.domains.storage.exceptions import FileSkippedException +from airweave.domains.storage.file_service import FileService +from airweave.domains.syncs.cursors.cursor import SyncCursor + +BASE = "https://gmail.googleapis.com/gmail/v1/users/me" +Get = Callable[..., Awaitable[dict]] + + +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 as exc: + if 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(blob.model_copy(update={"source_path": path})) + 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(self, boundary: str) -> tuple[list[str], str]: + """Drain history before emitting, so an expired boundary can restart cleanly.""" + params = {"startHistoryId": boundary, "maxResults": 500} + affected: dict[str, None] = {} + tokens: set[str] = set() + while True: + page = await self.get(f"{BASE}/history", params=dict(params)) + for entry in page.get("history", []): + # `messages` is the full changed-ID list; typed lists can overlap it. + for message in entry.get("messages", []): + affected[message["id"]] = None + for field in ("messagesAdded", "messagesDeleted", "labelsAdded", "labelsRemoved"): + for change in entry.get(field, []): + affected[change["message"]["id"]] = None + token = page.get("nextPageToken") + if not token: + return list(affected), str(page["historyId"]) + if token in tokens: + raise ValueError("Gmail history pagination repeated a token") + tokens.add(token) + params["pageToken"] = token + + async def enumerate(self) -> AsyncGenerator[CaptureRecord, None]: + """List every page in the provider's exact configured query scope.""" + params = {"maxResults": 500, "includeSpamTrash": "true"} + if self.query: + params["q"] = self.query + tokens: set[str] = set() + while True: + page = await self.get(f"{BASE}/messages", params=dict(params)) + for item in page.get("messages", []): + yield await self.message(item["id"]) + token = page.get("nextPageToken") + if not token: + return + if token in tokens: + raise ValueError("Gmail messages pagination repeated a token") + tokens.add(token) + params["pageToken"] = token + + async def saved_history(self, cursor: SyncCursor | None) -> tuple[list[str], str] | None: + """Only canonical unfiltered checkpoints can resume this record collection.""" + saved = cursor.data if cursor else {} + boundary = saved.get("history_id") if saved.get("canonical_query") == "" else None + if not boundary: + return None + try: + return await self.history(boundary) + except httpx.HTTPStatusError as exc: + if exc.response.status_code != 404: + raise + return None + + async def generate(self, cursor: SyncCursor | None) -> AsyncGenerator[SourceObservation, None]: + """Filtered queries reconcile each run; unfiltered mailboxes use Gmail history. + + Gmail queries cannot be faithfully evaluated against history JSON. Full query + enumeration is a best-effort changing mailbox view, not a snapshot. Never + advance an incremental cursor for such a view. External MIME bytes require + successful immutable storage; missing storage or oversize bodies remain partial. + """ + if self.query: + async for record in self.enumerate(): + yield record + yield CompletedScope(record_type="message") + if cursor: + cursor.update(history_id="", canonical_query=self.query) + return + + history = await self.saved_history(cursor) + if history is None: + profile = await self.get(f"{BASE}/profile") + boundary = str(profile["historyId"]) + async for record in self.enumerate(): + yield record + # If this boundary expires mid-crawl, fail the run. Starting another full + # crawl in this attempt would incorrectly retain earlier seen records. + history = await self.history(boundary) + full = True + else: + full = False + message_ids, latest = history + for message_id in message_ids: + yield await self.message(message_id) + if full: + yield CompletedScope(record_type="message") + if cursor: + cursor.update(history_id=latest, canonical_query="") diff --git a/backend/airweave/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index 86d9f17d3a..0619d8eb46 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -42,10 +42,15 @@ 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.token_providers.protocol import TokenProviderProtocol +from airweave.domains.entities.canonical.source import SourceObservation +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 ( @@ -57,6 +62,7 @@ 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.records.google_calendar import generate_calendar_observations from airweave.platform.sources.retry_helpers import ( retry_if_rate_limit_or_timeout, wait_rate_limit_with_backoff, @@ -77,7 +83,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 +96,22 @@ class GoogleCalendarSource(BaseSource): Google Calendar scheduling information for productivity and time management insights. """ + canonical_record_types = ("calendar", "event") + canonical_container_parents = {"event": "calendar"} + + async def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[SourceObservation, None]: + """Capture original masters, exceptions and cancellations; no recurrence expansion.""" + if node_selections: + raise ValueError("Calendar selection scopes are not supported by canonical capture") + async for observation in generate_calendar_observations(self._get, cursor): + yield observation + # ----------------------- # Construction / Config # ----------------------- @@ -96,7 +119,7 @@ class GoogleCalendarSource(BaseSource): async def create( cls, *, - auth: TokenProviderProtocol, + auth: SourceAuthProvider, logger: ContextualLogger, http_client: AirweaveHttpClient, config: GoogleCalendarConfig, @@ -119,13 +142,11 @@ 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), diff --git a/backend/airweave/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py index 18f3eefb58..3b673c647a 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -25,8 +25,13 @@ 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.requests import CaptureRecord +from airweave.domains.entities.canonical.source import SourceObservation 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 @@ -43,6 +48,8 @@ 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.records.google_drive import generate_drive_observations +from airweave.platform.sources.records.google_drive_content import capture_file_content from airweave.platform.sources.retry_helpers import ( retry_if_rate_limit_or_timeout, wait_rate_limit_with_backoff, @@ -77,11 +84,35 @@ class GoogleDriveSource(BaseSource): while maintaining proper organization and access permissions. """ + canonical_record_types = ("file",) + + async def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[SourceObservation, None]: + """Capture original file metadata; missing body bytes remain explicitly metadata-only.""" + if self.include_patterns or node_selections: + raise ValueError("Drive path selection is not yet supported by canonical capture") + async for observation in generate_drive_observations(self._get, cursor): + if isinstance(observation, CaptureRecord) and files is not None: + observation = await capture_file_content( + observation, + files=files, + get=self._get, + client=self.http_client, + auth=self.auth, + logger=self.logger, + ) + yield observation + @classmethod async def create( cls, *, - auth: TokenProviderProtocol, + auth: SourceAuthProvider, logger: ContextualLogger, http_client: AirweaveHttpClient, config: GoogleDriveConfig, @@ -123,13 +154,11 @@ 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) raise_for_status( diff --git a/backend/airweave/platform/sources/http_helpers.py b/backend/airweave/platform/sources/http_helpers.py index 593940f351..7743a5fb3b 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}", 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..aca9241213 --- /dev/null +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -0,0 +1,164 @@ +"""Original Calendar resources with per-calendar incremental synchronization.""" + +from collections.abc import AsyncGenerator, Awaitable, Callable +from datetime import datetime, timezone +from urllib.parse import quote + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + RemovedScope, + StartedScope, +) +from airweave.domains.entities.canonical.source import SourceObservation +from airweave.domains.sources.exceptions import ( + SourceEntityNotFoundError, + SourceGoneError, +) +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.platform.cursors.google_calendar import GoogleCalendarCursor + +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") + deleted = payload.get("status") == "cancelled" or payload.get("deleted") is True + return CaptureRecord( + identity=RecordIdentity(record_type=kind, native_id=native_id, container_id=calendar_id), + payload=payload, + kind="delete" if deleted else "upsert", + removal_reason="provider_deleted" if deleted else None, + observed_at=datetime.now(timezone.utc), + ) + + +async def capture_calendar( + get: GetJSON, calendar_id: str, token: str | None, tokens: dict[str, str] +) -> AsyncGenerator[SourceObservation, None]: + """A 410 resets only this calendar, including sightings from earlier delta pages.""" + full = token is None + if full: + yield StartedScope(record_type="event", container_id=calendar_id) + params: dict[str, str | int] = { + "maxResults": 2500, + "showDeleted": "true", + "singleEvents": "false", + } + if token: + params["syncToken"] = token + seen_pages: set[str] = set() + while True: + try: + page = CalendarPage.model_validate( + await get(f"{BASE}/calendars/{quote(calendar_id, safe='')}/events", params=params) + ) + except SourceGoneError: + if full: + raise + full = True + params.pop("syncToken", None) + params.pop("pageToken", None) + seen_pages.clear() + yield StartedScope(record_type="event", container_id=calendar_id) + continue + for item in page.items: + yield record("event", item, calendar_id) + if not page.nextPageToken: + tokens[calendar_id] = page.completed_token() + if full: + yield CompletedScope(record_type="event", container_id=calendar_id) + return + if page.nextPageToken in seen_pages: + raise ValueError("Calendar returned a repeated page token") + seen_pages.add(page.nextPageToken) + params["pageToken"] = page.nextPageToken + + +async def generate_calendar_observations( + get: GetJSON, cursor: SyncCursor | None +) -> AsyncGenerator[SourceObservation, None]: + """Refresh calendar membership; capture each accessible calendar independently.""" + previous = GoogleCalendarCursor.model_validate(cursor.data if cursor else {}) + tokens: dict[str, str] = {} + seen_calendars: set[str] = set() + seen_pages: set[str] = set() + params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"} + yield StartedScope(record_type="calendar") + while True: + page = CalendarPage.model_validate( + await get(f"{BASE}/users/me/calendarList", params=params) + ) + for item in page.items: + calendar_record = record("calendar", item) + calendar_id = calendar_record.identity.native_id + if calendar_id in seen_calendars: + raise ValueError("Calendar list repeated an identity across pages") + seen_calendars.add(calendar_id) + yield calendar_record + try: + async for observation in capture_calendar( + get, calendar_id, previous.calendar_tokens.get(calendar_id), tokens + ): + yield observation + except SourceEntityNotFoundError: + yield CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id=calendar_id), + payload={"id": calendar_id}, + kind="delete", + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) + yield RemovedScope( + record_type="event", + container_id=calendar_id, + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) + if not page.nextPageToken: + break + if page.nextPageToken in seen_pages: + raise ValueError("Calendar list returned a repeated page token") + seen_pages.add(page.nextPageToken) + params["pageToken"] = page.nextPageToken + for missing in previous.calendar_tokens.keys() - seen_calendars: + yield CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id=missing), + payload={"id": missing}, + kind="delete", + removal_reason="scope_removed", + observed_at=datetime.now(timezone.utc), + ) + yield RemovedScope( + record_type="event", + container_id=missing, + removal_reason="scope_removed", + observed_at=datetime.now(timezone.utc), + ) + yield CompletedScope(record_type="calendar") + if cursor is not None: + cursor.update(calendar_tokens=tokens) 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..fac15cd806 --- /dev/null +++ b/backend/airweave/platform/sources/records/google_drive.py @@ -0,0 +1,157 @@ +"""Drive metadata capture with boundary-before-scan and complete changes replay.""" + +from collections.abc import AsyncGenerator, Awaitable, Callable +from datetime import datetime, timezone + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + StartedScope, +) +from airweave.domains.entities.canonical.source import SourceObservation +from airweave.domains.sources.exceptions import SourceGoneError +from airweave.domains.syncs.cursors.cursor import SyncCursor + +BASE = "https://www.googleapis.com/drive/v3" +GetJSON = Callable[..., Awaitable[dict]] + + +class DrivePage(BaseModel): + """Original item dictionaries survive validation intact.""" + + model_config = ConfigDict(extra="ignore") + files: list[dict[str, JsonValue]] = Field(default_factory=list) + changes: list[dict[str, JsonValue]] = Field(default_factory=list) + nextPageToken: str | None = None + newStartPageToken: str | None = None + incompleteSearch: bool = False + + +class DriveScopeChanged(Exception): + """A shared-drive membership change requires fresh corpus enumeration.""" + + +class DriveCheckpoint(BaseModel): + """In-memory result of draining one change stream, never persisted prematurely.""" + + token: str | None = None + + +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), + ) + + +async def enumerate_files(get: GetJSON) -> AsyncGenerator[CaptureRecord, None]: + """Fail incomplete all-drive listings rather than deleting unenumerated records.""" + params: dict[str, str | int] = { + "pageSize": 1000, + "fields": "*", + "spaces": "drive", + "corpora": "allDrives", + "includeItemsFromAllDrives": "true", + "supportsAllDrives": "true", + } + seen: set[str] = set() + while True: + page = DrivePage.model_validate(await get(f"{BASE}/files", params=params)) + if page.incompleteSearch: + raise ValueError("Drive reported incomplete enumeration; checkpoint was not advanced") + for item in page.files: + yield file_record(item) + if not page.nextPageToken: + return + if page.nextPageToken in seen: + raise ValueError("Drive returned a repeated listing page token") + seen.add(page.nextPageToken) + params["pageToken"] = page.nextPageToken + + +async def replay_changes( + get: GetJSON, token: str, checkpoint: DriveCheckpoint +) -> AsyncGenerator[CaptureRecord, None]: + """Read every change page; metadata edits and trash state are real updates.""" + params: dict[str, str | int] = { + "pageToken": token, + "pageSize": 1000, + "fields": "*", + "spaces": "drive", + "includeRemoved": "true", + "includeCorpusRemovals": "true", + "includeItemsFromAllDrives": "true", + "supportsAllDrives": "true", + } + seen: set[str] = {token} + while True: + page = DrivePage.model_validate(await get(f"{BASE}/changes", params=params)) + for change in page.changes: + if change.get("changeType") == "drive": + # Drive membership changes require a new exhaustive file scan. + raise DriveScopeChanged("Shared-drive membership changed") + if change.get("removed") is True: + yield file_record({"id": change.get("fileId"), "change": change}, removed=True) + else: + item = change.get("file") + if not isinstance(item, dict): + raise ValueError("Drive file change omitted its current file resource") + yield file_record(item) + if not page.nextPageToken: + if not page.newStartPageToken: + raise ValueError("Drive change enumeration ended without newStartPageToken") + checkpoint.token = page.newStartPageToken + return + if page.nextPageToken in seen: + raise ValueError("Drive returned a repeated changes page token") + seen.add(page.nextPageToken) + params["pageToken"] = page.nextPageToken + + +async def generate_drive_observations( + get: GetJSON, cursor: SyncCursor | None +) -> AsyncGenerator[SourceObservation, None]: + """Only canonical checkpoints can skip bootstrap; legacy index cursors cannot.""" + state = cursor.data if cursor else {} + token = state.get("canonical_page_token") + for reset_attempt in range(2): + full = not token + if full: + boundary = await get( + f"{BASE}/changes/startPageToken", params={"supportsAllDrives": "true"} + ) + token = boundary.get("startPageToken") + if not isinstance(token, str) or not token: + raise ValueError("Drive omitted its starting change boundary") + yield StartedScope(record_type="file") + async for item in enumerate_files(get): + yield item + checkpoint = DriveCheckpoint() + try: + async for item in replay_changes(get, token, checkpoint): + yield item + except (DriveScopeChanged, SourceGoneError): + if reset_attempt: + raise + token = None + continue + if full: + yield CompletedScope(record_type="file") + if cursor is not None: + cursor.update(canonical_page_token=checkpoint.token) + return 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..b102949e60 --- /dev/null +++ b/backend/airweave/platform/sources/records/google_drive_content.py @@ -0,0 +1,59 @@ +"""Immutable Drive bodies captured before a record may claim content completeness.""" + +from urllib.parse import quote, urlencode + +from airweave.core.logging import ContextualLogger +from airweave.domains.entities.canonical.requests import 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_drive import BASE, GetJSON + + +async def capture_file_content( + record: CaptureRecord, + *, + files: FileService, + get: GetJSON, + client: AirweaveHttpClient, + auth: SourceAuthProvider, + logger: ContextualLogger, +) -> 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='')}" + if mime.startswith("application/vnd.google-apps."): + export = GOOGLE_EXPORT_FORMATS.get(mime) + if export is None: + return record + media_type = export[0] + download_url = url + "/export?" + urlencode({"mimeType": media_type}) + else: + media_type = mime + download_url = url + "?alt=media&supportsAllDrives=true" + try: + blob = await files.capture_canonical_url( + url=download_url, client=client, auth=auth, logger=logger, media_type=media_type + ) + except FileSkippedException: + # Declared/streamed size limits are a visible metadata-only record, not absent data. + 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" + ) + return record.model_copy( + update={"blobs": (blob,), "content_hash": blob.sha256, "completeness": "complete"} + ) diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 4e59082856..5ade2827f3 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -1,8 +1,10 @@ -"""Slack source implementation for federated search.""" +"""Slack capture of accessible conversations, with optional live search.""" from __future__ import annotations import re +from contextlib import aclosing +from datetime import datetime, timezone from typing import Any, AsyncGenerator, Dict, List, Optional from tenacity import retry, stop_after_attempt @@ -10,8 +12,18 @@ 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 TokenProviderProtocol +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + RemovedScope, + StartedScope, +) +from airweave.domains.sources.exceptions import SourceAuthError, 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 @@ -68,6 +80,15 @@ def _replace(m: re.Match[str]) -> str: return re.sub(r" +", " ", cleaned).strip() +class SlackApiError(ValueError): + """A Slack application error; unlike empty pages it cannot complete a scope.""" + + def __init__(self, code: str): + """Retain the safe provider code for explicit access-loss handling.""" + self.code = code + super().__init__(f"Slack request failed: {code}") + + @source( name="Slack", short_name="slack", @@ -81,22 +102,20 @@ 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.""" - 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. - """ + canonical_record_types = ("channel", "message") + canonical_container_parents = {"message": "channel"} @classmethod async def create( cls, *, - auth: TokenProviderProtocol, + auth: SourceAuthProvider, logger: ContextualLogger, http_client: AirweaveHttpClient, config: SlackConfig, @@ -116,13 +135,11 @@ async def create( ) 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 +147,140 @@ 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=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 + + async def _capture_pages( + self, operation: str, item_key: str, params: dict[str, Any] + ) -> AsyncGenerator[dict, None]: + """Drain a cursor endpoint; incomplete pagination is a failed capture.""" + page_cursor = "" + seen_cursors: set[str] = set() + while True: + page_params = {**params, "limit": 100} + if page_cursor: + page_params["cursor"] = page_cursor + page = await self._get(f"https://slack.com/api/{operation}", page_params) + items = page.get(item_key) + if not isinstance(items, list): + raise ValueError("Slack capture response lacks its record list") + for item in items: + if not isinstance(item, dict): + raise ValueError("Slack capture returned an invalid record") + yield item + page_cursor = page.get("response_metadata", {}).get("next_cursor", "").strip() + if not page_cursor: + if page.get("has_more"): + raise ValueError("Slack capture has more records without a cursor") + return + if page_cursor in seen_cursors: + raise ValueError("Slack capture repeated a pagination cursor") + seen_cursors.add(page_cursor) + + @staticmethod + def _capture_message(message: dict, channel_id: str) -> CaptureRecord: + """Keep the complete native response; file bytes are explicitly not captured yet.""" + 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", + ) + + async def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[CaptureRecord | CompletedScope | StartedScope | RemovedScope, None]: + """Reconcile complete accessible histories, preserving threads and channel identity. + + Slack offers no durable deletion cursor here. Every run re-enumerates history; + a failure aborts reconciliation. Lost channel access is never called deletion. + """ + if node_selections: + raise ValueError("Slack capture does not yet support selected conversation scopes") + seen_channels: set[str] = set() + async with aclosing( + self._capture_pages( + "conversations.list", + "channels", + {"types": "public_channel,private_channel,im,mpim", "exclude_archived": "false"}, + ) + ) as channels: + async for channel in channels: + channel_id = channel["id"] + seen_channels.add(channel_id) + yield CaptureRecord( + identity=RecordIdentity(record_type="channel", native_id=channel_id), + payload=channel, + observed_at=datetime.now(timezone.utc), + ) + yield StartedScope(record_type="message", container_id=channel_id) + async with aclosing( + self._capture_pages( + "conversations.history", "messages", {"channel": channel_id} + ) + ) as history: + async for message in history: + yield self._capture_message(message, channel_id) + if message.get("reply_count", 0): + async with aclosing( + self._capture_pages( + "conversations.replies", + "messages", + {"channel": channel_id, "ts": message["ts"]}, + ) + ) as replies: + async for reply in replies: + yield self._capture_message(reply, channel_id) + yield CompletedScope(record_type="message", container_id=channel_id) + + previous_channels = set(cursor.get().get("channel_ids", [])) if cursor else set() + for channel_id in previous_channels - seen_channels: + await self._confirm_channel_access_lost(channel_id) + observed_at = datetime.now(timezone.utc) + yield CaptureRecord( + identity=RecordIdentity(record_type="channel", native_id=channel_id), + payload={"id": channel_id}, + kind="delete", + removal_reason="access_revoked", + observed_at=observed_at, + ) + yield RemovedScope( + record_type="message", + container_id=channel_id, + removal_reason="access_revoked", + observed_at=observed_at, + ) + if cursor is not None: + cursor.update(channel_ids=sorted(seen_channels)) + + async def _confirm_channel_access_lost(self, channel_id: str) -> None: + """Require provider confirmation before hiding an omitted prior channel.""" + try: + await self._get("https://slack.com/api/conversations.info", {"channel": channel_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") # ------------------------------------------------------------------ # Federated search 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..4f8095ee7f --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py @@ -0,0 +1,371 @@ +"""Gmail capture races and failures against scripted provider responses.""" + +from uuid import uuid4 + +import httpx +import pytest + +from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.platform.cursors.gmail import GmailCursor +from airweave.platform.sources.gmail_capture import BASE, GmailCapture + + +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 + + +def cursor(**data): + return SyncCursor(uuid4(), cursor_schema=GmailCursor, cursor_data=data) + + +@pytest.mark.asyncio +async def test_bootstrap_boundary_precedes_pagination_and_replays_concurrent_edits(): + provider = Provider( + [ + ("/profile", {"historyId": "10"}), + ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "p2"}), + ("/messages/a", message("a", labelIds=["INBOX"])), + ("/messages", {"messages": [{"id": "b"}]}), + ("/messages/b", message("b")), + ( + "/history", + { + "history": [{"labelsRemoved": [{"message": {"id": "a"}}]}], + "nextPageToken": "h2", + "historyId": "11", + }, + ), + ( + "/history", + {"history": [{"messagesAdded": [{"message": {"id": "c"}}]}], "historyId": "12"}, + ), + ("/messages/a", message("a", labelIds=[])), + ("/messages/c", message("c")), + ] + ) + state = cursor() + results = [r async for r in GmailCapture(provider.get, None).generate(state)] + assert [r.payload["id"] for r in results[:-1]] == ["a", "b", "a", "c"] + assert results[2].payload["labelIds"] == [] + assert isinstance(results[-1], CompletedScope) + assert state.data["history_id"] == "12" + assert provider.calls[3][1]["pageToken"] == "p2" + assert provider.calls[6][1] == {"startHistoryId": "10", "maxResults": 500, "pageToken": "h2"} + assert not provider.steps + + +@pytest.mark.asyncio +async def test_incremental_labels_and_deleted_messages_fetch_current_state(): + provider = Provider( + [ + ( + "/history", + { + "historyId": "21", + "history": [ + { + "messagesDeleted": [{"message": {"id": "gone"}}], + "labelsAdded": [{"message": {"id": "a"}}], + "messages": [{"id": "a"}], + } + ], + }, + ), + ("/messages/a", message("a", labelIds=["STARRED"])), + ("/messages/gone", missing()), + ] + ) + state = cursor(history_id="20", canonical_query="") + results = [r async for r in GmailCapture(provider.get, None).generate(state)] + assert len(results) == 2 + assert results[1].kind == "delete" + assert results[1].removal_reason == "provider_deleted" + assert state.data["history_id"] == "21" + + +@pytest.mark.asyncio +async def test_expired_incremental_boundary_bootstraps_and_reconciles(): + provider = Provider( + [ + ("/history", missing()), + ("/profile", {"historyId": "30"}), + ("/messages", {}), + ("/history", {"historyId": "31"}), + ] + ) + state = cursor(history_id="20", canonical_query="") + results = [r async for r in GmailCapture(provider.get, None).generate(state)] + assert results == [CompletedScope(record_type="message")] + assert state.data["history_id"] == "31" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [RuntimeError("network"), missing()]) +async def test_failure_after_full_crawl_does_not_finish_scope_or_advance_checkpoint(error): + provider = Provider( + [ + ("/profile", {"historyId": "30"}), + ("/messages", {"messages": [{"id": "a"}]}), + ("/messages/a", message("a")), + ("/history", error), + ] + ) + state = cursor() + before = state.data + observed = [] + with pytest.raises(type(error)): + async for record in GmailCapture(provider.get, None).generate(state): + observed.append(record) + assert len(observed) == 1 + assert state.data == before + + +@pytest.mark.asyncio +async def test_filtered_query_always_reconciles_and_external_mime_bytes_are_partial(): + raw = message("a", payload={"parts": [{"body": {"attachmentId": "external", "size": 3}}]}) + provider = Provider( + [ + ("/messages", {"messages": [{"id": "a"}]}), + ("/messages/a", raw), + ] + ) + state = cursor(history_id="20", canonical_query="") + results = [ + r async for r in GmailCapture(provider.get, "from:person@example.com").generate(state) + ] + assert results[0].payload == raw + assert results[0].completeness == "partial" + assert results[0].blobs == () + assert results[-1] == CompletedScope(record_type="message") + assert provider.calls[0][1]["q"] == "from:person@example.com" + assert state.data["history_id"] == "" + + +@pytest.mark.asyncio +async def test_legacy_entity_cursor_cannot_skip_canonical_bootstrap(): + provider = Provider( + [ + ("/profile", {"historyId": "30"}), + ("/messages", {}), + ("/history", {"historyId": "31"}), + ] + ) + results = [r async for r in GmailCapture(provider.get, None).generate(cursor(history_id="20"))] + assert results == [CompletedScope(record_type="message")] + + +@pytest.mark.asyncio +async def test_repeated_page_token_fails_without_advancing_cursor(): + provider = Provider( + [ + ("/history", {"nextPageToken": "same"}), + ("/history", {"nextPageToken": "same"}), + ] + ) + state = cursor(history_id="20", canonical_query="") + with pytest.raises(ValueError, match="repeated"): + _ = [r async for r in GmailCapture(provider.get, None).generate(state)] + assert state.data["history_id"] == "20" + + +@pytest.mark.asyncio +async def test_incremental_detail_failure_retains_old_boundary(): + provider = Provider( + [ + ("/history", {"historyId": "21", "history": [{"messages": [{"id": "a"}]}]}), + ("/messages/a", RuntimeError("upstream unavailable")), + ] + ) + state = cursor(history_id="20", canonical_query="") + with pytest.raises(RuntimeError): + _ = [r async for r in GmailCapture(provider.get, None).generate(state)] + assert state.data["history_id"] == "20" + + +@pytest.mark.asyncio +async def test_filtered_page_failure_never_claims_completed_scope(): + provider = Provider( + [ + ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "p2"}), + ("/messages/a", message("a")), + ("/messages", RuntimeError("upstream unavailable")), + ] + ) + state = cursor(history_id="20", canonical_query="") + results = [] + with pytest.raises(RuntimeError): + async for record in GmailCapture(provider.get, "in:inbox").generate(state): + results.append(record) + assert len(results) == 1 + assert not any(isinstance(r, CompletedScope) for r in results) + assert state.data["history_id"] == "20" + + +@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", + "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].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}})), + ] + ) + state = cursor(history_id="20", canonical_query="") + attachment_get = AsyncMock(return_value={"data": "YWJj", "size": 3}) + capture = GmailCapture(provider.get, None, files=files, attachment_get=attachment_get) + with pytest.raises(OSError): + _ = [r async for r in capture.generate(state)] + assert state.data["history_id"] == "20" + + +@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 diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py new file mode 100644 index 0000000000..340757adde --- /dev/null +++ b/backend/airweave/platform/sources/wispr.py @@ -0,0 +1,245 @@ +"""Capture Wispr's available meeting representations through account-bound sessions.""" + +from __future__ import annotations + +import re +from contextlib import aclosing +from datetime import datetime, timezone +from typing import AsyncGenerator +from urllib.parse import quote + +from pydantic import BaseModel, ConfigDict, JsonValue + +from airweave.core.logging import ContextualLogger +from airweave.domains.browse_tree.types import NodeSelectionData +from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity +from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider +from airweave.domains.storage.file_service import FileService +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.platform.decorators import source +from airweave.platform.http_client.airweave_client import AirweaveHttpClient +from airweave.platform.http_client.composio_transport import ComposioTransport +from airweave.platform.sources._base import BaseSource +from airweave.schemas.source_connection import AuthenticationMethod + + +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]] + has_more: bool + next_cursor: str | None = None + truncated: bool = False + + +@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",) + canonical_container_parents: dict[str, str] = {} + + @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 + return instance + + async def _post(self, path: str, body: dict) -> dict: + response = await self.http_client.post( + "https://backend.composio.dev/api/v3.1/tool_router/" + path, + headers={"x-api-key": self._tool_auth.api_key.get_secret_value()}, + json=body, + ) + ComposioTransport._check_proxy_response(response, response.request) + return response.json() + + 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 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 actually read meetings.""" + MeetingPage.model_validate( + await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", {"limit": 1}) + ) + + async def _list_window( + self, + since: datetime | None, + until: datetime | None, + ) -> tuple[list[dict[str, JsonValue]], bool]: + arguments: dict[str, JsonValue] = {"limit": 200} + if since is not None: + arguments["since"] = since.isoformat() + if until is not None: + arguments["until"] = until.isoformat() + records: list[dict[str, JsonValue]] = [] + seen_cursors: set[str] = set() + while True: + page = MeetingPage.model_validate( + await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", arguments) + ) + records.extend(page.meetings) + if page.truncated: + return records, True + if not page.has_more: + return records, False + if not page.next_cursor or page.next_cursor in seen_cursors: + raise ValueError("Wispr returned incomplete or repeated pagination") + if len(records) >= 1000: + return records, True + seen_cursors.add(page.next_cursor) + arguments["cursor"] = page.next_cursor + + async def _list_all(self) -> AsyncGenerator[dict[str, JsonValue], None]: + windows: list[tuple[datetime | None, datetime | None, int]] = [(None, None, 0)] + seen_ids: set[str] = set() + while windows: + since, until, depth = windows.pop() + records, truncated = await self._list_window(since, until) + if truncated: + starts = [ + datetime.fromisoformat(str(row["start"]).replace("Z", "+00:00")) + for row in records + ] + if not starts or min(starts) == max(starts) or depth >= 32: + raise ValueError("Wispr query cap cannot be partitioned safely") + split = min(starts) + (max(starts) - min(starts)) / 2 + if (since is not None and split <= since) or (until is not None and split >= until): + raise ValueError("Wispr date partition made no progress") + windows.extend([(since, split, depth + 1), (split, until, depth + 1)]) + continue + for record in records: + identity = record.get("id") + if not isinstance(identity, str) or not identity: + raise ValueError("Wispr returned a meeting without identity") + if identity not in seen_ids: + seen_ids.add(identity) + yield record + + @staticmethod + def _next_offset(text: JsonValue, field: str) -> int | None: + if not isinstance(text, str): + raise ValueError("Wispr meeting text is not a string") + match = re.search( + rf"\(\.\.\.truncated, \d+ chars remaining; continue with " + rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", + text, + re.MULTILINE, + ) + return int(match.group(1)) if match else None + + async def _meeting(self, identity: str) -> dict[str, JsonValue]: + offsets = {"content": 0, "transcript": 0} + responses: list[JsonValue] = [] + for _ in range(100): + arguments = { + "meeting_id": identity, + **{ + "view_" + field: {"char_limit": 40000, "start_char": offset} + for field, offset in offsets.items() + }, + } + response = await self._execute("WISPR_FLOW_MCP_GET_MEETING", arguments) + if response.get("id") != identity: + raise ValueError("Wispr returned a different meeting identity") + if responses and response.get("modified_at") != responses[0]["response"].get( + "modified_at" + ): + raise ValueError("Wispr meeting changed during paginated capture") + responses.append({"requested_ranges": arguments, "response": response}) + following = {} + for field, offset in offsets.items(): + 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} + # Request only continuing ranges; complete ranges already remain in responses. + offsets = following + raise ValueError("Wispr meeting exceeds the bounded capture page limit") + + async def generate_observations( + self, + *, + cursor: SyncCursor | None = None, + files: FileService | None = None, + node_selections: list[NodeSelectionData] | None = None, + ) -> AsyncGenerator[CaptureRecord, None]: + """Capture full exposed ranges without claiming deletion or snapshot completeness.""" + if node_selections: + raise ValueError("Wispr capture does not support selected meeting scopes") + async with aclosing(self._list_all()) as meetings: + async for listing in meetings: + identity = str(listing["id"]) + payload = await self._meeting(identity) + payload["listing"] = listing + yield CaptureRecord( + identity=RecordIdentity(record_type="meeting", native_id=identity), + payload=payload, + completeness="partial", + observed_at=datetime.now(timezone.utc), + ) diff --git a/backend/alembic/versions/0001_canonical_records.py b/backend/alembic/versions/0001_canonical_records.py index 63f9b17d11..c439d2ff7b 100644 --- a/backend/alembic/versions/0001_canonical_records.py +++ b/backend/alembic/versions/0001_canonical_records.py @@ -30,6 +30,9 @@ def upgrade(): 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"), diff --git a/backend/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py new file mode 100644 index 0000000000..506b6546c3 --- /dev/null +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -0,0 +1,308 @@ +"""Exercise credential custody, provider errors, and actual proxy binary shapes.""" + +import json +from unittest.mock import AsyncMock + +import httpx +import pytest + +from airweave.core.shared_models import SourceConnectionErrorCategory +from airweave.domains.auth_provider.exceptions import ( + AuthProviderAuthError, + AuthProviderRateLimitError, + AuthProviderServerError, +) +from airweave.domains.sources.exceptions.classifier import classify_error +from airweave.domains.sources.token_providers.protocol import ( + ManagedAuthProvider, + authorization_headers, +) +from airweave.platform.http_client.composio_transport import ComposioProxyError, ComposioTransport + + +@pytest.mark.asyncio +async def test_account_binding_headers_and_repeated_query(): + async def proxy(request): + data = json.loads(request.content) + assert data["connected_account_id"] == "ca_bound" + assert data["endpoint"].endswith("?label=A&label=B") + assert all( + x["name"].lower() not in {"authorization", "cookie", "x-api-key"} + for x in data["parameters"] + ) + assert all(x["type"] == "header" and "in" not in x for x in data["parameters"]) + assert request.headers["x-api-key"] == "private-key" + return httpx.Response( + 200, + json={"status": 429, "data": {"error": "limited"}, "headers": {"retry-after": "17"}}, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key="private-key", + connected_account_id="ca_bound", + allowed_hosts={"gmail.googleapis.com"}, + ) + ) as client: + result = await client.get( + "https://gmail.googleapis.com/gmail/v1/users/me?label=A&label=B", + headers={"authorization": "must-not-forward", "cookie": "secret"}, + ) + assert result.status_code == 429 + assert result.headers["retry-after"] == "17" + assert not upstream.is_closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status,exception,category", + [ + ( + 401, + AuthProviderAuthError, + SourceConnectionErrorCategory.AUTH_PROVIDER_CREDENTIALS_INVALID, + ), + (429, AuthProviderRateLimitError, SourceConnectionErrorCategory.RATE_LIMITED), + (503, AuthProviderServerError, None), + ], +) +async def test_proxy_error_classification(status, exception, category): + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda _: httpx.Response(status)) + ) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"slack.com"}, + ) + ) as client: + with pytest.raises(exception) as caught: + await client.get("https://slack.com/api/auth.test") + assert classify_error(caught.value).category == category + + +@pytest.mark.asyncio +async def test_binary_envelope_download_does_not_leak_credentials(monkeypatch, caplog): + checked = [] + monkeypatch.setattr( + "airweave.platform.http_client.composio_transport.validate_url", checked.append + ) + + async def proxy(_): + return httpx.Response( + 200, + json={ + "status": 200, + "data": {}, + "binary_data": { + "url": "https://files.example/file?signature=secret", + "size": 7, + "content_type": "application/pdf", + "expires_at": "2026-10-01", + }, + "headers": {"content-encoding": "gzip"}, + }, + ) + + async def download(request): + assert "authorization" not in request.headers + assert "x-api-key" not in request.headers + assert "cookie" not in request.headers + return httpx.Response(200, content=b"%PDF-ok") + + async with ( + httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as upstream, + ): + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + download_transport=httpx.MockTransport(download), + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"www.googleapis.com"}, + ) + ) as client: + result = await client.get("https://www.googleapis.com/drive/v3/files/id/export") + assert result.content == b"%PDF-ok" + assert result.headers["content-type"] == "application/pdf" + assert "content-encoding" not in result.headers + assert checked == ["https://files.example/file?signature=secret"] + assert "signature=secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_reject_cross_provider_before_proxy(): + proxy = AsyncMock() + async with httpx.AsyncClient( + transport=ComposioTransport( + client=proxy, + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"slack.com"}, + ) + ) as client: + with pytest.raises(ComposioProxyError): + await client.get("https://evil.example/api") + proxy.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_explicit_managed_auth_has_no_token(): + auth = ManagedAuthProvider( + api_key="secret", connected_account_id="ca_bound", allowed_hosts=frozenset({"slack.com"}) + ) + assert await authorization_headers(auth) == {} + assert not hasattr(auth, "get_token") + assert "secret" not in repr(auth) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", ["plain text", {"items": [1]}]) +async def test_plaintext_and_json_envelopes(data): + async with httpx.AsyncClient( + transport=httpx.MockTransport( + lambda _: httpx.Response(200, json={"status": 200, "data": data}) + ) + ) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"www.googleapis.com"}, + ) + ) as client: + result = await client.get("https://www.googleapis.com/drive/v3/files/id/export") + assert (result.text if isinstance(data, str) else result.json()) == data + + +@pytest.mark.asyncio +async def test_download_rejects_private_redirect(): + from airweave.platform.http_client.composio_transport import BinaryData + from airweave.platform.utils.ssrf import SSRFViolation + + calls = [] + + async def download(request): + calls.append(str(request.url)) + return httpx.Response(302, headers={"location": "https://169.254.169.254/latest"}) + + async with ComposioTransport( + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"www.googleapis.com"}, + download_transport=httpx.MockTransport(download), + ) as transport: + with pytest.raises(SSRFViolation): + await transport._download( + BinaryData( + url="https://files.example/file", + size=1, + content_type="application/pdf", + expires_at="2026-10-01", + ) + ) + assert calls == ["https://files.example/file"] + + +@pytest.mark.asyncio +async def test_download_stream_enforces_bytes_even_when_size_is_wrong(): + from airweave.platform.http_client.composio_transport import DownloadStream + + response = httpx.Response(200, content=b"12345") + stream = DownloadStream(response, max_bytes=3) + try: + with pytest.raises(ComposioProxyError, match="size limit"): + async for _ in stream: + pass + finally: + await stream.aclose() + assert response.is_closed + + +@pytest.mark.asyncio +async def test_direct_oauth_authentication_still_works(): + from airweave.domains.sources.token_providers.static import StaticTokenProvider + + assert await authorization_headers(StaticTokenProvider("direct")) == { + "Authorization": "Bearer direct" + } + + +@pytest.mark.asyncio +async def test_managed_provider_account_binding(monkeypatch): + from airweave.domains.auth_provider.exceptions import AuthProviderConfigError + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "key"}, config={"account_id": "ca_bound"} + ) + fetch = AsyncMock(return_value={"toolkit": {"slug": "gmail"}, "status": "ACTIVE"}) + monkeypatch.setattr(provider, "_get_with_auth", fetch) + result = await provider.get_auth_result("gmail", ["access_token"]) + assert result.credentials is None + assert result.managed_auth.connected_account_id == "ca_bound" + assert result.managed_auth.allowed_hosts == {"gmail.googleapis.com"} + with pytest.raises(AuthProviderConfigError, match="toolkit"): + await provider.get_auth_result("google_drive", ["access_token"]) + + +@pytest.mark.asyncio +async def test_expired_account_is_auth_failure_not_generic_422(): + async with httpx.AsyncClient( + transport=httpx.MockTransport( + lambda _: httpx.Response( + 422, json={"error": {"code": 1706, "slug": "TOOL_AUTH_BadConnectedAccountState"}} + ) + ) + ) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"slack.com"}, + ) + ) as client: + with pytest.raises(AuthProviderAuthError, match="reauthorization"): + await client.get("https://slack.com/api/auth.test") + + +@pytest.mark.asyncio +async def test_proxy_envelope_is_bounded_before_json_decode(): + closed = [] + + class TooLarge(httpx.AsyncByteStream): + def __aiter__(self): + return self + + async def __anext__(self): + if closed: + raise StopAsyncIteration + return b"x" * (1024 * 1024 + 5) + + async def aclose(self): + closed.append(True) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda _: httpx.Response(200, stream=TooLarge())) + ) as upstream: + transport = ComposioTransport( + client=upstream, + api_key="key", + connected_account_id="ca_bound", + allowed_hosts={"gmail.googleapis.com"}, + ) + transport.MAX_BINARY_BYTES = 1 + async with httpx.AsyncClient(transport=transport) as client: + with pytest.raises(ComposioProxyError, match="size limit"): + await client.get("https://gmail.googleapis.com/gmail/v1/users/me/messages/id") + assert closed + # Flush HTTPX's nested async iterator finalizers before this test loop closes. + import asyncio + + await asyncio.get_running_loop().shutdown_asyncgens() diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py new file mode 100644 index 0000000000..9996c35ef7 --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -0,0 +1,122 @@ +"""Calendar races and resets, with no provider mutation or live credentials.""" + +from copy import deepcopy +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RemovedScope, + StartedScope, +) +from airweave.domains.sources.exceptions import SourceGoneError, SourceServerError +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.platform.cursors.google_calendar import GoogleCalendarCursor +from airweave.platform.sources.records.google_calendar import generate_calendar_observations + + +class ScriptedGet: + def __init__(self, replies): + self.replies = iter(replies) + self.calls = [] + + async def __call__(self, url, params): + self.calls.append((url, deepcopy(params))) + value = next(self.replies) + if isinstance(value, Exception): + raise value + return value + + +def cursor(tokens=None): + return SyncCursor(uuid4(), GoogleCalendarCursor, {"calendar_tokens": tokens or {}}) + + +async def test_full_pagination_retains_sparse_cancellation_and_container(): + get = ScriptedGet( + [ + {"items": [{"id": "team/a@example.com"}]}, + { + "items": [{"id": "series", "recurrence": ["RRULE:FREQ=WEEKLY"]}], + "nextPageToken": "page2", + }, + { + "items": [ + { + "id": "exception", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-10-01"}, + } + ], + "nextSyncToken": "boundary", + }, + ] + ) + state = cursor() + results = [item async for item in generate_calendar_observations(get, state)] + events = [ + item + for item in results + if isinstance(item, CaptureRecord) and item.identity.record_type == "event" + ] + assert all(item.identity.container_id == "team/a@example.com" for item in events) + assert events[1].kind == "delete" and events[1].payload["recurringEventId"] == "series" + assert "%2F" in get.calls[1][0] + assert get.calls[2][1]["pageToken"] == "page2" + assert state.data["calendar_tokens"] == {"team/a@example.com": "boundary"} + assert any(isinstance(item, CompletedScope) and item.record_type == "event" for item in results) + + +async def test_410_after_delta_page_resets_only_that_scope(): + get = ScriptedGet( + [ + {"items": [{"id": "a"}, {"id": "b"}]}, + {"items": [{"id": "stale"}], "nextPageToken": "delta2"}, + SourceGoneError("expired"), + {"items": [{"id": "current"}], "nextSyncToken": "a-new"}, + {"items": [], "nextSyncToken": "b-new"}, + ] + ) + state = cursor({"a": "a-old", "b": "b-old"}) + results = [item async for item in generate_calendar_observations(get, state)] + starts = [ + item.container_id + for item in results + if type(item) is StartedScope and item.record_type == "event" + ] + assert starts == ["a"] + assert get.calls[2][1]["syncToken"] == "a-old" + assert "syncToken" not in get.calls[3][1] and "pageToken" not in get.calls[3][1] + assert get.calls[4][1]["syncToken"] == "b-old" + assert state.data["calendar_tokens"] == {"a": "a-new", "b": "b-new"} + + +async def test_failure_keeps_original_cursor_and_does_not_complete_scope(): + get = ScriptedGet( + [ + {"items": [{"id": "a"}]}, + {"items": [{"id": "one"}], "nextPageToken": "next"}, + SourceServerError("temporary", status_code=503), + ] + ) + state = cursor() + observed = [] + with pytest.raises(SourceServerError): + async for item in generate_calendar_observations(get, state): + observed.append(item) + assert state.data["calendar_tokens"] == {} + assert not any(type(item) is CompletedScope for item in observed) + + +async def test_missing_calendar_removes_only_previous_calendar_children(): + get = ScriptedGet([{"items": [{"id": "kept"}]}, {"items": [], "nextSyncToken": "new"}]) + state = cursor({"kept": "old", "removed": "old-removed"}) + results = [item async for item in generate_calendar_observations(get, state)] + removed = [item for item in results if type(item) is RemovedScope] + assert [(item.record_type, item.container_id, item.removal_reason) for item in removed] == [ + ("event", "removed", "scope_removed") + ] + assert state.data["calendar_tokens"] == {"kept": "new"} diff --git a/backend/tests/unit/platform/sources/records/test_google_drive.py b/backend/tests/unit/platform/sources/records/test_google_drive.py new file mode 100644 index 0000000000..9dd8c63d9b --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_google_drive.py @@ -0,0 +1,108 @@ +"""Drive bootstrap/delta fidelity, scope completeness and checkpoint failure behavior.""" + +from copy import deepcopy +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.platform.cursors.google_drive import GoogleDriveCursor +from airweave.platform.sources.records.google_drive import generate_drive_observations + + +class ScriptedGet: + def __init__(self, replies): + self.replies = iter(replies) + self.calls = [] + + async def __call__(self, url, params): + self.calls.append((url, deepcopy(params))) + value = next(self.replies) + if isinstance(value, Exception): + raise value + return value + + +def state(token=""): + return SyncCursor(uuid4(), GoogleDriveCursor, {"canonical_page_token": token}) + + +async def test_initial_boundary_precedes_scan_and_replay_updates_metadata(): + get = ScriptedGet( + [ + {"startPageToken": "before"}, + { + "files": [{"id": "one", "name": "old", "parents": ["folder-a"]}], + "nextPageToken": "page2", + }, + {"files": []}, + { + "changes": [ + {"file": {"id": "one", "name": "new", "parents": ["folder-b"], "trashed": True}} + ], + "newStartPageToken": "after", + }, + ] + ) + cursor = state() + results = [item async for item in generate_drive_observations(get, cursor)] + records = [item for item in results if isinstance(item, CaptureRecord)] + assert get.calls[0][0].endswith("startPageToken") + assert get.calls[-1][1]["pageToken"] == "before" + assert [item.payload["name"] for item in records] == ["old", "new"] + assert records[-1].kind == "upsert" and records[-1].payload["trashed"] is True + assert all(item.completeness == "metadata_only" for item in records) + assert cursor.data["canonical_page_token"] == "after" + assert type(results[-1]) is CompletedScope + + +async def test_incremental_removal_pagination_and_no_full_scan(): + get = ScriptedGet( + [ + {"changes": [{"fileId": "gone", "removed": True}], "nextPageToken": "next"}, + {"changes": [], "newStartPageToken": "after"}, + ] + ) + cursor = state("before") + results = [item async for item in generate_drive_observations(get, cursor)] + assert len(results) == 1 and results[0].kind == "delete" + assert results[0].removal_reason == "scope_removed" + assert cursor.data["canonical_page_token"] == "after" + assert all(url.endswith("/changes") for url, _ in get.calls) + + +async def test_incomplete_enumeration_cannot_mark_scope_complete_or_save_cursor(): + get = ScriptedGet([{"startPageToken": "before"}, {"files": [], "incompleteSearch": True}]) + cursor = state() + observed = [] + with pytest.raises(ValueError, match="incomplete enumeration"): + async for item in generate_drive_observations(get, cursor): + observed.append(item) + assert cursor.data["canonical_page_token"] == "" + assert [type(item) for item in observed] == [StartedScope] + + +async def test_shared_drive_membership_change_restarts_with_new_boundary(): + get = ScriptedGet( + [ + {"changes": [{"changeType": "drive", "removed": True, "driveId": "lost"}]}, + {"startPageToken": "new-before"}, + {"files": []}, + {"changes": [], "newStartPageToken": "new-after"}, + ] + ) + cursor = state("old") + results = [item async for item in generate_drive_observations(get, cursor)] + assert [type(item) for item in results] == [StartedScope, CompletedScope] + assert cursor.data["canonical_page_token"] == "new-after" + + +async def test_failed_change_page_retains_saved_checkpoint(): + get = ScriptedGet( + [{"changes": [], "nextPageToken": "next"}, ConnectionError("provider unavailable")] + ) + cursor = state("before") + with pytest.raises(ConnectionError): + _ = [item async for item in generate_drive_observations(get, cursor)] + assert cursor.data["canonical_page_token"] == "before" diff --git a/backend/tests/unit/platform/sources/records/test_google_drive_content.py b/backend/tests/unit/platform/sources/records/test_google_drive_content.py new file mode 100644 index 0000000000..8d22752dd0 --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_google_drive_content.py @@ -0,0 +1,74 @@ +"""Body completeness requires durable bytes matching the observed file version.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from airweave.domains.entities.canonical.requests import BlobReference +from airweave.domains.storage.exceptions import FileSkippedException +from airweave.platform.sources.records.google_drive import file_record +from airweave.platform.sources.records.google_drive_content import capture_file_content + + +@pytest.fixture +def inputs(): + blob = BlobReference( + key="canonical/test/blobs/sha256/" + "a" * 64, + sha256="a" * 64, + size_bytes=3, + media_type="text/plain", + ) + return { + "files": SimpleNamespace(capture_canonical_url=AsyncMock(return_value=blob)), + "get": AsyncMock(return_value={"version": "7"}), + "client": object(), + "auth": object(), + "logger": object(), + } + + +async def test_complete_body_keeps_original_metadata_and_blob(inputs): + record = file_record( + {"id": "file/id", "mimeType": "text/plain", "version": "7", "name": "notes"} + ) + result = await capture_file_content(record, **inputs) + assert result.completeness == "complete" and result.blobs[0].size_bytes == 3 + assert result.payload == record.payload and result.content_hash == "a" * 64 + assert "file%2Fid" in inputs["files"].capture_canonical_url.call_args.kwargs["url"] + + +async def test_changed_version_cannot_claim_complete(inputs): + inputs["get"].return_value = {"version": "8"} + record = file_record({"id": "file", "mimeType": "text/plain", "version": "7"}) + with pytest.raises(ValueError, match="changed during content capture"): + await capture_file_content(record, **inputs) + + +async def test_oversized_body_remains_explicitly_metadata_only(inputs): + inputs["files"].capture_canonical_url.side_effect = FileSkippedException( + reason="too large", filename="file" + ) + record = file_record({"id": "file", "mimeType": "text/plain", "version": "7"}) + result = await capture_file_content(record, **inputs) + assert result.completeness == "metadata_only" and not result.blobs + + +async def test_storage_failure_is_not_silently_a_partial_success(inputs): + inputs["files"].capture_canonical_url.side_effect = ConnectionError("storage unavailable") + record = file_record({"id": "file", "mimeType": "text/plain", "version": "7"}) + with pytest.raises(ConnectionError): + await capture_file_content(record, **inputs) + + +async def test_native_document_uses_existing_export_format(inputs): + record = file_record( + {"id": "doc", "mimeType": "application/vnd.google-apps.document", "version": "7"} + ) + await capture_file_content(record, **inputs) + call = inputs["files"].capture_canonical_url.call_args.kwargs + assert "/export?" in call["url"] + assert ( + call["media_type"] + == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + ) diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py new file mode 100644 index 0000000000..360e1c8e19 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -0,0 +1,120 @@ +"""Thread preservation and failure-safe Slack capture.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.platform.sources.slack import SlackSource + + +def source(): + return SlackSource( + auth=StaticTokenProvider("token"), logger=MagicMock(), http_client=MagicMock() + ) + + +@pytest.mark.asyncio +async def test_channel_pagination_threads_and_native_payload(monkeypatch): + connector = source() + calls = [] + + async def get(url, params): + calls.append((url, params)) + if url.endswith("conversations.list"): + return {"channels": [{"id": "C1", "unknown_native_field": "kept"}]} + if url.endswith("conversations.history"): + if params.get("cursor"): + return {"messages": [{"ts": "2", "text": "second", "files": [{"id": "F1"}]}]} + return { + "messages": [{"ts": "1", "reply_count": 1, "blocks": [{"type": "rich_text"}]}], + "response_metadata": {"next_cursor": "next"}, + "has_more": True, + } + return { + "messages": [ + {"ts": "1", "reply_count": 1}, + {"ts": "1.1", "thread_ts": "1", "text": "reply"}, + ] + } + + monkeypatch.setattr(connector, "_get", get) + observations = [x async for x in connector.generate_observations()] + records = [x for x in observations if isinstance(x, CaptureRecord)] + assert records[0].payload["unknown_native_field"] == "kept" + assert any(x.identity.native_id == "1.1" and x.payload["thread_ts"] == "1" for x in records) + assert records[-1].completeness == "partial" + assert observations[-1] == CompletedScope(record_type="message", container_id="C1") + assert calls[0][1]["types"] == "public_channel,private_channel,im,mpim" + + +@pytest.mark.asyncio +async def test_failed_thread_never_completes_scope(monkeypatch): + connector = source() + get = AsyncMock( + side_effect=[ + {"channels": [{"id": "C1"}]}, + {"messages": [{"ts": "1", "reply_count": 1}]}, + RuntimeError("permission lost"), + ] + ) + monkeypatch.setattr(connector, "_get", get) + emitted = [] + with pytest.raises(RuntimeError, match="permission"): + async for observation in connector.generate_observations(): + emitted.append(observation) + assert not any(type(x) is CompletedScope for x in emitted) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "pages", + [ + [{"messages": [], "has_more": True}], + [{"messages": [], "response_metadata": {"next_cursor": "repeat"}}] * 2, + ], +) +async def test_incomplete_pagination_fails(monkeypatch, pages): + connector = source() + monkeypatch.setattr(connector, "_get", AsyncMock(side_effect=pages)) + with pytest.raises(ValueError): + _ = [x async for x in connector._capture_pages("conversations.history", "messages", {})] + + +@pytest.mark.asyncio +async def test_prior_channel_loss_requires_explicit_provider_confirmation(monkeypatch): + from uuid import uuid4 + + from airweave.domains.entities.canonical.requests import RemovedScope + from airweave.domains.syncs.cursors.cursor import SyncCursor + from airweave.platform.sources.slack import SlackApiError + + connector = source() + cursor = SyncCursor(uuid4(), cursor_data={"channel_ids": ["C1"]}) + monkeypatch.setattr( + connector, + "_get", + AsyncMock(side_effect=[{"channels": []}, SlackApiError("channel_not_found")]), + ) + observations = [x async for x in connector.generate_observations(cursor=cursor)] + assert observations[0].kind == "delete" + assert observations[0].removal_reason == "access_revoked" + assert isinstance(observations[1], RemovedScope) + assert cursor.get()["channel_ids"] == [] + + +@pytest.mark.asyncio +async def test_omitted_but_accessible_channel_does_not_advance_cursor(monkeypatch): + from uuid import uuid4 + + from airweave.domains.syncs.cursors.cursor import SyncCursor + + connector = source() + cursor = SyncCursor(uuid4(), cursor_data={"channel_ids": ["C1"]}) + monkeypatch.setattr( + connector, "_get", AsyncMock(side_effect=[{"channels": []}, {"channel": {"id": "C1"}}]) + ) + with pytest.raises(ValueError, match="omitted"): + _ = [x async for x in connector.generate_observations(cursor=cursor)] + assert cursor.get()["channel_ids"] == ["C1"] diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py new file mode 100644 index 0000000000..fd1ec6c69f --- /dev/null +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -0,0 +1,105 @@ +"""Wispr session binding, query caps, and exact transcript continuations.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider +from airweave.platform.sources.wispr import WisprSource + + +async def source(): + return await WisprSource.create( + auth=ManagedToolAuthProvider( + api_key="key", connected_account_id="ca_bound", user_id="user" + ), + logger=MagicMock(), + http_client=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_session_is_bound_and_never_auto_connects(monkeypatch): + connector = await source() + post = AsyncMock( + side_effect=[ + {"session_id": "session"}, + {"data": {"meetings": [], "has_more": False}, "error": None}, + ] + ) + monkeypatch.setattr(connector, "_post", post) + await connector.validate() + payload = post.call_args_list[0].args[1] + assert payload["connected_accounts"] == {"wispr_flow_mcp": ["ca_bound"]} + assert payload["manage_connections"] == {"enable": False} + assert "account" not in post.call_args_list[1].args[1] + + +@pytest.mark.asyncio +async def test_preserves_native_ranges_and_continues_exact_offset(monkeypatch): + connector = await source() + execute = AsyncMock( + side_effect=[ + { + "id": "m", + "modified_at": "same", + "content": "notes", + "transcript": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)\n\nProvider guidance follows." + ), + "native_extra": True, + }, + {"id": "m", "modified_at": "same", "transcript": "def", "content": "notes"}, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + result = await connector._meeting("m") + assert len(result["responses"]) == 2 + assert result["responses"][0]["response"]["native_extra"] is True + assert execute.call_args.args[1]["view_transcript"]["start_char"] == 3 + + +@pytest.mark.asyncio +async def test_changed_meeting_fails_instead_of_mixing_versions(monkeypatch): + connector = await source() + execute = AsyncMock( + side_effect=[ + { + "id": "m", + "modified_at": "old", + "content": "", + "transcript": ( + "(...truncated, 2 chars remaining; " + "continue with view_transcript.start_char=1...)" + ), + }, + {"id": "m", "modified_at": "new", "transcript": "new"}, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + with pytest.raises(ValueError, match="changed"): + await connector._meeting("m") + + +@pytest.mark.asyncio +async def test_query_cap_partitions_by_meeting_start_not_modified(monkeypatch): + connector = await source() + row1 = {"id": "1", "start": "2026-01-01T00:00:00Z", "modified_at": "2026-09-01T00:00:00Z"} + row2 = {"id": "2", "start": "2026-01-03T00:00:00Z"} + listing = AsyncMock(side_effect=[([row1, row2], True), ([row2], False), ([row1], False)]) + monkeypatch.setattr(connector, "_list_window", listing) + assert {row["id"] async for row in connector._list_all()} == {"1", "2"} + assert listing.call_args_list[1].args[0].isoformat() == "2026-01-02T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_unpartitionable_cap_fails(monkeypatch): + connector = await source() + monkeypatch.setattr( + connector, + "_list_window", + AsyncMock(return_value=([{"start": "2026-01-01T00:00:00Z"}], True)), + ) + with pytest.raises(ValueError, match="cap"): + _ = [row async for row in connector._list_all()] From 6380d4666e5279af79e6029aefabe9f415f77404 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 02:40:34 -0700 Subject: [PATCH 004/266] Publish canonical search generations with authoritative visibility checks --- Makefile | 28 +++ backend/airweave/api/v1/endpoints/search.py | 12 +- backend/airweave/core/container/factory.py | 2 + .../entities/canonical/blob_materializer.py | 43 ++++ .../entities/canonical/gmail_projection.py | 160 +++++++++++++ .../entities/canonical/projection_mappers.py | 224 ++++++++++++++++++ .../entities/canonical/projection_models.py | 61 +++++ .../entities/canonical/projection_store.py | 139 +++++++++++ .../domains/entities/canonical/projector.py | 150 ++++++++++++ .../domains/entities/canonical/query_store.py | 23 +- .../entities/canonical/tests/conftest.py | 1 + .../canonical/tests/test_gmail_projection.py | 186 +++++++++++++++ .../canonical/tests/test_metadata_counts.py | 74 ++++++ .../canonical/tests/test_projection.py | 149 ++++++++++++ .../canonical/tests/test_search_visibility.py | 109 +++++++++ .../canonical/tests/test_strict_projection.py | 47 ++++ .../airweave/domains/search/agentic/agent.py | 41 +++- .../domains/search/agentic/context_manager.py | 11 +- .../domains/search/agentic/service.py | 5 + .../airweave/domains/search/agentic/state.py | 3 + .../search/agentic/tests/test_agent.py | 111 ++++++++- .../agentic/tests/test_context_manager.py | 80 +++++-- .../search/agentic/tests/test_tools.py | 67 +++++- .../domains/search/agentic/tools/count.py | 11 +- .../domains/search/agentic/tools/read.py | 16 +- .../domains/search/agentic/tools/search.py | 7 + .../domains/search/agentic/tools/types.py | 2 + .../airweave/domains/search/browse/service.py | 31 +-- .../search/browse/tests/test_service.py | 65 ++++- .../search/builders/collection_metadata.py | 29 ++- .../domains/search/canonical_visibility.py | 99 ++++++++ .../domains/search/classic/service.py | 4 +- backend/airweave/domains/search/executor.py | 25 +- .../tests/test_collection_metadata_builder.py | 31 +++ .../domains/search/tests/test_executor.py | 74 +++--- .../search/tests/test_visible_vector_db.py | 163 +++++++++++++ .../airweave/domains/search/types/metadata.py | 35 ++- .../airweave/domains/search/types/results.py | 9 + .../domains/search/visible_vector_db.py | 138 +++++++++++ .../airweave/domains/storage/file_service.py | 5 +- backend/airweave/domains/storage/limits.py | 3 + .../domains/sync_pipeline/entity/pipeline.py | 38 +-- .../sync_pipeline/pipeline/text_builder.py | 20 +- .../sync_pipeline/processors/chunk_embed.py | 46 +++- .../sync_pipeline/processors/context.py | 26 ++ .../sync_pipeline/processors/entity_fields.py | 26 ++ .../domains/sync_pipeline/processors/utils.py | 10 +- .../activities/project_canonical_records.py | 81 +++++++ .../domains/temporal/worker/wiring.py | 25 ++ .../workflows/project_canonical_records.py | 49 ++++ .../workflows/run_source_connection.py | 47 ++-- .../tests/test_canonical_projection.py | 42 ++++ backend/airweave/models/entity.py | 2 + .../destinations/vespa/destination.py | 4 +- .../airweave/platform/entities/__init__.py | 7 +- .../platform/entities/google_drive.py | 13 + backend/airweave/platform/entities/slack.py | 9 + backend/airweave/platform/entities/wispr.py | 32 +++ .../{sources => http_client}/retry_helpers.py | 0 backend/airweave/platform/sources/airtable.py | 6 +- backend/airweave/platform/sources/apollo.py | 6 +- backend/airweave/platform/sources/asana.py | 6 +- backend/airweave/platform/sources/attio.py | 6 +- .../airweave/platform/sources/bitbucket.py | 6 +- backend/airweave/platform/sources/box.py | 6 +- backend/airweave/platform/sources/clickup.py | 6 +- backend/airweave/platform/sources/coda.py | 6 +- .../airweave/platform/sources/confluence.py | 6 +- backend/airweave/platform/sources/dropbox.py | 6 +- .../airweave/platform/sources/fireflies.py | 6 +- .../airweave/platform/sources/freshdesk.py | 6 +- backend/airweave/platform/sources/github.py | 6 +- backend/airweave/platform/sources/gitlab.py | 6 +- backend/airweave/platform/sources/gmail.py | 6 +- .../platform/sources/google_calendar.py | 8 +- .../airweave/platform/sources/google_docs.py | 6 +- .../airweave/platform/sources/google_drive.py | 8 +- .../platform/sources/google_slides.py | 6 +- backend/airweave/platform/sources/hubspot.py | 6 +- backend/airweave/platform/sources/intercom.py | 6 +- backend/airweave/platform/sources/jira.py | 6 +- backend/airweave/platform/sources/linear.py | 6 +- backend/airweave/platform/sources/monday.py | 6 +- backend/airweave/platform/sources/notion.py | 6 +- backend/airweave/platform/sources/onedrive.py | 8 +- backend/airweave/platform/sources/onenote.py | 6 +- .../platform/sources/outlook_calendar.py | 6 +- .../airweave/platform/sources/outlook_mail.py | 6 +- .../airweave/platform/sources/pipedrive.py | 6 +- .../airweave/platform/sources/powerpoint.py | 6 +- .../airweave/platform/sources/salesforce.py | 6 +- .../airweave/platform/sources/servicenow.py | 6 +- .../airweave/platform/sources/sharepoint.py | 8 +- .../sources/sharepoint2019v2/client.py | 2 +- .../sources/sharepoint_online/client.py | 2 +- .../sources/sharepoint_online/source.py | 8 +- backend/airweave/platform/sources/shopify.py | 6 +- backend/airweave/platform/sources/slab.py | 6 +- backend/airweave/platform/sources/slack.py | 6 +- backend/airweave/platform/sources/slite.py | 6 +- backend/airweave/platform/sources/stripe.py | 6 +- backend/airweave/platform/sources/teams.py | 6 +- backend/airweave/platform/sources/todoist.py | 6 +- backend/airweave/platform/sources/trello.py | 6 +- backend/airweave/platform/sources/word.py | 6 +- backend/airweave/platform/sources/zendesk.py | 6 +- backend/airweave/platform/sources/zoho_crm.py | 6 +- backend/airweave/platform/sources/zoom.py | 6 +- backend/airweave/schemas/search_v2.py | 10 +- .../versions/0002_projection_publication.py | 26 ++ .../entities/test_projection_mappers.py | 175 ++++++++++++++ .../unit/domains/storage/test_import_order.py | 27 +++ 112 files changed, 3140 insertions(+), 344 deletions(-) create mode 100644 Makefile create mode 100644 backend/airweave/domains/entities/canonical/blob_materializer.py create mode 100644 backend/airweave/domains/entities/canonical/gmail_projection.py create mode 100644 backend/airweave/domains/entities/canonical/projection_mappers.py create mode 100644 backend/airweave/domains/entities/canonical/projection_models.py create mode 100644 backend/airweave/domains/entities/canonical/projection_store.py create mode 100644 backend/airweave/domains/entities/canonical/projector.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_metadata_counts.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_search_visibility.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_strict_projection.py create mode 100644 backend/airweave/domains/search/canonical_visibility.py create mode 100644 backend/airweave/domains/search/tests/test_visible_vector_db.py create mode 100644 backend/airweave/domains/search/visible_vector_db.py create mode 100644 backend/airweave/domains/storage/limits.py create mode 100644 backend/airweave/domains/sync_pipeline/processors/context.py create mode 100644 backend/airweave/domains/sync_pipeline/processors/entity_fields.py create mode 100644 backend/airweave/domains/temporal/activities/project_canonical_records.py create mode 100644 backend/airweave/domains/temporal/workflows/project_canonical_records.py create mode 100644 backend/airweave/domains/temporal/workflows/tests/test_canonical_projection.py create mode 100644 backend/airweave/platform/entities/wispr.py rename backend/airweave/platform/{sources => http_client}/retry_helpers.py (100%) create mode 100644 backend/alembic/versions/0002_projection_publication.py create mode 100644 backend/tests/unit/domains/entities/test_projection_mappers.py create mode 100644 backend/tests/unit/domains/storage/test_import_order.py diff --git a/Makefile b/Makefile new file mode 100644 index 0000000000..b4458ebff5 --- /dev/null +++ b/Makefile @@ -0,0 +1,28 @@ +.DEFAULT_GOAL := help +BACKEND := backend +PYTHON := .venv/bin/python +PYTEST := .venv/bin/pytest + +.PHONY: help setup test-store test-capture test-search + +help: + @echo 'setup Install the locked backend and development dependencies' + @echo 'test-store Verify record transactions and publication on disposable PostgreSQL' + @echo 'test-capture Verify source capture, transport and blob handling' + @echo 'test-search Verify search services and visibility behavior' + @echo '' + @echo 'test-store requires CANONICAL_TEST_DATABASE_URL; it creates and removes test schemas.' + @echo 'Deployment is not configured yet. See backend/SERVICE_AUTH.md for service setup.' + +setup: + cd $(BACKEND) && POETRY_VIRTUALENVS_IN_PROJECT=true uvx --from poetry==2.3.2 poetry install --with dev,lint --no-root --no-interaction + +test-store: + @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests + +test-capture: + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_slack_capture.py tests/unit/platform/sources/test_wispr_capture.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + +test-search: + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search diff --git a/backend/airweave/api/v1/endpoints/search.py b/backend/airweave/api/v1/endpoints/search.py index 6ccc06903c..ed06bf8b8e 100644 --- a/backend/airweave/api/v1/endpoints/search.py +++ b/backend/airweave/api/v1/endpoints/search.py @@ -79,7 +79,7 @@ async def instant_search( ) results = await service.search(db, ctx, readable_id, request) - return SearchV2Response(results=results.results) + return SearchV2Response(**results.model_dump()) @router.post( @@ -139,7 +139,7 @@ async def classic_search( ) results = await service.search(db, ctx, readable_id, request) - return SearchV2Response(results=results.results) + return SearchV2Response(**results.model_dump()) @router.post( @@ -179,7 +179,7 @@ async def agentic_search( # Service emits SearchCompletedEvent internally results = await service.search(db, ctx, readable_id, request) truncated = results.results[: request.limit] if request.limit else results.results - return SearchV2Response(results=truncated) + return SearchV2Response(**results.model_copy(update={"results": truncated}).model_dump()) # ── Streaming agentic search ────────────────────────────────────────── @@ -407,7 +407,7 @@ async def admin_instant_search_as_user( results = await service.search( db, ctx, readable_id, request, user_principal_override=user_principal ) - return SearchV2Response(results=results.results) + return SearchV2Response(**results.model_dump()) @admin_router.post( @@ -438,7 +438,7 @@ async def admin_classic_search_as_user( results = await service.search( db, ctx, readable_id, request, user_principal_override=user_principal ) - return SearchV2Response(results=results.results) + return SearchV2Response(**results.model_dump()) @admin_router.post( @@ -470,4 +470,4 @@ async def admin_agentic_search_as_user( db, ctx, readable_id, request, user_principal_override=user_principal ) truncated = results.results[: request.limit] if request.limit else results.results - return SearchV2Response(results=truncated) + return SearchV2Response(**results.model_copy(update={"results": truncated}).model_dump()) diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 94ce802393..8250967620 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -1354,12 +1354,14 @@ def _create_search_services( reranker=reranker, executor=executor, vector_db=vector_db, + source_registry=source_registry, metadata_builder=metadata_builder, collection_repo=collection_repo, event_bus=event_bus, ) browse_service = BrowseService( vector_db=vector_db, + source_registry=source_registry, collection_repo=collection_repo, ) diff --git a/backend/airweave/domains/entities/canonical/blob_materializer.py b/backend/airweave/domains/entities/canonical/blob_materializer.py new file mode 100644 index 0000000000..3562243f8c --- /dev/null +++ b/backend/airweave/domains/entities/canonical/blob_materializer.py @@ -0,0 +1,43 @@ +"""Verify canonical blob identity before materializing disposable projection inputs.""" + +import hashlib +import re +from pathlib import Path + +import aiofiles + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import BlobReference +from airweave.domains.storage.limits import MAX_FILE_SIZE_BYTES +from airweave.domains.storage.protocols import StorageBackend + + +async def read_blob(record: SourceRecord, ref: BlobReference, storage: StorageBackend) -> bytes: + """Read only this sync's content-addressed bytes; reject corrupt references/content.""" + expected = f"canonical/{record.sync_id}/blobs/sha256/{ref.sha256}" + if ref.key != expected: + raise ValueError("Canonical blob key does not match source scope and digest") + if ref not in record.blobs: + raise ValueError("Canonical blob reference does not belong to the source record") + if ref.size_bytes > MAX_FILE_SIZE_BYTES: + raise ValueError("Canonical blob exceeds projection size limit") + content = await storage.read_file(ref.key) + if len(content) != ref.size_bytes or hashlib.sha256(content).hexdigest() != ref.sha256: + raise ValueError("Canonical blob bytes do not match recorded size and digest") + return content + + +async def write_blob(content: bytes, directory: Path, *, suffix: str) -> Path: + """Use trusted extension and content hash, never provider filenames, for local paths.""" + if not re.fullmatch(r"\.[a-z0-9]{1,12}", suffix): + raise ValueError("Projection filename suffix must be one simple extension") + if len(content) > MAX_FILE_SIZE_BYTES: + raise ValueError("Projection content exceeds size limit") + directory.mkdir(parents=True, exist_ok=True) + path = directory / (hashlib.sha256(content).hexdigest() + suffix) + # A disposable per-projection directory is owned by the caller, not provider data. + if path.is_symlink(): + raise ValueError("Projection path must not be a symbolic link") + async with aiofiles.open(path, "wb") as output: + await output.write(content) + return path diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py new file mode 100644 index 0000000000..580b5db5e9 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -0,0 +1,160 @@ +"""Deterministic Gmail search inputs from committed payloads and immutable blobs only.""" + +import base64 +import hashlib +import html +import mimetypes +from datetime import datetime, timezone +from email.message import Message +from email.utils import formataddr, getaddresses +from pathlib import Path + +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.gmail import GmailAttachmentEntity, GmailMessageEntity + + +def _header(part: dict, name: str) -> str: + return next( + (h["value"] for h in part.get("headers", []) if h["name"].lower() == name.lower()), "" + ) + + +def _attachment(part: dict) -> bool: + mime = part.get("mimeType", "").lower() + return ( + bool(part.get("filename")) + or _header(part, "Content-Disposition").lower().startswith("attachment") + or (not mime.startswith("multipart/") and mime not in {"text/plain", "text/html"}) + ) + + +def _has_html(part: dict) -> bool: + """Find an HTML alternative even when wrapped by multipart/related.""" + return part.get("mimeType") == "text/html" or any( + _has_html(child) for child in part.get("parts", []) + ) + + +async def _body( + part: dict, + path: str, + record: SourceRecord, + storage: StorageBackend, +) -> bytes: + body = part.get("body", {}) + if body.get("attachmentId"): + matching = [b for b in record.blobs if b.source_path == path + "/body"] + if len(matching) != 1: + raise ValueError("Gmail MIME body requires exactly one stored blob reference") + content = await read_blob(record, matching[0], storage) + else: + encoded = body.get("data", "") + content = base64.b64decode( + encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True + ) + if len(content) != int(body.get("size", 0)): + raise ValueError("Gmail MIME size differs from committed provider metadata") + return content + + +def _decode_text(content: bytes, part: dict) -> str: + header = Message() + header["Content-Type"] = _header(part, "Content-Type") or part["mimeType"] + return content.decode(header.get_content_charset() or "utf-8", errors="strict") + + +class GmailProjection: + """One message projection; per-call state never escapes into the source record.""" + + def __init__(self, record: SourceRecord, storage: StorageBackend, directory: Path): + """Bind committed source authority and disposable output directory.""" + self.record = record + self.storage = storage + self.directory = directory + self.attachments: list[GmailAttachmentEntity] = [] + + async def attachment(self, part: dict, path: str) -> None: + """Retain each MIME attachment identity, including duplicate filenames.""" + content = await _body(part, path, self.record, self.storage) + filename = part.get("filename") or "attachment" + mime = part.get("mimeType", "application/octet-stream") + suffix = Path(filename).suffix.lower() or mimetypes.guess_extension(mime) or ".bin" + local = await write_blob(content, self.directory, suffix=suffix) + message_id = self.record.identity.native_id + self.attachments.append( + GmailAttachmentEntity( + breadcrumbs=[ + Breadcrumb( + entity_id=f"msg_{message_id}", + name="Email message", + entity_type="GmailMessageEntity", + ) + ], + attachment_key=f"attach_{message_id}_{hashlib.sha256(path.encode()).hexdigest()}", + filename=filename, + message_id=message_id, + attachment_id=part.get("body", {}).get("attachmentId") + or part.get("partId") + or path, + thread_id=self.record.payload["threadId"], + url=f"https://mail.google.com/mail/u/0/#inbox/{message_id}", + size=len(content), + file_type=suffix.lstrip("."), + mime_type=mime, + local_path=str(local), + ) + ) + + async def render(self, part: dict, path: str) -> str: + """Select MIME alternatives, preserve mixed bodies, and materialize attachments.""" + mime = part.get("mimeType", "").lower() + if mime.startswith("multipart/"): + children = list(enumerate(part.get("parts", []))) + if mime == "multipart/alternative": + html_parts = [(i, p) for i, p in children if _has_html(p)] + plain_parts = [(i, p) for i, p in children if p.get("mimeType") == "text/plain"] + # Prefer one complete alternative, not duplicate plain and HTML text. + children = html_parts[-1:] or plain_parts[-1:] or children[-1:] + return "\n".join( + [await self.render(child, f"{path}/parts/{index}") for index, child in children] + ) + if _attachment(part): + await self.attachment(part, path) + return "" + content = await _body(part, path, self.record, self.storage) + text = _decode_text(content, part) + return text if mime == "text/html" else "
" + html.escape(text) + "
" + + async def map(self) -> tuple[BaseEntity, ...]: + """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 != "complete": + raise ValueError("Gmail projection requires a complete active 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 (entity, *self.attachments) + + +async def map_gmail( + record: SourceRecord, + storage: StorageBackend, + directory: Path, +) -> tuple[BaseEntity, ...]: + """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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py new file mode 100644 index 0000000000..9980254a3d --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -0,0 +1,224 @@ +"""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 + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.storage.protocols import StorageBackend +from airweave.platform.entities._base import BaseEntity +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 + + +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 + + +async def _drive( + record: SourceRecord, + storage: StorageBackend, + directory: Path, +) -> tuple[BaseEntity, ...]: + from airweave.domains.entities.canonical.blob_materializer import read_blob, write_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 ( + GoogleDriveFolderEntity( + folder_id=record.identity.native_id, + title=_string(data.get("name"), default="Untitled"), + description=_string(data.get("description")), + breadcrumbs=[], + ), + ) + 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: + raise ProjectionMappingError("Drive captured bytes have no known file format") + suffix = suffix.lower() + 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 (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}} + entity = SlackMessageEntity.from_api(metadata, breadcrumbs=[]) + 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_text(responses: list[JsonValue], field: str) -> str: + """Follow exactly the requested ranges; strip server continuation guidance.""" + expected = 0 + complete = False + pieces: list[str] = [] + for page in responses: + if not isinstance(page, dict): + raise ProjectionMappingError("Invalid Wispr response range") + arguments = page.get("requested_ranges") + response = page.get("response") + if not isinstance(arguments, dict) or not isinstance(response, dict): + raise ProjectionMappingError("Missing Wispr native range response") + window = arguments.get("view_" + field) + if window is None: + continue + if not isinstance(window, dict) or window.get("start_char") != expected or complete: + raise ProjectionMappingError("Wispr text ranges overlap or have a gap") + text = _string(response.get(field)) + marker = re.search( + rf"(?m)^\(\.\.\.truncated, \d+ chars remaining; continue with " + rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", + text, + ) + if marker: + following = int(marker.group(1)) + if following <= expected: + raise ProjectionMappingError("Wispr continuation does not advance") + pieces.append(text[: marker.start()].rstrip()) + expected = following + else: + pieces.append(text) + complete = True + if not complete: + raise ProjectionMappingError("Wispr text has an unfinished continuation") + return "\n".join(pieces) + + +def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: + if record.identity.record_type != "meeting": + raise ProjectionMappingError("Unsupported Wispr record type") + responses = record.payload.get("responses") + if not isinstance(responses, list) or not responses or not isinstance(responses[0], dict): + raise ProjectionMappingError("Wispr has no captured meeting content") + first = responses[0].get("response") + if not isinstance(first, dict): + raise ProjectionMappingError("Wispr meeting response is invalid") + return ( + WisprMeetingEntity( + meeting_id=record.identity.native_id, + title=_string(first.get("title"), default="Meeting"), + notes=_wispr_text(responses, "content"), + transcript=_wispr_text(responses, "transcript"), + summary=_string(first.get("summary")), + starts_at=_datetime(first.get("start")), + modified_at=_datetime(first.get("modified_at")), + share_link=_string(first.get("share_link")) or None, + breadcrumbs=[], + ), + ) + + +@asynccontextmanager +async def map_record( + record: SourceRecord, + source_name: str, + storage: StorageBackend, +) -> AsyncIterator[tuple[BaseEntity, ...]]: + """Keep verified local blob files alive only while strict projection consumes them.""" + if record.deleted_at is not None or record.content_access != "available": + raise ProjectionMappingError("Unavailable records cannot be projected") + with TemporaryDirectory(prefix="airweave-projection-") as temporary: + directory = Path(temporary) + if source_name == "gmail": + from airweave.domains.entities.canonical.gmail_projection import map_gmail + + entities = await map_gmail(record, storage, directory) + elif source_name == "google_drive": + entities = await _drive(record, storage, directory) + elif source_name == "google_calendar": + entities = _calendar(record) + elif source_name == "slack": + entities = _slack(record) + elif source_name == "wispr": + entities = _wispr(record) + else: + raise ProjectionMappingError("Source has no audited search projection mapper") + yield entities diff --git a/backend/airweave/domains/entities/canonical/projection_models.py b/backend/airweave/domains/entities/canonical/projection_models.py new file mode 100644 index 0000000000..d40352a021 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_models.py @@ -0,0 +1,61 @@ +"""Typed immutable search publication identities, independent of provider identity.""" + +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.domains.entities.canonical.models import SourceRecord + + +class ProjectionLocator(BaseModel): + """Stored in the existing original_entity_id field, never a native provider ID.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + record_id: UUID + revision: int = Field(ge=1) + pipeline_version: int = Field(ge=1) + generation: UUID + part_index: int = Field(ge=0) + + def encode(self) -> str: + """Encode a versioned, unambiguous publication locator.""" + return ( + f"canonical:v1:{self.record_id}:{self.revision}:{self.pipeline_version}:" + f"{self.generation}:{self.part_index}" + ) + + @classmethod + def parse(cls, value: str) -> "ProjectionLocator | None": + """None identifies legacy values; malformed canonical values fail closed.""" + if not value.startswith("canonical:"): + return None + parts = value.split(":") + if len(parts) != 7 or parts[:2] != ["canonical", "v1"]: + raise ValueError("Invalid canonical projection locator") + return cls( + record_id=parts[2], + revision=parts[3], + pipeline_version=parts[4], + generation=parts[5], + part_index=parts[6], + ) + + +class ProjectionWork(BaseModel): + """Snapshot plus publication pointer used for compare-and-swap.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + organization_id: UUID + record: SourceRecord + pipeline_version: int + previous_generation: UUID | None + + +class ProjectionBatchResult(BaseModel): + """One bounded scan page; failures remain pending for a later retry.""" + + after_id: UUID | None = None + has_more: bool = False + published: int = 0 + superseded: int = 0 + failed: int = 0 diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py new file mode 100644 index 0000000000..4b2342ac48 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -0,0 +1,139 @@ +"""Canonical pending work and atomic publication. No network I/O under row locks.""" + +from uuid import UUID + +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.db.unit_of_work import UnitOfWork +from airweave.domains.entities.canonical.projection_models import ProjectionLocator, ProjectionWork +from airweave.domains.entities.canonical.store import content_is_available, source_record +from airweave.models.entity import Entity +from airweave.models.sync import Sync + + +def publication_matches(locator: ProjectionLocator): + """Add request organization/source scope separately; gate before text leaves retrieval.""" + return and_( + Entity.id == locator.record_id, + Entity.record_revision == locator.revision, + Entity.indexed_revision == locator.revision, + Entity.indexed_pipeline_version == locator.pipeline_version, + Sync.index_pipeline_version == locator.pipeline_version, + Entity.indexed_generation == locator.generation, + Entity.deleted_at.is_(None), + content_is_available(), + ) + + +class CanonicalProjectionStore: + """Optimistic projection computation, serialized CAS against source capture.""" + + async def pending( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + *, + after_id: UUID | None = None, + limit: int = 25, + ) -> tuple[ProjectionWork, ...]: + """One bounded UUID page; deleted records publish an empty generation.""" + if not 1 <= limit <= 100: + raise ValueError("Projection batch size must be between 1 and 100") + statement = ( + select(Entity, Sync.index_pipeline_version) + .join(Sync, Sync.id == Entity.sync_id) + .where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Sync.organization_id == organization_id, + Entity.record_revision > 0, + 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), + ), + or_(Entity.deleted_at.is_not(None), content_is_available()), + ) + ) + if after_id is not None: + statement = statement.where(Entity.id > after_id) + rows = await db.execute(statement.order_by(Entity.id).limit(limit + 1)) + return tuple( + ProjectionWork( + organization_id=organization_id, + record=source_record(row), + pipeline_version=version, + previous_generation=row.indexed_generation, + ) + for row, version in rows + ) + + async def _current(self, db: AsyncSession, work: ProjectionWork) -> Entity | None: + # Capture locks Sync first too: parent visibility and revision cannot change + # between validation and publication. This transaction contains no provider I/O. + sync = await db.scalar( + select(Sync) + .where( + Sync.id == work.record.sync_id, + Sync.organization_id == work.organization_id, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + if sync is None or sync.index_pipeline_version != work.pipeline_version: + return None + entity = await db.scalar( + select(Entity) + .where( + Entity.id == work.record.id, + Entity.sync_id == work.record.sync_id, + Entity.organization_id == work.organization_id, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + if ( + entity is None + or entity.record_revision != work.record.revision + or entity.indexed_generation != work.previous_generation + ): + return None + if entity.deleted_at is None: + visible = await db.scalar(select(content_is_available()).where(Entity.id == entity.id)) + if not visible: + return None + return entity + + async def publish( + self, + db: AsyncSession, + work: ProjectionWork, + generation: UUID, + chunk_count: int, + ) -> bool: + """Publish only a complete current generation; late writers cannot replace it.""" + if chunk_count < 0 or (work.record.deleted_at is None and chunk_count == 0): + raise ValueError("Active records require a nonempty complete projection") + if work.record.deleted_at is not None and chunk_count != 0: + raise ValueError("Deleted records must publish no content") + async with UnitOfWork(db): + entity = await self._current(db, work) + if entity is None: + return False + entity.indexed_revision = work.record.revision + entity.indexed_pipeline_version = work.pipeline_version + entity.indexed_generation = generation + entity.indexed_chunk_count = chunk_count + entity.projection_error = None + await db.flush() + return True + + async def fail(self, db: AsyncSession, work: ProjectionWork, message: str) -> None: + """Record bounded diagnostics only while the same input remains current.""" + async with UnitOfWork(db): + entity = await self._current(db, work) + if entity is not None: + entity.projection_error = message[:1000] + await db.flush() diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py new file mode 100644 index 0000000000..ba17e95c10 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -0,0 +1,150 @@ +"""Bounded original-record projection using existing conversion/embed/feed primitives.""" + +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass +from uuid import UUID, uuid4 + +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.core.logging import ContextualLogger +from airweave.domains.entities.canonical.projection_models import ( + ProjectionBatchResult, + ProjectionLocator, + ProjectionWork, +) +from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore +from airweave.domains.storage.protocols import StorageBackend +from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor +from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields +from airweave.platform.destinations.vespa.destination import VespaDestination +from airweave.platform.entities._base import AirweaveSystemMetadata + + +@dataclass +class ProjectionContext: + """Projection processing has no source credentials or artificial capture job.""" + + logger: ContextualLogger + source_short_name: str + + +class StrictProjectionTracker: + """A replay publication cannot silently omit failed conversion inputs.""" + + async def record_skipped(self, count: int) -> None: + """Fail publication when an existing converter tries to skip required input.""" + if count: + raise ValueError("Required projection input could not be converted") + + +@dataclass +class ProjectionRuntime: + """Only processing diagnostics are required; source runtime is not recreated.""" + + entity_tracker: StrictProjectionTracker + + +class CanonicalProjector: + """Read snapshot, compute outside transaction, then atomically publish generation.""" + + def __init__( + self, + store: CanonicalProjectionStore, + sessions: Callable[[], AbstractAsyncContextManager[AsyncSession]], + processor: ChunkEmbedProcessor, + storage: StorageBackend, + ): + """Inject existing processor, storage, and transaction session ownership.""" + self._store = store + self._sessions = sessions + self._processor = processor + self._storage = storage + + async def project_one( + self, + work: ProjectionWork, + source_name: str, + destination: VespaDestination, + logger: ContextualLogger, + ) -> bool: + """Never feed native capture JSON; mapper emits explicit safe projection entities.""" + from airweave.domains.entities.canonical.projection_mappers import map_record + + generation = uuid4() + chunks = [] + if work.record.deleted_at is None: + async with map_record(work.record, source_name, self._storage) as mapped: + if not mapped: + raise ValueError("Projection mapper returned no required content") + for part_index, entity in enumerate(mapped): + populate_base_fields(entity) + if entity.airweave_system_metadata is None: + entity.airweave_system_metadata = AirweaveSystemMetadata() + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=part_index, + ) + entity.entity_id = locator.encode() + meta = entity.airweave_system_metadata + meta.source_name = source_name + meta.entity_type = type(entity).__name__ + meta.sync_id = work.record.sync_id + meta.sync_job_id = None + meta.db_entity_id = work.record.id + chunks = await self._processor.process( + list(mapped), + ProjectionContext(logger, source_name), + ProjectionRuntime(StrictProjectionTracker()), + strict=True, + ) + await destination.bulk_insert(chunks, strict=True) + async with self._sessions() as db: + return await self._store.publish(db, work, generation, len(chunks)) + + async def batch( + self, + organization_id: UUID, + sync_id: UUID, + source_name: str, + destination: VespaDestination, + logger: ContextualLogger, + *, + after_id: UUID | None = None, + limit: int = 25, + ) -> ProjectionBatchResult: + """Failed rows stay pending; other rows in the page continue to make progress.""" + async with self._sessions() as db: + pending = await self._store.pending( + db, + organization_id, + sync_id, + after_id=after_id, + limit=limit, + ) + published = superseded = failed = 0 + for work in pending[:limit]: + try: + if await self.project_one(work, source_name, destination, logger): + published += 1 + else: + superseded += 1 + except Exception as error: + failed += 1 + async with self._sessions() as db: + await self._store.fail(db, work, type(error).__name__) + logger.warning( + "Canonical projection failed for record %s (%s)", + work.record.id, + type(error).__name__, + ) + return ProjectionBatchResult( + after_id=pending[min(limit, len(pending)) - 1].record.id if pending else after_id, + has_more=len(pending) > limit, + published=published, + superseded=superseded, + failed=failed, + ) diff --git a/backend/airweave/domains/entities/canonical/query_store.py b/backend/airweave/domains/entities/canonical/query_store.py index c2436cfd6c..5931ccb8e5 100644 --- a/backend/airweave/domains/entities/canonical/query_store.py +++ b/backend/airweave/domains/entities/canonical/query_store.py @@ -2,13 +2,14 @@ from uuid import UUID -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import RecordFilters from airweave.domains.entities.canonical.store import ( SourceNotFound, + content_is_available, parent_is_visible, source_record, with_content_access, @@ -20,6 +21,26 @@ class CanonicalQueryStore: """No provider or index calls; reads only committed canonical rows.""" + async def captured_counts( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + ) -> dict[str, int]: + """Count currently visible captured records, independent of index publication.""" + rows = await db.execute( + select(Entity.entity_definition_short_name, func.count(Entity.id)) + .where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + ) + .group_by(Entity.entity_definition_short_name) + ) + return dict(rows.all()) + async def list_records( self, db: AsyncSession, diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 5333bd9594..89f7efc717 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -104,6 +104,7 @@ async def database(request): if getattr(request, "param", None) == "legacy": await connection.run_sync(seed_legacy) await connection.run_sync(migrate, "0001_canonical_records.py") + await connection.run_sync(migrate, "0002_projection_publication.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py new file mode 100644 index 0000000000..5110f672ab --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -0,0 +1,186 @@ +"""Replay source MIME content without provider requests or lossy snippet fallbacks.""" + +import base64 +import hashlib +from datetime import datetime, timezone +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.blob_materializer import read_blob +from airweave.domains.entities.canonical.gmail_projection import map_gmail +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity + + +def part(content=b"hello", mime="text/plain", **fields): + return { + "mimeType": mime, + "body": {"data": base64.urlsafe_b64encode(content).decode(), "size": len(content)}, + **fields, + } + + +def record(payload, **changes): + values = { + "id": uuid4(), + "sync_id": uuid4(), + "identity": RecordIdentity(record_type="message", native_id="m1"), + "revision": 1, + "payload": { + "id": "m1", + "threadId": "t1", + "internalDate": "1000", + "snippet": "must not substitute this snippet", + "payload": payload, + }, + "payload_schema_version": 1, + "capture_hash": "capture", + "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, + } + values.update(changes) + return SourceRecord(**values) + + +def blob(source, data, source_path="/payload/body"): + digest = hashlib.sha256(data).hexdigest() + return BlobReference( + key=f"canonical/{source.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(data), + source_path=source_path, + ) + + +@pytest.mark.asyncio +async def test_plain_charset_body_and_quoted_address_are_preserved(tmp_path): + source = record( + part( + "café ".encode("latin-1"), + headers=[ + {"name": "Content-Type", "value": "text/plain; charset=iso-8859-1"}, + {"name": "To", "value": '"Doe, Jane" , bob@example.com'}, + {"name": "Subject", "value": "Subject"}, + ], + ) + ) + before = source.model_dump() + entities = await map_gmail(source, AsyncMock(), tmp_path) + assert len(entities) == 1 + assert entities[0].to == ['"Doe, Jane" ', "bob@example.com"] + assert "café <hello>" in next(tmp_path.glob("*.html")).read_text() + 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) + text = next(tmp_path.glob("*.html")).read_text() + assert text == "

Full HTML body

" + assert entities[0].local_path is not None + + +@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": 11}}, + 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"),)}) + # Provider's declared size must match actual content. + 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) + 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_partial_capture_cannot_be_marked_successfully_projected(tmp_path): + with pytest.raises(ValueError, match="complete active"): + await map_gmail(record(part(), completeness="partial"), AsyncMock(), tmp_path) + + +@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) + 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" 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_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py new file mode 100644 index 0000000000..e830546730 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py @@ -0,0 +1,149 @@ +"""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 capture, observation +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 capture(database, service, fence, observation()) + work = (await pending(database, fence))[0] + store = CanonicalProjectionStore() + generation = uuid4() + async with database() as db: + assert await store.publish(db, work, generation, 3) + async with database() as db: + assert not await store.publish(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 store.publish(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 store.publish(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 + 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 store.publish(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 store.publish(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 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 store.publish(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 + + service, fence = source + 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): + 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] + + processor.process = process + destination = MagicMock() + destination.bulk_insert = AsyncMock() + storage = MagicMock() + projector = CanonicalProjector(CanonicalProjectionStore(), database, processor, storage) + assert await projector.project_one(work, "slack", destination, MagicMock()) + destination.bulk_insert.assert_awaited_once() + assert destination.bulk_insert.call_args.kwargs == {"strict": True} + assert not await pending(database, fence) 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..305a72365a --- /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 +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 store.publish(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_strict_projection.py b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py new file mode 100644 index 0000000000..0fb4d6b5b5 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py @@ -0,0 +1,47 @@ +"""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() 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..5a02d5c00f --- /dev/null +++ b/backend/airweave/domains/search/canonical_visibility.py @@ -0,0 +1,99 @@ +"""Reject stale index publications before any reranking or model consumption.""" + +from uuid import UUID + +from sqlalchemy import and_, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.projection_models import ProjectionLocator +from airweave.domains.entities.canonical.projection_store import publication_matches +from airweave.domains.search.types.results import SearchResult +from airweave.domains.sources.protocols import SourceRegistryProtocol +from airweave.models.entity import Entity +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( + getattr(registry.get(scope.short_name).source_class_ref, "canonical_record_types", ()) + ) + 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: set[tuple[UUID, UUID, int, int, UUID]] = set() + candidate_values = list(candidates.values()) + for start in range(0, len(candidate_values), 100): + predicates = [ + and_(Entity.sync_id == sync_id, publication_matches(locator)) + for sync_id, locator in 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, + ) + .join(Sync, Sync.id == Entity.sync_id) + .where( + Entity.organization_id == organization_id, + Sync.organization_id == organization_id, + or_(*predicates), + ) + ) + eligible.update(tuple(row) for row in rows) + return [ + result + for index, result in enumerate(results) + if index in legacy + or ( + index in candidates + and ( + candidates[index][0], + candidates[index][1].record_id, + candidates[index][1].revision, + candidates[index][1].pipeline_version, + candidates[index][1].generation, + ) + in eligible + ) + ] 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..4926362d65 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 ( @@ -130,12 +131,20 @@ async def execute( ) ) - vector_results = await vector_task + retrieved = await vector_task + 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, + 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,11 +156,15 @@ 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, 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, excluded_candidates=excluded, ) async def _execute_vector_search( @@ -279,7 +292,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/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..c7f72a8855 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,26 @@ 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] 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/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/results.py b/backend/airweave/domains/search/types/results.py index 367b94efd4..d950a08b70 100644 --- a/backend/airweave/domains/search/types/results.py +++ b/backend/airweave/domains/search/types/results.py @@ -239,6 +239,15 @@ class SearchResults(BaseModel): description="Search results ordered by relevance (highest first).", ) + 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/storage/file_service.py b/backend/airweave/domains/storage/file_service.py index 50a431a473..f3786540ad 100644 --- a/backend/airweave/domains/storage/file_service.py +++ b/backend/airweave/domains/storage/file_service.py @@ -28,12 +28,13 @@ 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, ) @@ -43,7 +44,7 @@ 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, 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/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/pipeline/text_builder.py b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py index 75a6d46901..5e01b89213 100644 --- a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py +++ b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py @@ -12,8 +12,10 @@ 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 +45,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. @@ -247,7 +249,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,7 +315,7 @@ 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", ) -> List[BaseEntity]: """Execute batch conversion for each converter group. @@ -340,7 +342,7 @@ async def _convert_sub_batch( self, converter: Any, sub_batch: List[Tuple[BaseEntity, str]], - sync_context: "SyncContext", + sync_context: "ProcessingContext", ) -> List[BaseEntity]: """Convert a sub-batch of entities using the given converter. @@ -403,8 +405,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/processors/chunk_embed.py b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py index 237d929e24..1d52e46bd6 100644 --- a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py +++ b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py @@ -23,8 +23,10 @@ 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: @@ -44,13 +46,19 @@ def __init__( 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) @@ -58,6 +66,8 @@ async def process( 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,12 +75,22 @@ 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" @@ -85,8 +105,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 +127,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 +154,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 +200,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 +233,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 +299,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 +345,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/temporal/activities/project_canonical_records.py b/backend/airweave/domains/temporal/activities/project_canonical_records.py new file mode 100644 index 0000000000..9d5756f6eb --- /dev/null +++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py @@ -0,0 +1,81 @@ +"""Project owned records without reconnecting or consulting the provider.""" + +from dataclasses import dataclass +from uuid import UUID + +from sqlalchemy import select +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 CanonicalSource +from airweave.domains.sources.protocols import SourceRegistryProtocol +from airweave.models.collection import Collection +from airweave.models.source_connection import SourceConnection +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) -> 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: + row = ( + await db.execute( + select(SourceConnection.short_name, Collection.id) + .join( + Collection, + Collection.readable_id == SourceConnection.readable_collection_id, + ) + .where( + SourceConnection.organization_id == organization, + SourceConnection.sync_id == sync, + Collection.organization_id == organization, + ) + ) + ).one_or_none() + if row is None: + return ProjectionBatchResult().model_dump(mode="json") + source_name, collection_id = row + entry = self.source_registry.get(source_name) + if not isinstance(entry.source_class_ref, CanonicalSource): + 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, + ) + 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: + result = await self.projector.batch( + organization, + sync, + source_name, + destination, + logger, + after_id=UUID(after_id) if after_id else None, + ) + return result.model_dump(mode="json") + finally: + await destination.close_connection() diff --git a/backend/airweave/domains/temporal/worker/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py index 58defe7211..7bd3b1e4ae 100644 --- a/backend/airweave/domains/temporal/worker/wiring.py +++ b/backend/airweave/domains/temporal/worker/wiring.py @@ -33,6 +33,14 @@ 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.project_canonical_records import ( + ProjectCanonicalRecordsActivity, + ) + email_service = container.email_service event_bus = container.event_bus sync_service = container.sync_service @@ -50,6 +58,19 @@ def create_activities() -> list: logger.debug("Wiring activities with container dependencies") return [ + 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 +122,12 @@ def get_workflows() -> list: CleanupSyncDataWorkflow, RunSourceConnectionWorkflow, ) + from airweave.domains.temporal.workflows.project_canonical_records import ( + ProjectCanonicalRecordsWorkflow, + ) return [ + ProjectCanonicalRecordsWorkflow, RunSourceConnectionWorkflow, CleanupStuckSyncJobsWorkflow, CleanupSyncDataWorkflow, 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..db6a689f7d --- /dev/null +++ b/backend/airweave/domains/temporal/workflows/project_canonical_records.py @@ -0,0 +1,49 @@ +"""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, + ) -> 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], + 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 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. + workflow.continue_as_new(args=[organization_id, sync_id, after_id, failures, sweep]) diff --git a/backend/airweave/domains/temporal/workflows/run_source_connection.py b/backend/airweave/domains/temporal/workflows/run_source_connection.py index f2966a83bc..3b10645b0f 100644 --- a/backend/airweave/domains/temporal/workflows/run_source_connection.py +++ b/backend/airweave/domains/temporal/workflows/run_source_connection.py @@ -31,6 +31,9 @@ CLASSIFIED_USER_ERROR_TYPE, ORPHANED_SYNC_ERROR_TYPE, ) + from airweave.domains.temporal.workflows.project_canonical_records import ( + ProjectCanonicalRecordsWorkflow, + ) # --------------------------------------------------------------------------- # Execution policies @@ -189,22 +192,34 @@ 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. + 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, + ) + ) # ------------------------------------------------------------------ # 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/models/entity.py b/backend/airweave/models/entity.py index f6c6bf0359..b22626bf23 100644 --- a/backend/airweave/models/entity.py +++ b/backend/airweave/models/entity.py @@ -60,6 +60,8 @@ class Entity(OrganizationBase): 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) diff --git a/backend/airweave/platform/destinations/vespa/destination.py b/backend/airweave/platform/destinations/vespa/destination.py index b92cd14d79..5c40fbfd4e 100644 --- a/backend/airweave/platform/destinations/vespa/destination.py +++ b/backend/airweave/platform/destinations/vespa/destination.py @@ -140,7 +140,7 @@ 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: + async def bulk_insert(self, entities: List[BaseEntity], *, strict: bool = False) -> None: """Transform entities and batch feed to Vespa. Args: @@ -162,6 +162,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/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index 7496964bb6..bad98fcf8b 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -147,6 +147,7 @@ GoogleDriveDriveEntity, GoogleDriveFileDeletionEntity, GoogleDriveFileEntity, + GoogleDriveFolderEntity, ) from .google_slides import ( GoogleSlidesPresentationEntity, @@ -287,7 +288,7 @@ SlabPostEntity, SlabTopicEntity, ) -from .slack import SlackMessageEntity +from .slack import SlackChannelEntity, SlackMessageEntity from .slite import SliteNoteEntity from .stripe import ( StripeBalanceEntity, @@ -335,6 +336,7 @@ TrelloMemberEntity, ) from .web import WebFileEntity +from .wispr import WisprMeetingEntity from .word import WordDocumentEntity from .zendesk import ( ZendeskAttachmentEntity, @@ -518,6 +520,7 @@ GoogleDocsDocumentEntity, ], "google_drive": [ + GoogleDriveFolderEntity, GoogleDriveDriveEntity, GoogleDriveFileDeletionEntity, GoogleDriveFileEntity, @@ -651,6 +654,7 @@ SlabTopicEntity, ], "slack": [ + SlackChannelEntity, SlackMessageEntity, ], "stripe": [ @@ -704,6 +708,7 @@ "web": [ WebFileEntity, ], + "wispr": [WisprMeetingEntity], "word": [ WordDocumentEntity, ], diff --git a/backend/airweave/platform/entities/google_drive.py b/backend/airweave/platform/entities/google_drive.py index 801f75c736..ac6e7e5468 100644 --- a/backend/airweave/platform/entities/google_drive.py +++ b/backend/airweave/platform/entities/google_drive.py @@ -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/slack.py b/backend/airweave/platform/entities/slack.py index e79b9c2a65..84b59c007b 100644 --- a/backend/airweave/platform/entities/slack.py +++ b/backend/airweave/platform/entities/slack.py @@ -132,3 +132,12 @@ 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) diff --git a/backend/airweave/platform/entities/wispr.py b/backend/airweave/platform/entities/wispr.py new file mode 100644 index 0000000000..cbb03236d6 --- /dev/null +++ b/backend/airweave/platform/entities/wispr.py @@ -0,0 +1,32 @@ +"""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 "" diff --git a/backend/airweave/platform/sources/retry_helpers.py b/backend/airweave/platform/http_client/retry_helpers.py similarity index 100% rename from backend/airweave/platform/sources/retry_helpers.py rename to backend/airweave/platform/http_client/retry_helpers.py 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..1d915a09d3 100644 --- a/backend/airweave/platform/sources/attio.py +++ b/backend/airweave/platform/sources/attio.py @@ -36,12 +36,12 @@ AttioRecordEntity, ) 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 _API = "https://api.attio.com/v2" 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..1a8c4d0c6c 100644 --- a/backend/airweave/platform/sources/github.py +++ b/backend/airweave/platform/sources/github.py @@ -34,12 +34,12 @@ GitHubRepositoryEntity, ) 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/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 a739f224c0..aec5ec7e48 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -42,12 +42,12 @@ GmailThreadEntity, ) from airweave.platform.http_client.airweave_client import AirweaveHttpClient +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.http_helpers import raise_for_status -from airweave.platform.sources.retry_helpers import ( - wait_rate_limit_with_backoff, -) from airweave.platform.utils.filename_utils import safe_filename from airweave.schemas.source_connection import AuthenticationMethod, OAuthType diff --git a/backend/airweave/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index 0619d8eb46..ac13bb9709 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -60,13 +60,13 @@ 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.records.google_calendar import generate_calendar_observations -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_calendar import generate_calendar_observations from airweave.schemas.source_connection import AuthenticationMethod, OAuthType 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 3b673c647a..54d3b6bd76 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -46,14 +46,14 @@ _parse_drive_dt, ) 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, +) from airweave.platform.sources._base import BaseSource from airweave.platform.sources.http_helpers import raise_for_status from airweave.platform.sources.records.google_drive import generate_drive_observations from airweave.platform.sources.records.google_drive_content import capture_file_content -from airweave.platform.sources.retry_helpers import ( - retry_if_rate_limit_or_timeout, - wait_rate_limit_with_backoff, -) from airweave.schemas.source_connection import AuthenticationMethod, OAuthType 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/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..35d40bf5a9 100644 --- a/backend/airweave/platform/sources/linear.py +++ b/backend/airweave/platform/sources/linear.py @@ -38,12 +38,12 @@ LinearUserEntity, ) 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" 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..9643bff6d6 100644 --- a/backend/airweave/platform/sources/notion.py +++ b/backend/airweave/platform/sources/notion.py @@ -32,12 +32,12 @@ _extract_rich_text_plain, ) 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/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..9d656b6abd 100644 --- a/backend/airweave/platform/sources/outlook_calendar.py +++ b/backend/airweave/platform/sources/outlook_calendar.py @@ -30,12 +30,12 @@ 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.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_mail.py b/backend/airweave/platform/sources/outlook_mail.py index cdc3ce6ccc..cf92280568 100644 --- a/backend/airweave/platform/sources/outlook_mail.py +++ b/backend/airweave/platform/sources/outlook_mail.py @@ -34,12 +34,12 @@ 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.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.filename_utils import safe_filename from airweave.schemas.source_connection import AuthenticationMethod, OAuthType 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/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 5ade2827f3..0796304bba 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -33,12 +33,12 @@ 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 raise_for_status from airweave.schemas.source_connection import AuthenticationMethod, OAuthType # Pattern for Slack mrkdwn special sequences: 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..ffa92f798e 100644 --- a/backend/airweave/platform/sources/stripe.py +++ b/backend/airweave/platform/sources/stripe.py @@ -49,12 +49,12 @@ _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.schemas.source_connection import AuthenticationMethod _parse_unix_timestamp = _parse_stripe_ts 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/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/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/search_v2.py b/backend/airweave/schemas/search_v2.py index 9eaa2d9537..3eae466de4 100644 --- a/backend/airweave/schemas/search_v2.py +++ b/backend/airweave/schemas/search_v2.py @@ -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,8 @@ class SearchV2Response(BaseModel): results: list[SearchResult] = Field( default_factory=list, description="Search results ordered by relevance." ) + + 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/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/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py new file mode 100644 index 0000000000..47a4fe7cd2 --- /dev/null +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -0,0 +1,175 @@ +"""Search projections preserve canonical data and cannot fetch missing source content.""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import ( + ProjectionMappingError, + map_record, +) +from airweave.domains.entities.canonical.requests import RecordIdentity + + +def record(record_type, payload, *, native_id="id", container_id=None): + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity( + record_type=record_type, native_id=native_id, container_id=container_id + ), + revision=1, + payload=payload, + payload_schema_version=1, + 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_calendar_keeps_master_recurrence_and_native_container(): + payload = { + "id": "event", + "summary": "Weekly", + "start": {"date": "2026-10-01"}, + "end": {"date": "2026-10-02"}, + "recurrence": ["RRULE:FREQ=WEEKLY"], + "unknown_native": "keep", + } + item = record("event", payload, native_id="event", container_id="calendar") + async with map_record(item, "google_calendar", AsyncMock()) as entities: + assert entities[0].calendar_key == "calendar" + assert entities[0].recurrence == ["RRULE:FREQ=WEEKLY"] + assert entities[0].start_date == "2026-10-01" + assert item.payload == payload + + +@pytest.mark.asyncio +async def test_slack_history_shape_uses_container_without_provider_lookup(): + item = record( + "message", + {"ts": "123.456", "text": "hello", "edited": {"ts": "124.0"}}, + native_id="123.456", + container_id="C1", + ) + storage = AsyncMock() + async with map_record(item, "slack", storage) as entities: + assert entities[0].channel_id == "C1" + assert entities[0].text == "hello" + assert entities[0].updated_at.timestamp() == 124 + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_drive_missing_blob_fails_without_remote_fallback(): + item = record("file", {"id": "file", "name": "file.pdf", "mimeType": "application/pdf"}) + with pytest.raises(ProjectionMappingError, match="blob"): + async with map_record(item, "google_drive", AsyncMock()): + pass + + +@pytest.mark.asyncio +async def test_wispr_ranges_do_not_duplicate_default_content(): + first = { + "id": "meeting", + "title": "Planning", + "summary": "summary", + "content": "notes", + "transcript": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)\n\nProvider guidance." + ), + } + second = {"id": "meeting", "content": "notes", "transcript": "def"} + payload = { + "responses": [ + { + "requested_ranges": { + "view_content": {"start_char": 0}, + "view_transcript": {"start_char": 0}, + }, + "response": first, + }, + {"requested_ranges": {"view_transcript": {"start_char": 3}}, "response": second}, + ] + } + async with map_record(record("meeting", payload), "wispr", AsyncMock()) as entities: + assert entities[0].notes == "notes" + assert entities[0].transcript == "abc\ndef" + assert "guidance" not in entities[0].transcript + + +@pytest.mark.asyncio +async def test_wispr_unfinished_ranges_fail_instead_of_indexing_truncated_body(): + payload = { + "responses": [ + { + "requested_ranges": { + "view_content": {"start_char": 0}, + "view_transcript": {"start_char": 0}, + }, + "response": { + "content": "", + "transcript": ( + "(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)" + ), + }, + } + ] + } + with pytest.raises(ProjectionMappingError, match="unfinished"): + async with map_record(record("meeting", payload), "wispr", AsyncMock()): + pass + + +@pytest.mark.asyncio +async def test_drive_uses_verified_owned_bytes_and_cleans_materialization(): + import hashlib + from pathlib import Path + + from airweave.domains.entities.canonical.requests import BlobReference + + content = b"owned document body" + item = record("file", {"id": "file", "name": "../../unsafe.txt", "mimeType": "text/plain"}) + digest = hashlib.sha256(content).hexdigest() + blob = BlobReference( + key=f"canonical/{item.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type="text/plain", + ) + item = item.model_copy(update={"blobs": (blob,)}) + storage = AsyncMock() + storage.read_file.return_value = content + async with map_record(item, "google_drive", storage) as entities: + path = Path(entities[0].local_path) + assert path.read_bytes() == content + assert path.name == digest + ".txt" + assert entities[0].mime_type == "text/plain" + assert not path.exists() + storage.read_file.assert_awaited_once_with(blob.key) + + +def test_added_typed_entities_are_registered(): + from airweave.domains.entities.registry import EntityDefinitionRegistry + from airweave.platform.entities.google_drive import GoogleDriveFolderEntity + from airweave.platform.entities.slack import SlackChannelEntity + from airweave.platform.entities.wispr import WisprMeetingEntity + + registry = EntityDefinitionRegistry() + registry.build() + for entity in (WisprMeetingEntity, GoogleDriveFolderEntity, SlackChannelEntity): + assert registry.get_short_name_by_class(entity) diff --git a/backend/tests/unit/domains/storage/test_import_order.py b/backend/tests/unit/domains/storage/test_import_order.py new file mode 100644 index 0000000000..86e1e80caf --- /dev/null +++ b/backend/tests/unit/domains/storage/test_import_order.py @@ -0,0 +1,27 @@ +"""Fresh-interpreter storage imports must not depend on eager source catalogue order.""" + +import subprocess +import sys +from pathlib import Path + + +def test_file_service_imports_before_source_catalogue(): + root = Path(__file__).resolve().parents[4] + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import conftest; " + "from airweave.domains.storage.file_service import FileService; " + "from airweave.platform.sources import GmailSource, GoogleDriveSource; " + "assert FileService and GmailSource and GoogleDriveSource" + ), + ], + cwd=root, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stderr From 302aad8281221b57a95140cfbf7250f08e4713a3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 02:42:49 -0700 Subject: [PATCH 005/266] Verify durable live capture and retain oversized Drive export metadata --- .../sources/records/google_drive_content.py | 28 ++ backend/tests/live/README.md | 44 +++ backend/tests/live/canonical_capture.py | 298 ++++++++++++++++++ .../entities/test_projection_mappers.py | 41 +++ .../records/test_google_drive_content.py | 31 +- 5 files changed, 440 insertions(+), 2 deletions(-) create mode 100644 backend/tests/live/README.md create mode 100644 backend/tests/live/canonical_capture.py diff --git a/backend/airweave/platform/sources/records/google_drive_content.py b/backend/airweave/platform/sources/records/google_drive_content.py index b102949e60..338756488b 100644 --- a/backend/airweave/platform/sources/records/google_drive_content.py +++ b/backend/airweave/platform/sources/records/google_drive_content.py @@ -2,6 +2,8 @@ from urllib.parse import quote, urlencode +import httpx + from airweave.core.logging import ContextualLogger from airweave.domains.entities.canonical.requests import CaptureRecord from airweave.domains.sources.token_providers.protocol import SourceAuthProvider @@ -46,6 +48,11 @@ async def capture_file_content( blob = await files.capture_canonical_url( url=download_url, client=client, auth=auth, logger=logger, media_type=media_type ) + except httpx.HTTPStatusError as error: + if not _export_limit_reached(error, mime): + raise + logger.info("Drive export exceeds provider size limit; retained metadata only") + return record except FileSkippedException: # Declared/streamed size limits are a visible metadata-only record, not absent data. return record @@ -57,3 +64,24 @@ async def capture_file_content( return record.model_copy( update={"blobs": (blob,), "content_hash": blob.sha256, "completeness": "complete"} ) + + +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/tests/live/README.md b/backend/tests/live/README.md new file mode 100644 index 0000000000..38c09e96c9 --- /dev/null +++ b/backend/tests/live/README.md @@ -0,0 +1,44 @@ +# Durable live capture proof + +`canonical_capture.py` is opt-in and makes read-only Gmail/Drive 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. 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 sample without a blob fails. 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. diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py new file mode 100644 index 0000000000..e0859b5a77 --- /dev/null +++ b/backend/tests/live/canonical_capture.py @@ -0,0 +1,298 @@ +"""Opt-in durable provider proof against an isolated, private local test PostgreSQL. + +Run from backend with PYTHONPATH=. and explicit environment variables documented in +README.md. No provider writes; this script writes real private records only into a +unique disposable schema and a mode0700 temporary blob directory. Both are removed. +""" + +import asyncio +import hashlib +import json +import logging +import os +import secrets +from contextlib import aclosing +from pathlib import Path +from tempfile import TemporaryDirectory +from uuid import uuid4 + +# The application DSN is deliberately made unusable. Only the explicit test DSN below +# is ever handed to an engine. Provider credentials remain subprocess environment data. +os.environ.update( + ENVIRONMENT="local", + TESTING="true", + AUTH_ENABLED="false", + AUTH_MODE="local", + POSTGRES_HOST="unused-live-probe.invalid", + POSTGRES_USER="unused", + POSTGRES_PASSWORD="unused", +) +import httpx # noqa: E402 +from sqlalchemy import text # noqa: E402 +from sqlalchemy.engine import make_url # noqa: E402 +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine # noqa: E402 + +import conftest # noqa: E402,F401 - local test settings must precede Airweave imports +from airweave.adapters.storage.filesystem import FilesystemBackend # noqa: E402 +from airweave.domains.entities.canonical.query import CanonicalQueryService # noqa: E402 +from airweave.domains.entities.canonical.query_models import RecordListQuery # noqa: E402 +from airweave.domains.entities.canonical.query_store import CanonicalQueryStore # noqa: E402 +from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord # noqa: E402 +from airweave.domains.entities.canonical.service import CanonicalCaptureService # noqa: E402 +from airweave.domains.entities.canonical.store import CanonicalRecordStore # noqa: E402 +from airweave.domains.entities.canonical.tests.conftest import migrate # noqa: E402 +from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider # noqa: E402 +from airweave.domains.storage.file_service import FileService # noqa: E402 +from airweave.domains.storage.paths import StoragePaths # noqa: E402 +from airweave.models import Organization, Sync, SyncJob # noqa: E402 +from airweave.platform.configs.config import GmailConfig, GoogleDriveConfig # noqa: E402 +from airweave.platform.http_client.airweave_client import AirweaveHttpClient # noqa: E402 +from airweave.platform.http_client.composio_transport import ComposioTransport # noqa: E402 +from airweave.platform.sources import GmailSource, GoogleDriveSource # noqa: E402 + +logging.disable(logging.CRITICAL) +PROBE_STAGE = "initialization" + + +def test_database_url() -> str: + """Reject network/application databases even if a caller supplies the wrong URL.""" + value = os.environ["CANONICAL_TEST_DATABASE_URL"] + url = make_url(value) + socket = url.query.get("host", "") + if ( + url.drivername != "postgresql+asyncpg" + or url.database != "sync_tests" + or url.username != "sync_test" + or url.host + or not socket.startswith("/") + ): + raise ValueError("An explicit private Unix-socket sync_tests database is required") + directory = Path(socket).resolve() + if ( + not directory.is_dir() + or directory.stat().st_uid != os.getuid() + or directory.stat().st_mode & 0o077 + ): + raise ValueError("The database socket directory must belong to the current user") + return value + + +async def fetch_records(name, account, expected_email, key, fence, root): + """Exercise actual source capture, bounded before any complete-scope claim.""" + host = "gmail.googleapis.com" if name == "gmail" else "www.googleapis.com" + auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) + transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) + transport.MAX_BINARY_BYTES = 10 * 1024 * 1024 + storage = FilesystemBackend(root / "blobs") + files = FileService(fence.job_id, storage, sync_id=fence.sync_id) + files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 + async with httpx.AsyncClient(transport=transport, timeout=180) as client: + wrapped = AirweaveHttpClient( + client, fence.organization_id, name, feature_flag_enabled=False + ) + if name == "gmail": + source = await GmailSource.create( + auth=auth, + logger=logging.getLogger("probe"), + http_client=wrapped, + config=GmailConfig(gmail_query="newer_than:7d smaller:5M"), + ) + profile = await source._get("https://gmail.googleapis.com/gmail/v1/users/me/profile") + email = profile.get("emailAddress", "") + else: + source = await GoogleDriveSource.create( + auth=auth, + logger=logging.getLogger("probe"), + http_client=wrapped, + config=GoogleDriveConfig(), + ) + profile = await source._get( + "https://www.googleapis.com/drive/v3/about", params={"fields": "user(emailAddress)"} + ) + email = profile.get("user", {}).get("emailAddress", "") + if email.casefold() != expected_email.casefold(): + raise ValueError("The connected account did not match the explicitly selected identity") + records = [] + async with ( + asyncio.timeout(180), + aclosing(source.generate_observations(files=files)) as stream, + ): + async for item in stream: + # Deliberately ignore Started/CompletedScope and cursor checkpoints: this + # selected sample does not prove exhaustive source enumeration. + if isinstance(item, CaptureRecord): + records.append(item) + if (len(records) >= 3 and any(x.blobs for x in records)) or len(records) >= 25: + break + if not records or not any(x.blobs for x in records): + raise ValueError("Bounded source sample did not contain a real owned blob") + return records, storage + + +async def verify(name, account, email, key, sessions, engine, root): + """Commit, disconnect, reopen, read/list/page, verify disk bytes, then replay.""" + organization_id, sync_id, job_id = uuid4(), uuid4(), uuid4() + async with sessions() as db: + db.add(Organization(id=organization_id, name="Private live capture proof")) + await db.flush() + db.add(Sync(id=sync_id, organization_id=organization_id, name="Bounded live " + name)) + await db.flush() + db.add( + SyncJob(id=job_id, organization_id=organization_id, sync_id=sync_id, status="running") + ) + await db.commit() + store = CanonicalRecordStore() + service = CanonicalCaptureService(store) + query = CanonicalQueryService(store, CanonicalQueryStore(), secrets.token_urlsafe(32)) + async with sessions() as db: + fence = await service.activate_writer( + db, organization_id, sync_id, job_id, attempt_id=uuid4(), attempt_number=1 + ) + global PROBE_STAGE + PROBE_STAGE = name + ":provider_capture" + records, _ = await fetch_records(name, account, email, key, fence, root) + PROBE_STAGE = name + ":durable_verification" + async with sessions() as db: + committed = await service.capture(db, CaptureBatch(fence=fence, records=tuple(records))) + await engine.dispose() # Force readback over fresh PostgreSQL connections. + storage = FilesystemBackend(root / "blobs") # New filesystem storage instance. + expected = {item.identity.entity_key: item for item in records} + listed = [] + cursor = None + while True: + async with sessions() as db: + page = await query.list_records( + db, organization_id, sync_id, RecordListQuery(limit=2, cursor=cursor) + ) + listed.extend(page.records) + if not page.has_more: + break + cursor = page.next_cursor + assert len(listed) == len(expected) + total_bytes = blob_count = 0 + for stored in listed: + assert stored.payload == expected[stored.identity.entity_key].payload + async with sessions() as db: + reread = await query.read(db, organization_id, sync_id, stored.id) + assert reread.payload == stored.payload and reread.revision == 1 + for blob in reread.blobs: + content = await storage.read_file(blob.key) + assert len(content) == blob.size_bytes + assert hashlib.sha256(content).hexdigest() == blob.sha256 + total_bytes += len(content) + blob_count += 1 + changes = [] + change_cursor = None + while True: + async with sessions() as db: + page = await query.changes(db, organization_id, sync_id, cursor=change_cursor, limit=2) + changes.extend(page.changes) + change_cursor = page.next_cursor + if not page.has_more: + break + assert len(changes) == len(expected) == len(committed.changes) + async with sessions() as db: + replay = await service.capture(db, CaptureBatch(fence=fence, records=tuple(records))) + assert replay.unchanged == len(records) and not replay.changes + async with sessions() as db: + tail = await query.changes(db, organization_id, sync_id, cursor=change_cursor) + assert not tail.changes + return { + "provider": name, + "records": len(listed), + "journal_changes": len(changes), + "blobs": blob_count, + "blob_bytes": total_bytes, + "blob_sha_verified": True, + "replay_new_changes": 0, + "fresh_connections_readback": True, + "full_scope_completed": False, + "checkpoint_saved": False, + } + + +async def main(): + """Allocate and destroy only private disposable local resources.""" + url = test_database_url() + key = os.environ["COMPOSIO_API_KEY"] + selected = os.environ.get("LIVE_PROVIDERS", "gmail,google_drive").split(",") + if not selected or set(selected) - {"gmail", "google_drive"}: + raise ValueError("LIVE_PROVIDERS must select gmail and/or google_drive") + inputs = [ + (name, os.environ[variable], os.environ["LIVE_EXPECTED_EMAIL"]) + for name, variable in ( + ("gmail", "LIVE_GMAIL_ACCOUNT_ID"), + ("google_drive", "LIVE_DRIVE_ACCOUNT_ID"), + ) + if name in selected + ] + schema = "canonical_live_" + uuid4().hex + admin = create_async_engine(url) + engine = create_async_engine(url, connect_args={"server_settings": {"search_path": schema}}) + old_path = StoragePaths.TEMP_PROCESSING + old_umask = os.umask(0o077) + created = False + results = [] + try: + async with admin.begin() as connection: + await connection.execute(text(f'CREATE SCHEMA "{schema}"')) + created = True + async with engine.begin() as connection: + for migration in ( + "0000_baseline.py", + "0001_canonical_records.py", + "0002_projection_publication.py", + ): + await connection.run_sync(migrate, migration) + with TemporaryDirectory(prefix="airweave-private-live-") as directory: + root = Path(directory) + os.chmod(root, 0o700) + StoragePaths.TEMP_PROCESSING = str(root / "downloads") + sessions = async_sessionmaker(engine, expire_on_commit=False) + for name, account, email in inputs: + result = await verify(name, account, email, key, sessions, engine, root) + results.append(result) + print(json.dumps({"verified_provider": result}), flush=True) + finally: + StoragePaths.TEMP_PROCESSING = old_path + os.umask(old_umask) + await engine.dispose() + if created: + async with admin.begin() as connection: + await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + await admin.dispose() + print( + json.dumps( + {"verification": results, "schema_removed": True, "blob_directory_removed": True} + ) + ) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except Exception as error: + # Never print SQL parameters, native payloads, credentials, or exception tracebacks. + details = {"failed": True, "error_type": type(error).__name__, "stage": PROBE_STAGE} + if isinstance(error, httpx.HTTPStatusError): + details["http_status"] = error.response.status_code + details["host"] = error.request.url.host + try: + native = error.response.json().get("error", {}) + reasons = [item.get("reason") for item in native.get("errors", [])] + allowed = { + "exportSizeLimitExceeded", + "cannotDownloadFile", + "insufficientFilePermissions", + "insufficientPermissions", + "fileNotDownloadable", + "downloadRestrictedForRevision", + "forbidden", + "rateLimitExceeded", + "userRateLimitExceeded", + } + details["recognized_provider_reasons"] = [r for r in reasons if r in allowed] + except (ValueError, AttributeError, TypeError): + pass + print(json.dumps(details)) + raise SystemExit(1) from None diff --git a/backend/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 47a4fe7cd2..508849d4ec 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -173,3 +173,44 @@ def test_added_typed_entities_are_registered(): registry.build() for entity in (WisprMeetingEntity, GoogleDriveFolderEntity, SlackChannelEntity): assert registry.get_short_name_by_class(entity) + + +async def test_owned_drive_spreadsheet_uses_real_converter(): + import hashlib + from io import BytesIO + from pathlib import Path + + import openpyxl + + from airweave.domains.converters.registry import ConverterRegistry + from airweave.domains.entities.canonical.requests import BlobReference + + workbook = openpyxl.Workbook() + workbook.active.append(["Project", "Status"]) + workbook.active.append(["Owned capture", "Verified"]) + buffer = BytesIO() + workbook.save(buffer) + content = buffer.getvalue() + digest = hashlib.sha256(content).hexdigest() + mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + item = record( + "file", + {"id": "sheet", "name": "Planning", "mimeType": "application/vnd.google-apps.spreadsheet"}, + ) + blob = BlobReference( + key=f"canonical/{item.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type=mime, + ) + item = item.model_copy(update={"blobs": (blob,)}) + storage = AsyncMock() + storage.read_file.return_value = content + async with map_record(item, "google_drive", storage) as entities: + path = entities[0].local_path + assert Path(path).suffix == ".xlsx" + converter = ConverterRegistry().for_extension(".xlsx") + assert converter is not None + result = await converter.convert_batch([path]) + assert "Owned capture" in result[path] and "Verified" in result[path] + assert not Path(path).exists() diff --git a/backend/tests/unit/platform/sources/records/test_google_drive_content.py b/backend/tests/unit/platform/sources/records/test_google_drive_content.py index 8d22752dd0..a03eebbb9f 100644 --- a/backend/tests/unit/platform/sources/records/test_google_drive_content.py +++ b/backend/tests/unit/platform/sources/records/test_google_drive_content.py @@ -1,8 +1,9 @@ """Body completeness requires durable bytes matching the observed file version.""" from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock +import httpx import pytest from airweave.domains.entities.canonical.requests import BlobReference @@ -24,7 +25,7 @@ def inputs(): "get": AsyncMock(return_value={"version": "7"}), "client": object(), "auth": object(), - "logger": object(), + "logger": Mock(), } @@ -72,3 +73,29 @@ async def test_native_document_uses_existing_export_format(inputs): call["media_type"] == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" ) + + +@pytest.mark.parametrize( + "reason", ["exportSizeLimitExceeded", "insufficientPermissions", "unknown"] +) +async def test_provider_export_limit_is_narrow_and_preserves_native_payload(inputs, reason): + response = httpx.Response( + 403, + json={"error": {"errors": [{"reason": reason}]}}, + request=httpx.Request("GET", "https://www.googleapis.com/drive/v3/files/test/export"), + ) + failure = httpx.HTTPStatusError( + "provider rejection", request=response.request, response=response + ) + inputs["files"].capture_canonical_url.side_effect = failure + record = file_record( + {"id": "doc", "mimeType": "application/vnd.google-apps.document", "version": "7"} + ) + if reason != "exportSizeLimitExceeded": + with pytest.raises(httpx.HTTPStatusError): + await capture_file_content(record, **inputs) + return + result = await capture_file_content(record, **inputs) + assert result is record and result.completeness == "metadata_only" and not result.blobs + assert result.payload == record.payload + inputs["get"].assert_not_awaited() From ae07e72367673bac95161c4d1559487647dfeba4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 02:54:25 -0700 Subject: [PATCH 006/266] Prepare private Porter application deployment and direct process startup --- Makefile | 8 +++- backend/entrypoint.sh | 60 +++------------------------ deploy/README.md | 95 +++++++++++++++++++++++++++++++++++++++++++ porter.yaml | 54 ++++++++++++++++++++++++ 4 files changed, 161 insertions(+), 56 deletions(-) create mode 100644 deploy/README.md create mode 100644 porter.yaml diff --git a/Makefile b/Makefile index b4458ebff5..9677c725ec 100644 --- a/Makefile +++ b/Makefile @@ -3,16 +3,20 @@ BACKEND := backend PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest -.PHONY: help setup test-store test-capture test-search +.PHONY: help setup test-store test-capture test-search validate-deploy help: @echo 'setup Install the locked backend and development dependencies' @echo 'test-store Verify record transactions and publication on disposable PostgreSQL' @echo 'test-capture Verify source capture, transport and blob handling' @echo 'test-search Verify search services and visibility behavior' + @echo 'validate-deploy Validate the Porter application manifest locally (no deployment)' @echo '' @echo 'test-store requires CANONICAL_TEST_DATABASE_URL; it creates and removes test schemas.' - @echo 'Deployment is not configured yet. See backend/SERVICE_AUTH.md for service setup.' + @echo 'See deploy/README.md for staging prerequisites and explicit deployment steps.' + +validate-deploy: + porter apply validate -f porter.yaml setup: cd $(BACKEND) && POETRY_VIRTUALENVS_IN_PROJECT=true uvx --from poetry==2.3.2 poetry install --with dev,lint --no-root --no-interaction diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 1226cf135b..5284e3fca5 100644 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -1,55 +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) -" - -echo "Starting against the provisioned database..." -cd /app -# Migrations and bootstrap are explicit deployment jobs, never replica startup. - -# Start application -echo "Starting application..." -poetry run uvicorn airweave.main:app --host 0.0.0.0 --port 8001 +# 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/deploy/README.md b/deploy/README.md new file mode 100644 index 0000000000..70a5da55e4 --- /dev/null +++ b/deploy/README.md @@ -0,0 +1,95 @@ +# Owned source store staging + +Status: application manifest prepared; dependencies and a running cluster are +still required. Local manifest validation is not a deployment or capacity test. + +The target is Porter project 19566, cluster 6108, AWS us-east-1. No production +cutover. Keep the full incremental bill below $500/month before credits. The +manifest requests 1.5 CPU and 3 GiB for application processes only; PostgreSQL, +Vespa, Temporal, embedding inference, Redis, Svix, disks and networking are +additional. Confirm the total node allocation and price before provisioning. + +## One release image, two processes + +`porter.yaml` builds the backend once and runs the API and Temporal worker. +The API is cluster-private, with no public load balancer. Initially verify it +through an authenticated operator port-forward; the Almanac bridge needs an +explicit reachable, authenticated network path before it can use this endpoint. +Do not mistake `private: true` for a private load balancer reachable from other +clouds. Do not expose PostgreSQL, Vespa or Temporal publicly. + +The entrypoint executes Uvicorn directly so termination reaches the application. +It never prints database URLs, runs migrations, or creates organizations. Worker +shutdown gets 90 seconds inside a 120-second pod grace period. Interrupted jobs +must recover through Temporal and the existing writer fencing, not an assumption +that shutdown always finishes. Test termination under actual capture before +claiming recovery is verified in this deployment. + +## Prerequisites + +Provision isolated PostgreSQL databases for application and Temporal state; +private Redis, Temporal and Svix; persistent Vespa with the matching application +package; local MiniLM inference; and a private S3 bucket for canonical bytes. +Use the repository's existing dependency configurations as inputs, not its local +Compose file as a production deployment: that file enables local auth and exposes +ports. Dependency image pins, persistent volumes, backups and a restore exercise +remain deployment work. The stock readiness route gates PostgreSQL only; verify +index, worker and capture readiness separately. + +Create the Porter environment group `almanac-source-store-staging` with only the +owned service's values. Supply database connection settings, Redis endpoint, +Temporal endpoint/namespace, Vespa endpoint, TEXT2VEC_INFERENCE_URL, +STORAGE_AWS_BUCKET, SVIX_URL, SVIX_JWT_SECRET and the application's required +encryption/signing secrets. Inspect `backend/airweave/core/config/settings.py` +for exact current fields. Do not copy the main product database from Doppler. +Use a least-privilege workload IAM role for this bucket and configure the Porter +AWS-role connection once that role exists. Do not add static root AWS keys. + +The manifest uses local embeddings to avoid assuming third-party model credits. +Semantic quality and resource use still need measurement. Do not claim generative +search is operational without separately configured model credentials and a live +search test. + +## Database setup and deployment + +Validate locally from the repository root: + +```sh +make validate-deploy +``` + +Against the explicitly verified isolated database, run the committed image as a +one-off operator job with migration credentials: + +```sh +alembic upgrade head +python -m airweave.db.bootstrap +``` + +See `backend/SERVICE_AUTH.md` for organization/key provisioning. Store each key in +secret management. Runtime database credentials should not have migration +privileges. Migrations are intentionally not automatic predeploy hooks: an image +rollback does not reverse database changes. + +After dependencies, IAM, secrets, cluster quota and the reviewed dev commit are +verified, use the existing Porter credentials without printing them: + +```sh +doppler run -p almanac -c dev -- porter apply --project 19566 --cluster 6108 --dry-run -f porter.yaml +doppler run -p almanac -c dev -- porter apply --project 19566 --cluster 6108 --wait -f porter.yaml +``` + +The Doppler invocation authenticates the operator CLI; it does not authorize +injecting every local variable into the application environment group. Build and +deploy from a clean reviewed dev checkout, not a feature branch with dirty work. +Record the resulting image digest and deployed commit. + +Before calling staging working, prove unauthorized requests fail; verify account +scope; capture actual provider records and bytes; read them from a new process; +query the live index; edit/delete disposable source fixtures where permitted; +interrupt and resume a worker; and restore isolated database/blob backups. Record +provider limitations separately from successful checks. Never infer full source +coverage from one healthy request. + +Manifest syntax reference: https://docs.porter.run/applications/configuration-as-code/reference +Private service semantics: https://docs.porter.run/applications/configuration-as-code/services/web-service diff --git a/porter.yaml b/porter.yaml new file mode 100644 index 0000000000..dbfd4eeb43 --- /dev/null +++ b/porter.yaml @@ -0,0 +1,54 @@ +# Application services only. Provision isolated dependencies and the environment +# group described in deploy/README.md before applying. Never attach almanac/dev +# wholesale: its database belongs to the existing product. +version: v2 +name: almanac-source-store +build: + method: docker + context: ./backend + dockerfile: ./backend/Dockerfile +envGroups: + - almanac-source-store-staging +env: + ENVIRONMENT: dev + LOCAL_DEVELOPMENT: "false" + AUTH_MODE: api_key + ANALYTICS_ENABLED: "false" + STRIPE_ENABLED: "false" + STORAGE_BACKEND: aws + STORAGE_AWS_REGION: us-east-1 + DENSE_EMBEDDER: local_minilm + EMBEDDING_DIMENSIONS: "384" + SPARSE_EMBEDDER: fastembed_bm25 + SYNC_MAX_WORKERS: "2" + SYNC_THREAD_POOL_SIZE: "8" + TEMPORAL_GRACEFUL_SHUTDOWN_TIMEOUT: "90" +services: + - name: api + type: web + run: ./entrypoint.sh + port: 8001 + private: true + instances: 1 + cpuCores: 0.5 + ramMegabytes: 1024 + terminationGracePeriodSeconds: 60 + readinessCheck: + enabled: true + httpPath: /health/ready + timeoutSeconds: 10 + initialDelaySeconds: 15 + livenessCheck: + enabled: true + httpPath: /health/live + timeoutSeconds: 5 + initialDelaySeconds: 30 + - name: worker + type: worker + run: python -m airweave.domains.temporal.worker + instances: 1 + cpuCores: 1 + ramMegabytes: 2048 + terminationGracePeriodSeconds: 120 +autoRollback: + enabled: true From 31885dc18064d13a1c109420d13ef8614278e0f1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:02:41 -0700 Subject: [PATCH 007/266] Serve scoped mail originals and reclaim retired search generations --- backend/airweave/adapters/storage/aws_s3.py | 22 +- .../airweave/adapters/storage/azure_blob.py | 20 +- .../airweave/adapters/storage/filesystem.py | 14 +- backend/airweave/adapters/storage/gcp_gcs.py | 23 +- backend/airweave/api/v1/endpoints/records.py | 57 +++- .../entities/canonical/blob_materializer.py | 14 +- .../domains/entities/canonical/calendar.py | 19 ++ .../entities/canonical/projection_gc.py | 139 +++++++++ .../entities/canonical/projection_mappers.py | 13 + .../entities/canonical/projection_models.py | 45 +++ .../entities/canonical/projection_store.py | 101 ++++++- .../domains/entities/canonical/projector.py | 36 ++- .../domains/entities/canonical/query.py | 114 +++++++- .../entities/canonical/query_models.py | 29 ++ .../domains/entities/canonical/query_store.py | 49 +++- .../entities/canonical/tests/conftest.py | 2 + .../entities/canonical/tests/helpers.py | 36 +++ .../tests/test_calendar_exclusions.py | 147 ++++++++++ .../canonical/tests/test_mail_http.py | 192 ++++++++++++ .../canonical/tests/test_projection.py | 46 ++- .../canonical/tests/test_projection_gc.py | 276 ++++++++++++++++++ .../canonical/tests/test_search_visibility.py | 4 +- .../canonical/tests/test_strict_projection.py | 23 ++ .../airweave/domains/storage/exceptions.py | 4 + backend/airweave/domains/storage/protocols.py | 5 +- .../cleanup_projection_generations.py | 45 +++ .../domains/temporal/worker/wiring.py | 4 + .../workflows/cleanup_stuck_sync_jobs.py | 10 + .../workflows/tests/test_cleanup_workflows.py | 8 +- backend/airweave/models/__init__.py | 3 + .../airweave/models/projection_generation.py | 31 ++ .../platform/destinations/vespa/client.py | 13 +- .../destinations/vespa/destination.py | 20 ++ .../sources/records/google_calendar.py | 15 +- .../versions/0003_mail_thread_index.py | 24 ++ .../versions/0004_projection_generation.py | 47 +++ .../entities/test_projection_mappers.py | 2 +- .../unit/domains/storage/test_bounded_read.py | 73 +++++ .../sources/records/test_google_calendar.py | 23 +- 39 files changed, 1698 insertions(+), 50 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/calendar.py create mode 100644 backend/airweave/domains/entities/canonical/projection_gc.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_mail_http.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection_gc.py create mode 100644 backend/airweave/domains/temporal/activities/cleanup_projection_generations.py create mode 100644 backend/airweave/models/projection_generation.py create mode 100644 backend/alembic/versions/0003_mail_thread_index.py create mode 100644 backend/alembic/versions/0004_projection_generation.py create mode 100644 backend/tests/unit/domains/storage/test_bounded_read.py diff --git a/backend/airweave/adapters/storage/aws_s3.py b/backend/airweave/adapters/storage/aws_s3.py index ef61ac7915..bb2f8371e7 100644 --- a/backend/airweave/adapters/storage/aws_s3.py +++ b/backend/airweave/adapters/storage/aws_s3.py @@ -12,6 +12,7 @@ from airweave.domains.storage.exceptions import ( StorageException, StorageNotFoundError, + StorageReadLimitExceeded, ) from airweave.domains.storage.protocols import StorageBackend @@ -127,15 +128,30 @@ async def write_file(self, path: str, content: bytes) -> None: except Exception as e: raise StorageException(f"Failed to write file to {path}: {e}") - async def read_file(self, path: str) -> bytes: + async def read_file(self, path: str, *, max_bytes: int | None = None) -> bytes: """Read binary content from S3.""" + from aiohttp import ClientError as HttpClientError + from botocore.exceptions import BotoCoreError, ClientError + + if max_bytes is not None and max_bytes < 0: + raise ValueError("max_bytes must be nonnegative") key = self._resolve(path) try: client = await self._get_client() response = await client.get_object(Bucket=self.bucket, Key=key) async with response["Body"] as stream: - return await stream.read() - except Exception as e: + if max_bytes is None: + return await stream.read() + content = bytearray() + while len(content) <= max_bytes: + chunk = await stream.read(min(65536, max_bytes + 1 - len(content))) + if not chunk: + return bytes(content) + content.extend(chunk) + raise StorageReadLimitExceeded("Stored object exceeds read limit") + except StorageReadLimitExceeded: + raise + except (OSError, HttpClientError, BotoCoreError, ClientError) as e: if "NoSuchKey" in str(e) or "404" in str(e): raise StorageNotFoundError(f"Path not found: {path}") raise StorageException(f"Failed to read file from {path}: {e}") diff --git a/backend/airweave/adapters/storage/azure_blob.py b/backend/airweave/adapters/storage/azure_blob.py index c47b625d62..5060f810de 100644 --- a/backend/airweave/adapters/storage/azure_blob.py +++ b/backend/airweave/adapters/storage/azure_blob.py @@ -12,6 +12,7 @@ from airweave.domains.storage.exceptions import ( StorageException, StorageNotFoundError, + StorageReadLimitExceeded, ) from airweave.domains.storage.protocols import StorageBackend @@ -133,19 +134,30 @@ async def write_file(self, path: str, content: bytes) -> None: except Exception as e: raise StorageException(f"Failed to write file to {path}: {e}") - async def read_file(self, path: str) -> bytes: + async def read_file(self, path: str, *, max_bytes: int | None = None) -> bytes: """Read binary content from Azure Blob.""" + from azure.core.exceptions import AzureError + + if max_bytes is not None and max_bytes < 0: + raise ValueError("max_bytes must be nonnegative") blob_path = self._resolve(path) try: container_client = await self._get_container_client() blob_client = container_client.get_blob_client(blob_path) if not await blob_client.exists(): raise StorageNotFoundError(f"Path not found: {path}") - download_stream = await blob_client.download_blob() - return await download_stream.readall() + download_stream = await blob_client.download_blob( + **({"offset": 0, "length": max_bytes + 1} if max_bytes is not None else {}) + ) + content = await download_stream.readall() + if max_bytes is not None and len(content) > max_bytes: + raise StorageReadLimitExceeded("Stored object exceeds read limit") + return content except StorageNotFoundError: raise - except Exception as e: + except StorageReadLimitExceeded: + raise + except (OSError, AzureError) as e: raise StorageException(f"Failed to read file from {path}: {e}") async def exists(self, path: str) -> bool: diff --git a/backend/airweave/adapters/storage/filesystem.py b/backend/airweave/adapters/storage/filesystem.py index f4b892ef94..0f8e0d8490 100644 --- a/backend/airweave/adapters/storage/filesystem.py +++ b/backend/airweave/adapters/storage/filesystem.py @@ -19,6 +19,7 @@ from airweave.domains.storage.exceptions import ( StorageException, StorageNotFoundError, + StorageReadLimitExceeded, ) from airweave.domains.storage.protocols import StorageBackend @@ -95,8 +96,10 @@ async def write_file(self, path: str, content: bytes) -> None: except Exception as e: raise StorageException(f"Failed to write file to {path}: {e}") - async def read_file(self, path: str) -> bytes: + async def read_file(self, path: str, *, max_bytes: int | None = None) -> bytes: """Read binary content from filesystem.""" + if max_bytes is not None and max_bytes < 0: + raise ValueError("max_bytes must be nonnegative") full_path = self._resolve(path) if not full_path.exists(): @@ -104,8 +107,13 @@ async def read_file(self, path: str) -> bytes: try: async with aiofiles.open(full_path, "rb") as f: - return await f.read() - except Exception as e: + content = await f.read() if max_bytes is None else await f.read(max_bytes + 1) + if max_bytes is not None and len(content) > max_bytes: + raise StorageReadLimitExceeded("Stored object exceeds read limit") + return content + except StorageReadLimitExceeded: + raise + except OSError as e: raise StorageException(f"Failed to read file from {path}: {e}") async def exists(self, path: str) -> bool: diff --git a/backend/airweave/adapters/storage/gcp_gcs.py b/backend/airweave/adapters/storage/gcp_gcs.py index 900d5755a8..a666007e2a 100644 --- a/backend/airweave/adapters/storage/gcp_gcs.py +++ b/backend/airweave/adapters/storage/gcp_gcs.py @@ -12,6 +12,7 @@ from airweave.domains.storage.exceptions import ( StorageException, StorageNotFoundError, + StorageReadLimitExceeded, ) from airweave.domains.storage.protocols import StorageBackend @@ -142,8 +143,13 @@ def _write(): except Exception as e: raise StorageException(f"Failed to write file to {path}: {e}") - async def read_file(self, path: str) -> bytes: + async def read_file(self, path: str, *, max_bytes: int | None = None) -> bytes: """Read binary content from GCS.""" + from google.api_core.exceptions import GoogleAPIError + from requests.exceptions import RequestException + + if max_bytes is not None and max_bytes < 0: + raise ValueError("max_bytes must be nonnegative") import asyncio blob_name = self._resolve(path) @@ -154,7 +160,16 @@ def _read(): bucket = self._get_bucket() blob = bucket.blob(blob_name) try: - return blob.download_as_bytes() + content = blob.download_as_bytes( + **( + {"start": 0, "end": max_bytes, "raw_download": True} + if max_bytes is not None + else {} + ) + ) + if max_bytes is not None and len(content) > max_bytes: + raise StorageReadLimitExceeded("Stored object exceeds read limit") + return content except NotFound: raise StorageNotFoundError(f"Path not found: {path}") @@ -162,7 +177,9 @@ def _read(): return await asyncio.to_thread(_read) except StorageNotFoundError: raise - except Exception as e: + except StorageReadLimitExceeded: + raise + except (OSError, GoogleAPIError, RequestException) as e: raise StorageException(f"Failed to read file from {path}: {e}") async def exists(self, path: str) -> bool: diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index a4b752e44a..f47059306c 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -3,17 +3,19 @@ from typing import Literal from uuid import UUID -from fastapi import Depends, Query, Request +from fastapi import Depends, Path, Query, Request, Response from fastapi.responses import JSONResponse from sqlalchemy.ext.asyncio import AsyncSession from airweave.api import deps from airweave.api.context import ApiContext from airweave.api.router import TrailingSlashRouter +from airweave.core.container import Container from airweave.db.session import get_db from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query import CanonicalQueryService from airweave.domains.entities.canonical.query_models import ( + MailThreadPage, RecordChangePage, RecordFilters, RecordListQuery, @@ -26,10 +28,16 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> JSONResponse: """Preserve machine-readable recovery codes without exposing SQL/provider diagnostics.""" - status = 404 if error.code in {"source_not_found", "record_not_found"} else 400 + status = { + "source_not_found": 404, + "record_not_found": 404, + "blob_not_found": 404, + "stale_record_revision": 409, + "blob_unavailable": 503, + }.get(error.code, 400) return JSONResponse( status_code=status, - content={"error": {"code": error.code, "message": str(error), "retryable": False}}, + content={"error": {"code": error.code, "message": str(error), "retryable": status == 503}}, ) @@ -77,3 +85,46 @@ async def read_record( ) -> SourceRecord: """Read current committed provider state, including tombstones and completeness.""" return await service.read(db, ctx.organization.id, sync_id, record_id) + + +@router.get("/{sync_id}/mail/threads/{thread_id}", response_model=MailThreadPage) +async def mail_thread( + sync_id: UUID, + thread_id: str = Path(min_length=1, max_length=512), + cursor: str | None = None, + limit: int = Query(100, ge=1, le=100), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), +) -> MailThreadPage: + """Read observed Gmail messages in chronological order without querying the provider.""" + return await service.mail_thread( + db, ctx.organization.id, sync_id, thread_id, cursor=cursor, limit=limit + ) + + +@router.get("/{sync_id}/records/{record_id}/blobs/{sha256}") +async def record_blob( + sync_id: UUID, + record_id: UUID, + sha256: str = Path(pattern=r"^[a-f0-9]{64}$"), + revision: int = Query(ge=1), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), + container: Container = Depends(deps.get_container), +) -> Response: + """Return verified owned bytes only; clients cannot supply storage keys or remote URLs.""" + content = await service.blob( + db, ctx.organization.id, sync_id, record_id, revision, sha256, container.storage_backend + ) + return Response( + content, + media_type="application/octet-stream", + headers={ + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "Content-Disposition": "attachment", + "ETag": '"' + sha256 + '"', + }, + ) diff --git a/backend/airweave/domains/entities/canonical/blob_materializer.py b/backend/airweave/domains/entities/canonical/blob_materializer.py index 3562243f8c..cefb1bc956 100644 --- a/backend/airweave/domains/entities/canonical/blob_materializer.py +++ b/backend/airweave/domains/entities/canonical/blob_materializer.py @@ -12,18 +12,22 @@ from airweave.domains.storage.protocols import StorageBackend +class BlobIntegrityError(ValueError): + """Committed reference or bytes violate canonical ownership/integrity.""" + + async def read_blob(record: SourceRecord, ref: BlobReference, storage: StorageBackend) -> bytes: """Read only this sync's content-addressed bytes; reject corrupt references/content.""" expected = f"canonical/{record.sync_id}/blobs/sha256/{ref.sha256}" if ref.key != expected: - raise ValueError("Canonical blob key does not match source scope and digest") + raise BlobIntegrityError("Canonical blob key does not match source scope and digest") if ref not in record.blobs: - raise ValueError("Canonical blob reference does not belong to the source record") + raise BlobIntegrityError("Canonical blob reference does not belong to the source record") if ref.size_bytes > MAX_FILE_SIZE_BYTES: - raise ValueError("Canonical blob exceeds projection size limit") - content = await storage.read_file(ref.key) + raise BlobIntegrityError("Canonical blob exceeds projection size limit") + content = await storage.read_file(ref.key, max_bytes=ref.size_bytes) if len(content) != ref.size_bytes or hashlib.sha256(content).hexdigest() != ref.sha256: - raise ValueError("Canonical blob bytes do not match recorded size and digest") + raise BlobIntegrityError("Canonical blob bytes do not match recorded size and digest") return content diff --git a/backend/airweave/domains/entities/canonical/calendar.py b/backend/airweave/domains/entities/canonical/calendar.py new file mode 100644 index 0000000000..d8ca49d279 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/calendar.py @@ -0,0 +1,19 @@ +"""Google Calendar recurrence exclusions are scheduling state, not live meetings.""" + +from collections.abc import Mapping + + +def is_cancelled_recurring_event(payload: Mapping[str, object]) -> bool: + """Recognize the provider-guaranteed identity of a cancelled series occurrence.""" + if payload.get("status") != "cancelled": + return False + recurring_id = payload.get("recurringEventId") + original = payload.get("originalStartTime") + return ( + isinstance(recurring_id, str) + and bool(recurring_id) + and isinstance(original, dict) + and any( + isinstance(original.get(key), str) and original[key] for key in ("date", "dateTime") + ) + ) 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..c7614c1b98 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_gc.py @@ -0,0 +1,139 @@ +"""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.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() + return ProjectionCleanupPage( + generation=generation, + attempt=row.gc_attempt, + cursor=row.delete_cursor, + 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) + if row.delete_cursor >= len(row.documents): + 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_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 9980254a3d..d1993b05ee 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -12,6 +12,7 @@ from pydantic import JsonValue +from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.storage.protocols import StorageBackend from airweave.platform.entities._base import BaseEntity @@ -196,6 +197,15 @@ def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: ) +def excluded_from_search(record: SourceRecord, source_name: str) -> bool: + """Retained calendar exclusions intentionally publish no searchable meeting.""" + return ( + source_name == "google_calendar" + and record.identity.record_type == "event" + and is_cancelled_recurring_event(record.payload) + ) + + @asynccontextmanager async def map_record( record: SourceRecord, @@ -205,6 +215,9 @@ async def map_record( """Keep verified local blob files alive only while strict projection consumes them.""" if record.deleted_at is not None or record.content_access != "available": raise ProjectionMappingError("Unavailable records cannot be projected") + if excluded_from_search(record, source_name): + yield () + return with TemporaryDirectory(prefix="airweave-projection-") as temporary: directory = Path(temporary) if source_name == "gmail": diff --git a/backend/airweave/domains/entities/canonical/projection_models.py b/backend/airweave/domains/entities/canonical/projection_models.py index d40352a021..f2fecac0dc 100644 --- a/backend/airweave/domains/entities/canonical/projection_models.py +++ b/backend/airweave/domains/entities/canonical/projection_models.py @@ -1,5 +1,6 @@ """Typed immutable search publication identities, independent of provider identity.""" +import re from uuid import UUID from pydantic import BaseModel, ConfigDict, Field @@ -59,3 +60,47 @@ class ProjectionBatchResult(BaseModel): published: int = 0 superseded: int = 0 failed: int = 0 + + +class ProjectionDocument(BaseModel): + """Exact remote identifier, without searchable text or embeddings.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + schema_name: str = Field(min_length=1) + document_id: str = Field(min_length=1) + + +class ProjectionCleanupPage(BaseModel): + """One idempotent deletion page of an irrevocably retired generation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + generation: UUID + attempt: UUID + cursor: int + documents: tuple[ProjectionDocument, ...] + + +def scope_projection_document_id(sync_id: UUID, collection_id: UUID, document_id: str) -> str: + """Bind canonical remote IDs to the destination without changing search locators.""" + scoped = f"{sync_id}_{collection_id}_{document_id}" + projection_document_locator(scoped, sync_id, collection_id) + return scoped + + +def projection_document_locator( + document_id: str, sync_id: UUID, collection_id: UUID +) -> ProjectionLocator: + """Reject foreign prefixes and ambiguous substring matches before manifest commit.""" + prefix = f"{sync_id}_{collection_id}_" + if not document_id.startswith(prefix): + raise ValueError("Projection document scope does not match destination") + match = re.fullmatch( + r"[A-Za-z][A-Za-z0-9_]*_(canonical:v1:[^_]+)__chunk_(0|[1-9][0-9]*)", + document_id[len(prefix) :], + ) + if match is None: + raise ValueError("Invalid canonical remote document ID") + locator = ProjectionLocator.parse(match[1]) + if locator is None or locator.encode() != match[1]: + raise ValueError("Noncanonical projection locator encoding") + return locator diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index 4b2342ac48..c115ee3589 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -1,14 +1,22 @@ """Canonical pending work and atomic publication. No network I/O under row locks.""" +from datetime import datetime, timedelta, timezone from uuid import UUID from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork -from airweave.domains.entities.canonical.projection_models import ProjectionLocator, ProjectionWork +from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event +from airweave.domains.entities.canonical.projection_models import ( + ProjectionDocument, + ProjectionLocator, + ProjectionWork, + projection_document_locator, +) from airweave.domains.entities.canonical.store import content_is_available, source_record from airweave.models.entity import Entity +from airweave.models.projection_generation import ProjectionGeneration from airweave.models.sync import Sync @@ -106,6 +114,64 @@ async def _current(self, db: AsyncSession, work: ProjectionWork) -> Entity | Non return None return entity + async def prepare( + self, + db: AsyncSession, + work: ProjectionWork, + generation: UUID, + collection_id: UUID, + documents: tuple[ProjectionDocument, ...], + ) -> bool: + """Commit exact immutable deletion identities before the first remote feed.""" + identities = set() + for document in documents: + try: + locator = projection_document_locator( + document.document_id, work.record.sync_id, collection_id + ) + valid = ( + locator.record_id == work.record.id + and locator.revision == work.record.revision + and locator.pipeline_version == work.pipeline_version + and locator.generation == generation + ) + except (ValueError, IndexError, AttributeError): + valid = False + if not valid or (document.schema_name, document.document_id) in identities: + raise ValueError("Projection manifest must contain unique IDs of this generation") + identities.add((document.schema_name, document.document_id)) + manifest = [document.model_dump() for document in documents] + async with UnitOfWork(db): + if await self._current(db, work) is None: + return False + prior = await db.get(ProjectionGeneration, generation) + if prior is not None: + if ( + prior.documents != manifest + or prior.record_id != work.record.id + or prior.revision != work.record.revision + or prior.pipeline_version != work.pipeline_version + or prior.organization_id != work.organization_id + or prior.collection_id != collection_id + ): + raise ValueError("Projection manifest is immutable") + return prior.retired_at is None + db.add( + ProjectionGeneration( + id=generation, + organization_id=work.organization_id, + sync_id=work.record.sync_id, + collection_id=collection_id, + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + documents=manifest, + next_gc_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + ) + await db.flush() + return True + async def publish( self, db: AsyncSession, @@ -114,14 +180,41 @@ async def publish( chunk_count: int, ) -> bool: """Publish only a complete current generation; late writers cannot replace it.""" - if chunk_count < 0 or (work.record.deleted_at is None and chunk_count == 0): + exclusion = work.record.identity.record_type == "event" and is_cancelled_recurring_event( + work.record.payload + ) + no_content = work.record.deleted_at is not None or exclusion + if chunk_count < 0 or (not no_content and chunk_count == 0): raise ValueError("Active records require a nonempty complete projection") - if work.record.deleted_at is not None and chunk_count != 0: - raise ValueError("Deleted records must publish no content") + if no_content and chunk_count != 0: + raise ValueError("Deleted records and recurrence exclusions must publish no content") async with UnitOfWork(db): entity = await self._current(db, work) if entity is None: return False + attempt = await db.scalar( + select(ProjectionGeneration) + .where( + ProjectionGeneration.id == generation, + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + if ( + attempt is None + or attempt.retired_at is not None + or attempt.organization_id != work.organization_id + or attempt.record_id != work.record.id + or attempt.revision != work.record.revision + or attempt.pipeline_version != work.pipeline_version + or len(attempt.documents) != chunk_count + ): + return False + if work.previous_generation is not None: + previous = await db.get(ProjectionGeneration, work.previous_generation) + if previous is not None: + previous.retired_at = datetime.now(timezone.utc) + previous.next_gc_at = datetime.now(timezone.utc) entity.indexed_revision = work.record.revision entity.indexed_pipeline_version = work.pipeline_version entity.indexed_generation = generation diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index ba17e95c10..66277fa480 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -10,8 +10,10 @@ from airweave.core.logging import ContextualLogger from airweave.domains.entities.canonical.projection_models import ( ProjectionBatchResult, + ProjectionDocument, ProjectionLocator, ProjectionWork, + scope_projection_document_id, ) from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore from airweave.domains.storage.protocols import StorageBackend @@ -69,11 +71,17 @@ async def project_one( logger: ContextualLogger, ) -> bool: """Never feed native capture JSON; mapper emits explicit safe projection entities.""" - from airweave.domains.entities.canonical.projection_mappers import map_record + from airweave.domains.entities.canonical.projection_mappers import ( + excluded_from_search, + map_record, + ) generation = uuid4() chunks = [] - if work.record.deleted_at is None: + no_documents = work.record.deleted_at is not None or excluded_from_search( + work.record, source_name + ) + if not no_documents: async with map_record(work.record, source_name, self._storage) as mapped: if not mapped: raise ValueError("Projection mapper returned no required content") @@ -101,7 +109,29 @@ async def project_one( ProjectionRuntime(StrictProjectionTracker()), strict=True, ) - await destination.bulk_insert(chunks, strict=True) + prepared = destination.prepare_documents(chunks) + for group in prepared.values(): + for document in group: + document.id = scope_projection_document_id( + work.record.sync_id, destination.collection_id, document.id + ) + manifest = tuple( + ProjectionDocument(schema_name=doc.schema_name, document_id=doc.id) + for group in prepared.values() + for doc in group + ) + async with self._sessions() as db: + if not await self._store.prepare( + db, work, generation, destination.collection_id, manifest + ): + return False + await destination.feed_prepared(prepared) + if no_documents: + async with self._sessions() as db: + if not await self._store.prepare( + db, work, generation, destination.collection_id, () + ): + return False async with self._sessions() as db: return await self._store.publish(db, work, generation, len(chunks)) diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py index 3e649aa924..07895e1690 100644 --- a/backend/airweave/domains/entities/canonical/query.py +++ b/backend/airweave/domains/entities/canonical/query.py @@ -6,8 +6,11 @@ from pydantic import ValidationError from sqlalchemy.ext.asyncio import AsyncSession +from airweave.domains.entities.canonical.blob_materializer import BlobIntegrityError, read_blob from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import ( + MailThreadCursor, + MailThreadPage, RecordChangePage, RecordCursor, RecordListQuery, @@ -15,6 +18,8 @@ ) from airweave.domains.entities.canonical.query_store import CanonicalQueryStore from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.storage.exceptions import StorageException +from airweave.domains.storage.protocols import StorageBackend class InvalidRecordCursor(CanonicalStoreError): @@ -29,6 +34,24 @@ class RecordNotFound(CanonicalStoreError): code = "record_not_found" +class StaleRecordRevision(CanonicalStoreError): + """The client must reread the current record before using its blob reference.""" + + code = "stale_record_revision" + + +class BlobNotFound(CanonicalStoreError): + """No authorized committed bytes match this record and digest.""" + + code = "blob_not_found" + + +class BlobUnavailable(CanonicalStoreError): + """Committed bytes are missing or fail integrity verification.""" + + code = "blob_unavailable" + + class CanonicalQueryService: """Own read semantics; cursor signatures never replace tenant authorization.""" @@ -40,7 +63,7 @@ def __init__( self.queries = queries self.signing_key = signing_key - def _encode(self, cursor: RecordCursor) -> str: + def _encode(self, cursor: RecordCursor | MailThreadCursor) -> str: return jwt.encode(cursor.model_dump(mode="json"), self.signing_key, algorithm="HS256") def _decode(self, token: str, organization_id: UUID, sync_id: UUID, mode: str) -> RecordCursor: @@ -133,3 +156,92 @@ async def changes( has_more=page.has_more, high_watermark=page.high_watermark, ) + + async def mail_thread( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + thread_id: str, + *, + cursor: str | None = None, + limit: int = 100, + ) -> MailThreadPage: + """Return stored messages only; observation time never proves complete coverage.""" + if not thread_id or len(thread_id) > 512 or not 1 <= limit <= 100: + raise ValueError("Invalid thread query bounds") + position = None + if cursor is not None: + try: + position = MailThreadCursor.model_validate( + jwt.decode(cursor, self.signing_key, algorithms=["HS256"]) + ) + except (JWTError, ValidationError, ValueError) as exc: + raise InvalidRecordCursor("Invalid thread cursor; restart this query") from exc + if (position.organization_id, position.sync_id, position.thread_id) != ( + organization_id, + sync_id, + thread_id, + ): + raise InvalidRecordCursor("Cursor belongs to a different account or thread") + rows = await self.queries.mail_thread( + db, + organization_id, + sync_id, + thread_id, + after_id=position.after_id if position else None, + after_created_at=position.after_created_at if position else None, + limit=limit, + ) + more = len(rows) > limit + messages = rows[:limit] + next_cursor = None + if more: + last = messages[-1] + next_cursor = self._encode( + MailThreadCursor( + organization_id=organization_id, + sync_id=sync_id, + thread_id=thread_id, + after_created_at=last.source_created_at, + after_id=last.id, + ) + ) + return MailThreadPage( + thread_id=thread_id, messages=messages, next_cursor=next_cursor, has_more=more + ) + + async def blob( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + record_id: UUID, + revision: int, + sha256: str, + storage: StorageBackend, + ) -> bytes: + """Resolve only committed references and recheck visibility after storage I/O.""" + record = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(record, revision) + ref = next((item for item in record.blobs if item.sha256 == sha256), None) + if ref is None: + raise BlobNotFound("Blob reference not found on this record") + try: + content = await read_blob(record, ref, storage) + except (StorageException, BlobIntegrityError) as exc: + raise BlobUnavailable("Committed blob is unavailable; retry later") from exc + # Drop identity-map values before the second READ COMMITTED visibility check. + db.expire_all() + current = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(current, revision) + if ref not in current.blobs: + raise BlobNotFound("Blob reference no longer belongs to this record") + return content + + @staticmethod + def _check_blob_record(record: SourceRecord, revision: int) -> None: + if record.content_access != "available" or record.deleted_at is not None: + raise RecordNotFound("Record content is not currently available") + if record.revision != revision: + raise StaleRecordRevision("Record changed; reread before requesting its blob") diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index 20a289e3f6..ea36e14e9f 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -1,5 +1,6 @@ """Typed exact-record queries; search ranking is a separate index concern.""" +from datetime import datetime from typing import Literal from uuid import UUID @@ -60,3 +61,31 @@ class RecordCursor(BaseModel): after_id: UUID | None = None after_sequence: int = Field(default=0, ge=0) high_watermark: int | None = Field(default=None, ge=0) + + +class MailThreadPage(BaseModel): + """Observed messages only; neither a complete thread nor a provider snapshot.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + thread_id: str + messages: tuple[SourceRecord, ...] + next_cursor: str | None + has_more: bool + consistency: Literal["live"] = "live" + order: Literal["source_created_at_asc_nulls_last_id_asc"] = ( + "source_created_at_asc_nulls_last_id_asc" + ) + coverage: Literal["stored_messages_only"] = "stored_messages_only" + + +class MailThreadCursor(BaseModel): + """Signed thread and account traversal position, independent of capture containers.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + purpose: Literal["canonical_mail_thread"] = "canonical_mail_thread" + version: Literal[1] = 1 + organization_id: UUID + sync_id: UUID + thread_id: str + after_created_at: datetime | None + after_id: UUID diff --git a/backend/airweave/domains/entities/canonical/query_store.py b/backend/airweave/domains/entities/canonical/query_store.py index 5931ccb8e5..d6bbf7d4cc 100644 --- a/backend/airweave/domains/entities/canonical/query_store.py +++ b/backend/airweave/domains/entities/canonical/query_store.py @@ -1,8 +1,9 @@ """Read-only SQL for exact canonical record listing.""" +from datetime import datetime from uuid import UUID -from sqlalchemy import func, select +from sqlalchemy import and_, func, literal_column, or_, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.models import SourceRecord @@ -80,3 +81,49 @@ async def list_records( ) for row in rows ) + + async def mail_thread( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + thread_id: str, + *, + after_id: UUID | None, + after_created_at: datetime | None, + limit: int, + ) -> tuple[SourceRecord, ...]: + """Use native Gmail thread identity, keeping mailbox capture containers unchanged.""" + scope = await db.scalar( + select(Sync.id).where(Sync.id == sync_id, Sync.organization_id == organization_id) + ) + if scope is None: + raise SourceNotFound("Source does not exist in this organization") + statement = select(Entity).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.entity_definition_short_name == "message", + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + Entity.source_payload.op("->>")(literal_column("'threadId'")) == thread_id, + ) + if after_id is not None: + if after_created_at is None: + statement = statement.where( + Entity.source_created_at.is_(None), Entity.id > after_id + ) + else: + statement = statement.where( + or_( + Entity.source_created_at > after_created_at, + Entity.source_created_at.is_(None), + and_(Entity.source_created_at == after_created_at, Entity.id > after_id), + ) + ) + rows = await db.scalars( + statement.order_by(Entity.source_created_at.asc().nulls_last(), Entity.id).limit( + limit + 1 + ) + ) + return tuple(source_record(row) for row in rows) diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 89f7efc717..267c88d6fa 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -105,6 +105,8 @@ async def database(request): await connection.run_sync(seed_legacy) await connection.run_sync(migrate, "0001_canonical_records.py") await connection.run_sync(migrate, "0002_projection_publication.py") + await connection.run_sync(migrate, "0003_mail_thread_index.py") + await connection.run_sync(migrate, "0004_projection_generation.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/helpers.py b/backend/airweave/domains/entities/canonical/tests/helpers.py index 30a8a73d9f..940a9cc5a0 100644 --- a/backend/airweave/domains/entities/canonical/tests/helpers.py +++ b/backend/airweave/domains/entities/canonical/tests/helpers.py @@ -21,3 +21,39 @@ def observation(native_id="one", container_id=None, **changes): async def capture(database, service, fence, *records): async with database() as db: return await service.capture(db, CaptureBatch(fence=fence, records=records)) + + +async def publish_prepared(store, db, work, generation, count, collection_id=None): + """Synthetic exact manifests exercise the same mandatory pre-feed transaction.""" + from uuid import uuid4 + + from airweave.domains.entities.canonical.projection_models import ( + ProjectionDocument, + ProjectionLocator, + scope_projection_document_id, + ) + + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=0, + ).encode() + collection_id = collection_id or uuid4() + prepared = await store.prepare( + db, + work, + generation, + collection_id, + tuple( + ProjectionDocument( + schema_name="base_entity", + document_id=scope_projection_document_id( + work.record.sync_id, collection_id, f"Entity_{locator}__chunk_{index}" + ), + ) + for index in range(count) + ), + ) + return prepared and await store.publish(db, work, generation, count) diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py new file mode 100644 index 0000000000..cf37578f57 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -0,0 +1,147 @@ +"""Cancelled instances remain recurrence facts, without becoming searchable meetings.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.projection_mappers import map_record +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_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 +from airweave.models import Entity, ProjectionGeneration +from airweave.platform.sources.records.google_calendar import record + + +async def test_cancelled_exception_read_project_reinstate_and_revoke(database, source): + capture, fence = source + parent = record("calendar", {"id": "cal"}) + native = { + "id": "instance", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-10-01"}, + } + excluded = record("event", native, "cal").model_copy(update={"parent": parent.identity}) + async with database() as db: + result = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, excluded))) + row = result.changes[-1].record + queries = CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "test-key") + projections = CanonicalProjectionStore() + async with database() as db: + current = await queries.read(db, fence.organization_id, fence.sync_id, row.id) + assert current.deleted_at is None and current.payload == native + page = await queries.list_records( + db, + fence.organization_id, + fence.sync_id, + RecordListQuery(filters=RecordFilters(record_type="event")), + ) + assert len(page.records) == 1 and page.records[0].payload["status"] == "cancelled" + work = next( + w + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + if w.record.id == row.id + ) + async with map_record(current, "google_calendar", AsyncMock()) as mapped: + assert mapped == () + processor = MagicMock() + processor.process = AsyncMock() + destination = MagicMock() + destination.collection_id = uuid4() + destination.feed_prepared = AsyncMock() + projector = CanonicalProjector(projections, database, processor, AsyncMock()) + assert await projector.project_one(work, "google_calendar", destination, MagicMock()) + processor.process.assert_not_awaited() + destination.feed_prepared.assert_not_awaited() + async with database() as db: + assert all( + w.record.id != row.id + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + ) + manifest = await db.scalar( + select(ProjectionGeneration).where(ProjectionGeneration.record_id == row.id) + ) + assert manifest.documents == [] + assert (await db.get(Entity, row.id)).indexed_chunk_count == 0 + restored = record( + "event", + { + **native, + "status": "confirmed", + "summary": "Restored", + "start": {"date": "2026-10-01"}, + "end": {"date": "2026-10-02"}, + }, + "cal", + ).model_copy(update={"parent": parent.identity}) + async with database() as db: + result = await capture.capture(db, CaptureBatch(fence=fence, records=(restored,))) + assert result.changes[0].record.id == row.id and result.changes[0].record.revision == 2 + assert any( + w.record.id == row.id + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + ) + current = await queries.read(db, fence.organization_id, fence.sync_id, row.id) + async with map_record(current, "google_calendar", AsyncMock()) as mapped: + assert len(mapped) == 1 and mapped[0].status == "confirmed" + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(excluded,))) + await capture.capture( + db, + CaptureBatch( + fence=fence, + records=( + parent.model_copy( + update={"kind": "delete", "removal_reason": "access_revoked"} + ), + ), + ), + ) + hidden = await queries.read(db, fence.organization_id, fence.sync_id, row.id) + assert hidden.content_access == "unavailable" and hidden.payload == {} and not hidden.blobs + page = await queries.list_records( + db, + fence.organization_id, + fence.sync_id, + RecordListQuery(filters=RecordFilters(record_type="event")), + ) + assert not page.records + assert all( + w.record.id != row.id + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + ) + + +async def test_live_event_cannot_publish_empty_or_exclusion_nonempty(database, source): + capture, fence = source + fact = record( + "event", + { + "id": "cancelled", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-10-01"}, + }, + "cal", + ) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(fact,))) + store = CanonicalProjectionStore() + work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] + with pytest.raises(ValueError, match="exclusions"): + await store.publish(db, work, uuid4(), 1) + active = work.model_copy( + update={ + "record": work.record.model_copy( + update={"payload": {"id": "cancelled", "status": "confirmed"}} + ) + } + ) + with pytest.raises(ValueError, match="nonempty"): + await store.publish(db, active, uuid4(), 0) 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_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py index e830546730..6a94751a20 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py @@ -11,7 +11,7 @@ publication_matches, ) from airweave.domains.entities.canonical.requests import RecordIdentity -from airweave.domains.entities.canonical.tests.helpers import capture, observation +from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared from airweave.models import Entity, Sync pytestmark = pytest.mark.integration @@ -29,9 +29,9 @@ async def test_publication_compare_and_swap_revision_generation_and_pipeline(dat store = CanonicalProjectionStore() generation = uuid4() async with database() as db: - assert await store.publish(db, work, generation, 3) + assert await publish_prepared(store, db, work, generation, 3) async with database() as db: - assert not await store.publish(db, work, uuid4(), 3) # timed-out prior feed + 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, @@ -52,13 +52,13 @@ async def test_publication_compare_and_swap_revision_generation_and_pipeline(dat await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator))) is None ) - assert not await store.publish(db, work, uuid4(), 3) + 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 store.publish(db, changed, uuid4(), 3) + assert not await publish_prepared(store, db, changed, uuid4(), 3) assert (await pending(database, fence))[0].pipeline_version == 2 @@ -78,11 +78,11 @@ async def test_parent_access_loss_fences_publication_and_tombstone_is_empty(data ) store = CanonicalProjectionStore() async with database() as db: - assert not await store.publish(db, child_work, uuid4(), 2) + 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 store.publish(db, rows[0], uuid4(), 0) + assert await publish_prepared(store, db, rows[0], uuid4(), 0) assert not await pending(database, fence) @@ -95,7 +95,7 @@ async def test_failed_work_retries_and_stale_failure_preserves_publication(datab await store.fail(db, work, "conversion failed") assert len(await pending(database, fence)) == 1 async with database() as db: - assert await store.publish(db, work, uuid4(), 1) + 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: @@ -113,6 +113,7 @@ 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 await capture( @@ -140,10 +141,33 @@ async def process(entities, context, runtime, *, strict): processor.process = process destination = MagicMock() - destination.bulk_insert = AsyncMock() + destination.collection_id = uuid4() + 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() projector = CanonicalProjector(CanonicalProjectionStore(), database, processor, storage) assert await projector.project_one(work, "slack", destination, MagicMock()) - destination.bulk_insert.assert_awaited_once() - assert destination.bulk_insert.call_args.kwargs == {"strict": True} + destination.feed_prepared.assert_awaited_once() assert not await pending(database, fence) 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..73a635ce87 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py @@ -0,0 +1,276 @@ +"""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 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 + 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 = uuid4() + 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)) + 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: + 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 + 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, uuid4(), ()) + 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_search_visibility.py b/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py index 305a72365a..a48cdb9fc7 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py +++ b/backend/airweave/domains/entities/canonical/tests/test_search_visibility.py @@ -7,7 +7,7 @@ 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 +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, @@ -73,7 +73,7 @@ async def test_search_rejects_stale_foreign_malformed_and_legacy_publications(da work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] generation = uuid4() async with database() as db: - assert await store.publish(db, work, generation, 1) + assert await publish_prepared(store, db, work, generation, 1) locator = ProjectionLocator( record_id=work.record.id, revision=1, diff --git a/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py index 0fb4d6b5b5..55eb2877b8 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_strict_projection.py @@ -45,3 +45,26 @@ async def test_strict_destination_cannot_ack_dropped_transformations(): 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/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/protocols.py b/backend/airweave/domains/storage/protocols.py index f36c7d2c84..202db64196 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 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..c6a46941b9 --- /dev/null +++ b/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py @@ -0,0 +1,45 @@ +"""Bounded reclamation on the existing system maintenance workflow.""" + +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.platform.destinations.vespa.client import VespaClient + + +class CleanupProjectionGenerationsActivity: + """A tick deletes at most 100 exact document IDs; failures remain durable.""" + + @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] + ) + 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) + finally: + await client.close() diff --git a/backend/airweave/domains/temporal/worker/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py index 7bd3b1e4ae..6fff40c9f9 100644 --- a/backend/airweave/domains/temporal/worker/wiring.py +++ b/backend/airweave/domains/temporal/worker/wiring.py @@ -37,6 +37,9 @@ def create_activities() -> list: 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.project_canonical_records import ( ProjectCanonicalRecordsActivity, ) @@ -58,6 +61,7 @@ def create_activities() -> list: logger.debug("Wiring activities with container dependencies") return [ + CleanupProjectionGenerationsActivity().run, ProjectCanonicalRecordsActivity( projector=CanonicalProjector( CanonicalProjectionStore(), 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..28b391ec0c 100644 --- a/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py +++ b/backend/airweave/domains/temporal/workflows/cleanup_stuck_sync_jobs.py @@ -9,6 +9,9 @@ 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, + ) _CLEANUP_TIMEOUT = timedelta(minutes=5) _CLEANUP_RETRY = RetryPolicy( @@ -35,3 +38,10 @@ async def run(self) -> None: 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, + ) 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..5c980701d7 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 @@ -27,13 +28,17 @@ @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", ()) + 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)], + activities=[mock_cleanup_stuck_sync_jobs(recorder), cleanup_generations], ): await env.client.execute_workflow( CleanupStuckSyncJobsWorkflow.run, @@ -41,6 +46,7 @@ async def test_cleanup_stuck_sync_jobs_calls_activity(): task_queue=TASK_QUEUE, ) + assert recorder.called("cleanup_generations") assert recorder.called("cleanup_stuck") assert recorder.call_count("cleanup_stuck") == 1 diff --git a/backend/airweave/models/__init__.py b/backend/airweave/models/__init__.py index b995014cad..43f5ccae23 100644 --- a/backend/airweave/models/__init__.py +++ b/backend/airweave/models/__init__.py @@ -39,6 +39,7 @@ "Collection", "Entity", "EntityChange", + "ProjectionGeneration", "EntityCount", "Connection", "ConnectionInitSession", @@ -63,3 +64,5 @@ "UserOrganization", "VectorDbDeploymentMetadata", ] + +from .projection_generation import ProjectionGeneration diff --git a/backend/airweave/models/projection_generation.py b/backend/airweave/models/projection_generation.py new file mode 100644 index 0000000000..c87764ef1d --- /dev/null +++ b/backend/airweave/models/projection_generation.py @@ -0,0 +1,31 @@ +"""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) + 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/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 5c40fbfd4e..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") + 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 diff --git a/backend/airweave/platform/sources/records/google_calendar.py b/backend/airweave/platform/sources/records/google_calendar.py index aca9241213..f3abb57f47 100644 --- a/backend/airweave/platform/sources/records/google_calendar.py +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue +from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event from airweave.domains.entities.canonical.requests import ( CaptureRecord, CompletedScope, @@ -47,7 +48,19 @@ def record( native_id = payload.get("id") if not isinstance(native_id, str) or not native_id: raise ValueError("Calendar returned a resource without an ID") - deleted = payload.get("status") == "cancelled" or payload.get("deleted") is True + 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. + deleted = (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, 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/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 508849d4ec..405f32df0c 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -160,7 +160,7 @@ async def test_drive_uses_verified_owned_bytes_and_cleans_materialization(): assert path.name == digest + ".txt" assert entities[0].mime_type == "text/plain" assert not path.exists() - storage.read_file.assert_awaited_once_with(blob.key) + storage.read_file.assert_awaited_once_with(blob.key, max_bytes=blob.size_bytes) def test_added_typed_entities_are_registered(): diff --git a/backend/tests/unit/domains/storage/test_bounded_read.py b/backend/tests/unit/domains/storage/test_bounded_read.py new file mode 100644 index 0000000000..a9ccf78b63 --- /dev/null +++ b/backend/tests/unit/domains/storage/test_bounded_read.py @@ -0,0 +1,73 @@ +"""Actual oversized storage objects cannot bypass a canonical reference size bound.""" + +import asyncio +from unittest.mock import AsyncMock, Mock + +import pytest + +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 StorageReadLimitExceeded + + +async def test_filesystem_reads_actual_bound_and_preserves_unbounded_callers(tmp_path): + storage = FilesystemBackend(tmp_path) + await storage.write_file("blob", b"12345") + with pytest.raises(StorageReadLimitExceeded): + await storage.read_file("blob", max_bytes=3) + assert await storage.read_file("blob", max_bytes=5) == b"12345" + assert await storage.read_file("blob") == b"12345" + await storage.write_file("empty", b"") + assert await storage.read_file("empty", max_bytes=0) == b"" + + +@pytest.mark.parametrize("cancel", [False, True]) +async def test_s3_stream_closes_on_oversize_and_cancellation(cancel): + storage = S3Backend(bucket="synthetic", region="us-east-1") + stream = AsyncMock() + stream.__aenter__.return_value = stream + stream.read.side_effect = asyncio.CancelledError() if cancel else None + stream.read.return_value = b"1234" + client = AsyncMock() + client.get_object.return_value = {"Body": stream} + storage._get_client = AsyncMock(return_value=client) + with pytest.raises(asyncio.CancelledError if cancel else StorageReadLimitExceeded): + await storage.read_file("blob", max_bytes=3) + stream.read.assert_awaited_once_with(4) + stream.__aexit__.assert_awaited_once() + + +async def test_azure_range_is_bounded_before_sdk_buffers(): + storage = AzureBlobBackend.__new__(AzureBlobBackend) + storage._resolve = lambda path: path + downloader = AsyncMock() + downloader.readall.return_value = b"1234" + blob = AsyncMock() + blob.download_blob.return_value = downloader + container = Mock() + container.get_blob_client.return_value = blob + storage._get_container_client = AsyncMock(return_value=container) + with pytest.raises(StorageReadLimitExceeded): + await storage.read_file("blob", max_bytes=3) + blob.download_blob.assert_awaited_once_with(offset=0, length=4) + + +async def test_gcs_range_is_bounded_before_sdk_buffers(): + storage = GCSBackend(bucket="synthetic") + blob = Mock() + blob.download_as_bytes.return_value = b"1234" + bucket = Mock() + bucket.blob.return_value = blob + storage._get_bucket = Mock(return_value=bucket) + with pytest.raises(StorageReadLimitExceeded): + await storage.read_file("blob", max_bytes=3) + blob.download_as_bytes.assert_called_once_with(start=0, end=3, raw_download=True) + + +async def test_storage_does_not_disguise_programming_errors(): + storage = S3Backend(bucket="synthetic", region="us-east-1") + storage._get_client = AsyncMock(side_effect=TypeError("programming error")) + with pytest.raises(TypeError): + await storage.read_file("blob", max_bytes=3) diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py index 9996c35ef7..8d6fdd3152 100644 --- a/backend/tests/unit/platform/sources/records/test_google_calendar.py +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -63,7 +63,7 @@ async def test_full_pagination_retains_sparse_cancellation_and_container(): if isinstance(item, CaptureRecord) and item.identity.record_type == "event" ] assert all(item.identity.container_id == "team/a@example.com" for item in events) - assert events[1].kind == "delete" and events[1].payload["recurringEventId"] == "series" + assert events[1].kind == "upsert" and events[1].payload["recurringEventId"] == "series" assert "%2F" in get.calls[1][0] assert get.calls[2][1]["pageToken"] == "page2" assert state.data["calendar_tokens"] == {"team/a@example.com": "boundary"} @@ -120,3 +120,24 @@ async def test_missing_calendar_removes_only_previous_calendar_children(): ("event", "removed", "scope_removed") ] assert state.data["calendar_tokens"] == {"kept": "new"} + + +def test_only_recurring_cancellations_are_retained_and_reinstatement_keeps_identity(): + from airweave.platform.sources.records.google_calendar import record + + cancelled = { + "id": "exception", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"dateTime": "2026-10-01T10:00:00Z"}, + } + exclusion = record("event", cancelled, "calendar") + active = record("event", {**cancelled, "status": "confirmed"}, "calendar") + assert exclusion.kind == active.kind == "upsert" + assert exclusion.identity == active.identity + assert exclusion.payload == cancelled + assert record("event", {"id": "single", "status": "cancelled"}, "calendar").kind == "delete" + with pytest.raises(ValueError, match="original occurrence"): + record( + "event", {"id": "bad", "status": "cancelled", "recurringEventId": "series"}, "calendar" + ) From 2220b8a34471423109eb091ebdf2b46d88488e80 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:08:41 -0700 Subject: [PATCH 008/266] Verify durable Calendar originals with the live capture harness --- backend/tests/live/README.md | 11 +++-- backend/tests/live/canonical_capture.py | 57 +++++++++++++++++++++---- 2 files changed, 56 insertions(+), 12 deletions(-) diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 38c09e96c9..fb890a6ddc 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -1,6 +1,6 @@ # Durable live capture proof -`canonical_capture.py` is opt-in and makes read-only Gmail/Drive provider requests. +`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. @@ -17,7 +17,9 @@ Requirements: 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. Only selected provider account variables are required. +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: @@ -27,7 +29,10 @@ 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 sample without a blob fails. No scope-completion observations +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. diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index e0859b5a77..b5286bce8c 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -28,6 +28,7 @@ POSTGRES_PASSWORD="unused", ) import httpx # noqa: E402 +from pydantic import ValidationError # noqa: E402 from sqlalchemy import text # noqa: E402 from sqlalchemy.engine import make_url # noqa: E402 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine # noqa: E402 @@ -35,7 +36,10 @@ import conftest # noqa: E402,F401 - local test settings must precede Airweave imports from airweave.adapters.storage.filesystem import FilesystemBackend # noqa: E402 from airweave.domains.entities.canonical.query import CanonicalQueryService # noqa: E402 -from airweave.domains.entities.canonical.query_models import RecordListQuery # noqa: E402 +from airweave.domains.entities.canonical.query_models import ( # noqa: E402 + RecordFilters, + RecordListQuery, +) from airweave.domains.entities.canonical.query_store import CanonicalQueryStore # noqa: E402 from airweave.domains.entities.canonical.requests import CaptureBatch, CaptureRecord # noqa: E402 from airweave.domains.entities.canonical.service import CanonicalCaptureService # noqa: E402 @@ -45,10 +49,15 @@ from airweave.domains.storage.file_service import FileService # noqa: E402 from airweave.domains.storage.paths import StoragePaths # noqa: E402 from airweave.models import Organization, Sync, SyncJob # noqa: E402 -from airweave.platform.configs.config import GmailConfig, GoogleDriveConfig # noqa: E402 +from airweave.platform.configs.config import ( # noqa: E402 + GmailConfig, + GoogleCalendarConfig, + GoogleDriveConfig, +) from airweave.platform.http_client.airweave_client import AirweaveHttpClient # noqa: E402 from airweave.platform.http_client.composio_transport import ComposioTransport # noqa: E402 from airweave.platform.sources import GmailSource, GoogleDriveSource # noqa: E402 +from airweave.platform.sources.google_calendar import GoogleCalendarSource # noqa: E402 logging.disable(logging.CRITICAL) PROBE_STAGE = "initialization" @@ -99,6 +108,17 @@ async def fetch_records(name, account, expected_email, key, fence, root): ) profile = await source._get("https://gmail.googleapis.com/gmail/v1/users/me/profile") email = profile.get("emailAddress", "") + elif name == "google_calendar": + source = await GoogleCalendarSource.create( + auth=auth, + logger=logging.getLogger("probe"), + http_client=wrapped, + config=GoogleCalendarConfig(), + ) + profile = await source._get( + "https://www.googleapis.com/calendar/v3/users/me/calendarList/primary" + ) + email = profile.get("id", "") else: source = await GoogleDriveSource.create( auth=auth, @@ -122,9 +142,16 @@ async def fetch_records(name, account, expected_email, key, fence, root): # selected sample does not prove exhaustive source enumeration. if isinstance(item, CaptureRecord): records.append(item) - if (len(records) >= 3 and any(x.blobs for x in records)) or len(records) >= 25: + if len(records) >= 25 or ( + name != "google_calendar" + and len(records) >= 3 + and any(x.blobs for x in records) + ): break - if not records or not any(x.blobs for x in records): + if name == "google_calendar": + if not any(x.identity.record_type == "event" for x in records): + raise ValueError("Bounded Calendar sample did not contain an event") + elif not records or not any(x.blobs for x in records): raise ValueError("Bounded source sample did not contain a real owned blob") return records, storage @@ -151,10 +178,11 @@ async def verify(name, account, email, key, sessions, engine, root): global PROBE_STAGE PROBE_STAGE = name + ":provider_capture" records, _ = await fetch_records(name, account, email, key, fence, root) - PROBE_STAGE = name + ":durable_verification" + PROBE_STAGE = name + ":durable_capture" async with sessions() as db: committed = await service.capture(db, CaptureBatch(fence=fence, records=tuple(records))) await engine.dispose() # Force readback over fresh PostgreSQL connections. + PROBE_STAGE = name + ":durable_readback" storage = FilesystemBackend(root / "blobs") # New filesystem storage instance. expected = {item.identity.entity_key: item for item in records} listed = [] @@ -162,7 +190,10 @@ async def verify(name, account, email, key, sessions, engine, root): while True: async with sessions() as db: page = await query.list_records( - db, organization_id, sync_id, RecordListQuery(limit=2, cursor=cursor) + db, + organization_id, + sync_id, + RecordListQuery(limit=2, cursor=cursor, filters=RecordFilters(state="all")), ) listed.extend(page.records) if not page.has_more: @@ -203,7 +234,7 @@ async def verify(name, account, email, key, sessions, engine, root): "journal_changes": len(changes), "blobs": blob_count, "blob_bytes": total_bytes, - "blob_sha_verified": True, + "blob_sha_verified": True if blob_count else None, "replay_new_changes": 0, "fresh_connections_readback": True, "full_scope_completed": False, @@ -216,13 +247,14 @@ async def main(): url = test_database_url() key = os.environ["COMPOSIO_API_KEY"] selected = os.environ.get("LIVE_PROVIDERS", "gmail,google_drive").split(",") - if not selected or set(selected) - {"gmail", "google_drive"}: - raise ValueError("LIVE_PROVIDERS must select gmail and/or google_drive") + if not selected or set(selected) - {"gmail", "google_drive", "google_calendar"}: + raise ValueError("LIVE_PROVIDERS must select gmail, google_drive and/or google_calendar") inputs = [ (name, os.environ[variable], os.environ["LIVE_EXPECTED_EMAIL"]) for name, variable in ( ("gmail", "LIVE_GMAIL_ACCOUNT_ID"), ("google_drive", "LIVE_DRIVE_ACCOUNT_ID"), + ("google_calendar", "LIVE_CALENDAR_ACCOUNT_ID"), ) if name in selected ] @@ -242,6 +274,8 @@ async def main(): "0000_baseline.py", "0001_canonical_records.py", "0002_projection_publication.py", + "0003_mail_thread_index.py", + "0004_projection_generation.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: @@ -274,6 +308,11 @@ async def main(): except Exception as error: # Never print SQL parameters, native payloads, credentials, or exception tracebacks. details = {"failed": True, "error_type": type(error).__name__, "stage": PROBE_STAGE} + if isinstance(error, ValidationError): + # Types only: validation input/context can contain entire original messages. + details["validation_error_types"] = sorted( + {item["type"] for item in error.errors(include_input=False, include_context=False)} + ) if isinstance(error, httpx.HTTPStatusError): details["http_status"] = error.response.status_code details["host"] = error.request.url.host From 4ec405bebaf22beacb935a97aad0e055c87aaf61 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:21:18 -0700 Subject: [PATCH 009/266] test: hand durable Gmail samples to the Almanac HTTP reader --- backend/tests/live/README.md | 28 ++++ backend/tests/live/almanac_handoff.py | 205 ++++++++++++++++++++++++ backend/tests/live/canonical_capture.py | 16 ++ 3 files changed, 249 insertions(+) create mode 100644 backend/tests/live/almanac_handoff.py diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index fb890a6ddc..75866f7e4a 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -47,3 +47,31 @@ 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. diff --git a/backend/tests/live/almanac_handoff.py b/backend/tests/live/almanac_handoff.py new file mode 100644 index 0000000000..541a5bf16f --- /dev/null +++ b/backend/tests/live/almanac_handoff.py @@ -0,0 +1,205 @@ +"""Optional cross-environment proof through the real owned-service HTTP router.""" + +import asyncio +import base64 +import hashlib +import json +import os +import secrets +import socket +from email.header import decode_header, make_header +from email.message import Message +from email.utils import getaddresses +from pathlib import Path +from types import SimpleNamespace +from uuid import uuid4 + +import uvicorn +from fastapi import FastAPI, HTTPException, Request + +from airweave.api import deps +from airweave.api.v1.api import api_router +from airweave.api.v1.endpoints.records import record_error_response +from airweave.db.session import get_db +from airweave.domains.entities.canonical.store import CanonicalStoreError + + +def digest(value): + return hashlib.sha256(value.encode()).hexdigest() + + +async def expectation(record, storage): + """Fingerprint native data independently before crossing the HTTP/process boundary.""" + payload = record.payload + headers = {h["name"].lower(): h["value"] for h in payload["payload"].get("headers", [])} + bodies = {"text/plain": [], "text/html": []} + body_blobs = 0 + + async def walk(part, path): + nonlocal body_blobs + metadata = Message() + for header in part.get("headers", []): + metadata[header["name"]] = header["value"] + if part.get("filename") or metadata.get_content_disposition() == "attachment": + return + if part.get("mimeType") in bodies: + body = part.get("body", {}) + if "data" in body: + value = body["data"] + raw = base64.urlsafe_b64decode(value + "=" * (-len(value) % 4)) + else: + blob = next((b for b in record.blobs if b.source_path == path + "/body"), None) + raw = await storage.read_file(blob.key, max_bytes=blob.size_bytes) if blob else b"" + body_blobs += int(blob is not None) + bodies[part["mimeType"]].append(raw.decode(metadata.get_content_charset() or "utf-8")) + for index, child in enumerate(part.get("parts", [])): + await walk(child, f"{path}/parts/{index}") + + await walk(payload["payload"], "/payload") + return { + "record_id": str(record.id), + "native_id": record.identity.native_id, + "revision": record.revision, + "observed_at": record.observed_at.isoformat(), + "payload_digest": digest(json.dumps(payload, sort_keys=True, separators=(",", ":"))), + "subject_digest": digest(str(make_header(decode_header(headers.get("subject", ""))))), + "sender_digest": digest( + json.dumps( + [ + [str(make_header(decode_header(name))), address] + for name, address in getaddresses([headers.get("from", "")]) + if address + ] + ) + ), + # Alternative bodies select one representation; mixed bodies may join leaves. + # Exact MIME semantics are separately tested; here every rendered byte must + # correspond to the native captured representation across process/HTTP boundaries. + "text_digests": [ + digest(x) for x in [*bodies["text/plain"], "\n".join(bodies["text/plain"])] + ], + "html_digests": [digest(x) for x in [*bodies["text/html"], "\n".join(bodies["text/html"])]], + "body_blobs": body_blobs, + } + + +async def verify_almanac_reader( + *, records, storage, query, sessions, organization_id, sync_id, root +): + python = Path(os.environ["LIVE_ALMANAC_PYTHON"]).absolute() + repository = Path(os.environ["LIVE_ALMANAC_ROOT"]).resolve() + consumer = repository / "backend/tests/live/read_owned_gmail.py" + if not python.is_file() or not consumer.is_file(): + raise ValueError("Explicit Almanac Python and consumer checkout are required") + candidates = [r for r in records if r.blobs and r.completeness == "complete"] + if not candidates: + raise ValueError("A complete native message with a stored blob is required") + thread_id = candidates[0].payload["threadId"] + selected = [r for r in records if r.payload.get("threadId") == thread_id] + expected = [await expectation(record, storage) for record in selected] + probe_key = secrets.token_urlsafe(32) + app = FastAPI() + app.include_router(api_router) + app.add_exception_handler(CanonicalStoreError, record_error_response) + + async def context(request: Request): + if not secrets.compare_digest(request.headers.get("x-api-key", ""), probe_key): + raise HTTPException(status_code=401) + return SimpleNamespace(organization=SimpleNamespace(id=organization_id)) + + async def database(): + async with sessions() as db: + yield db + + app.dependency_overrides[deps.get_context] = context + app.dependency_overrides[get_db] = database + app.dependency_overrides[deps.get_canonical_query_service] = lambda: query + app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(storage_backend=storage) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(128) + listener.setblocking(False) + server = uvicorn.Server(uvicorn.Config(app, log_config=None, access_log=False, lifespan="off")) + task = asyncio.create_task(server.serve(sockets=[listener])) + process = None + manifest = root / "almanac-handoff.json" + try: + async with asyncio.timeout(20): + while not server.started: + if task.done(): + await task + raise RuntimeError("Local API server did not start") + await asyncio.sleep(0.01) + manifest.write_text( + json.dumps( + { + "origin": f"http://127.0.0.1:{listener.getsockname()[1]}", + "api_key": probe_key, + "organization_id": str(organization_id), + "sync_id": str(sync_id), + "account_id": str(uuid4()), + "source_connection_id": str(uuid4()), + "thread_id": thread_id, + "expected": expected, + } + ) + ) + manifest.chmod(0o600) + # Provider keys/identities do not enter the consumer subprocess environment. + environment = { + "PATH": os.environ.get("PATH", ""), + "PYTHONUNBUFFERED": "1", + "PYTHONPATH": str(repository / "backend") + + os.pathsep + + str(repository / "shared/revtext"), + } + process = await asyncio.create_subprocess_exec( + str(python), + str(consumer), + str(manifest), + env=environment, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + async with asyncio.timeout(120): + output, _ = await process.communicate() + return consumer_result(output, process.returncode) + finally: + try: + await stop_probe(process, server, task) + finally: + listener.close() + manifest.unlink(missing_ok=True) + + +async def stop_probe(process, server, task): + """Reap child/server before the caller destroys their private storage.""" + if process is not None and process.returncode is None: + process.kill() + await process.wait() + server.should_exit = True + try: + async with asyncio.timeout(10): + await task + except TimeoutError: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +def consumer_result(output, returncode): + """Accept only the consumer's structured result; suppress private diagnostics.""" + if not output.strip(): + raise RuntimeError("Consumer interpreter failed before structured verification") + result = json.loads(output) + if returncode or not result.get("verified"): + safe_stage = result.get("stage") + if safe_stage in { + "handoff", + "thread_read", + "exact_record_read", + "blob_read", + "context_negative_check", + }: + print(json.dumps({"consumer_failed_stage": safe_stage}), flush=True) + raise RuntimeError("Almanac consumer proof failed; private diagnostics suppressed") + return result diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index b5286bce8c..d07b3cb690 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -228,7 +228,23 @@ async def verify(name, account, email, key, sessions, engine, root): async with sessions() as db: tail = await query.changes(db, organization_id, sync_id, cursor=change_cursor) assert not tail.changes + consumer = None + if name == "gmail" and os.environ.get("LIVE_ALMANAC_PYTHON"): + from almanac_handoff import verify_almanac_reader + + PROBE_STAGE = name + ":almanac_http_reader" + await engine.dispose() + consumer = await verify_almanac_reader( + records=listed, + storage=storage, + query=query, + sessions=sessions, + organization_id=organization_id, + sync_id=sync_id, + root=root, + ) return { + "almanac_reader": consumer, "provider": name, "records": len(listed), "journal_changes": len(changes), From 6e6e030d0a3d54b3d0e806642e7f6f466c616dd6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:22:24 -0700 Subject: [PATCH 010/266] Add scoped indexed original search with final visibility checks --- backend/airweave/api/v1/endpoints/records.py | 13 + backend/airweave/core/container/container.py | 2 + backend/airweave/core/container/factory.py | 3 + .../canonical/tests/test_owned_search.py | 364 ++++++++++++++++++ .../vector_db/tests/test_partial_coverage.py | 38 ++ .../search/adapters/vector_db/vespa_client.py | 9 +- backend/airweave/domains/search/executor.py | 37 +- .../airweave/domains/search/fakes/executor.py | 1 + backend/airweave/domains/search/owned.py | 337 ++++++++++++++++ .../airweave/domains/search/owned_models.py | 80 ++++ backend/airweave/domains/search/protocols.py | 1 + .../airweave/domains/search/types/results.py | 8 +- backend/airweave/schemas/search_v2.py | 7 +- 13 files changed, 886 insertions(+), 14 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_owned_search.py create mode 100644 backend/airweave/domains/search/adapters/vector_db/tests/test_partial_coverage.py create mode 100644 backend/airweave/domains/search/owned.py create mode 100644 backend/airweave/domains/search/owned_models.py diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index f47059306c..f3f20e90a7 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -23,9 +23,22 @@ ) from airweave.domains.entities.canonical.store import CanonicalStoreError +from airweave.domains.search.owned_models import OwnedSearchRequest, OwnedSearchResponse + router = TrailingSlashRouter() +@router.post("/search", response_model=OwnedSearchResponse) +async def search_records( + request: OwnedSearchRequest, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + container: Container = Depends(deps.get_container), +) -> OwnedSearchResponse: + """Retrieve bounded indexed originals; no provider requests or agent execution.""" + return await container.owned_search.search(db, ctx, request) + + async def record_error_response(request: Request, error: CanonicalStoreError) -> JSONResponse: """Preserve machine-readable recovery codes without exposing SQL/provider diagnostics.""" status = { diff --git a/backend/airweave/core/container/container.py b/backend/airweave/core/container/container.py index 7958740565..a86adde5f4 100644 --- a/backend/airweave/core/container/container.py +++ b/backend/airweave/core/container/container.py @@ -74,6 +74,7 @@ OrganizationServiceProtocol, UserOrganizationRepositoryProtocol, ) +from airweave.domains.search.owned import OwnedSearchService from airweave.domains.search.protocols import ( AgenticSearchServiceProtocol, BrowseServiceProtocol, @@ -251,6 +252,7 @@ async def my_endpoint(event_bus: EventBus = Inject(EventBus)): connect_service: ConnectServiceProtocol # Search domain (v2 tiers) + owned_search: OwnedSearchService instant_search: InstantSearchServiceProtocol classic_search: ClassicSearchServiceProtocol agentic_search: AgenticSearchServiceProtocol diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 8250967620..7485ccfce4 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -113,6 +113,7 @@ from airweave.domains.search.config import SearchConfig from airweave.domains.search.executor import SearchPlanExecutor from airweave.domains.search.instant.service import InstantSearchService +from airweave.domains.search.owned import OwnedSearchService from airweave.domains.source_connections.create import SourceConnectionCreationService from airweave.domains.source_connections.delete import SourceConnectionDeletionService from airweave.domains.source_connections.repository import SourceConnectionRepository @@ -613,6 +614,7 @@ def create_container(settings: Settings) -> Container: organization_service=org_service, user_service=user_service, email_service=email_service, + owned_search=search_deps["owned_search"], instant_search=search_deps["instant_search"], classic_search=search_deps["classic_search"], agentic_search=search_deps["agentic_search"], @@ -1366,6 +1368,7 @@ def _create_search_services( ) return { + "owned_search": OwnedSearchService(executor, source_registry), "instant_search": instant_search, "classic_search": classic_search, "agentic_search": agentic_search, diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py new file mode 100644 index 0000000000..597c5e689e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -0,0 +1,364 @@ +"""Real SQL authority and actual HTTP/executor with fake search/embedding transport.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select, update + +from airweave.api import deps +from airweave.api.v1.endpoints.records import router +from airweave.db.session import get_db +from airweave.domains.embedders.fakes.embedder import FakeSparseEmbedder +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.entities.canonical.tests.test_search_visibility import hit +from airweave.domains.search.adapters.vector_db.fakes.vector_db import FakeVectorDB +from airweave.domains.search.executor import SearchPlanExecutor +from airweave.domains.search.owned import OwnedSearchService +from airweave.domains.search.types import SearchResults +from airweave.models.collection import Collection +from airweave.models.source_connection import SourceConnection +from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata + + +@pytest.fixture +async def indexed(database, source): + service, fence = source + await capture(database, service, fence, observation()) + async with database() as db: + deployment = VectorDbDeploymentMetadata( + dense_embedder="fake", embedding_dimensions=3, sparse_embedder="fake" + ) + db.add(deployment) + await db.flush() + collection = Collection( + name="Test", + readable_id="owned-search", + organization_id=fence.organization_id, + vector_db_deployment_metadata_id=deployment.id, + ) + db.add(collection) + await db.flush() + connection = SourceConnection( + name="Test", + short_name="gmail", + 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() + 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=1, + pipeline_version=1, + generation=generation, + part_index=0, + ) + return fence, locator, connection + + +@pytest.fixture +async def http_search(database, indexed): + fence, locator, connection = indexed + registry = SimpleNamespace( + get=lambda _: SimpleNamespace( + source_class_ref=SimpleNamespace(canonical_record_types=("event", "message")) + ) + ) + vector = FakeVectorDB() + dense = AsyncMock() + dense.embed_many.side_effect = AssertionError("Keyword search must not call dense embedding") + # Any accidental federation discovery fails instead of quietly returning empty. + providers = MagicMock() + providers.get_all_by_collection = AsyncMock(side_effect=AssertionError("No provider discovery")) + executor = SearchPlanExecutor( + dense_embedder=dense, + sparse_embedder=FakeSparseEmbedder(), + vector_db=vector, + sc_repo=providers, + source_registry=registry, + source_lifecycle=MagicMock(), + access_broker=MagicMock(), + ) + executor._discover_federated_sources = AsyncMock(side_effect=AssertionError("No federation")) + service = OwnedSearchService(executor, registry) + ctx = SimpleNamespace(organization=SimpleNamespace(id=fence.organization_id)) + 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_container] = lambda: SimpleNamespace(owned_search=service) + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + yield client, vector, ctx, executor, dense + + +async def test_owned_http_deduplicates_and_preserves_partial_engine(database, indexed, http_search): + fence, locator, connection = indexed + client, vector, _, executor, dense = http_search + first = hit(fence, locator.encode()) + second = first.model_copy( + update={"entity_id": "second-chunk", "textual_representation": "Other evidence"} + ) + vector.seed_results( + SearchResults(results=[first, second], engine_partial=True, engine_coverage_percent=75) + ) + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert len(body["items"]) == 1 + item = body["items"][0] + assert item["identity"] == {"record_type": "event", "native_id": "one", "container_id": None} + assert item["source_connection_id"] == str(connection.id) + assert len(item["excerpts"]) == 2 and "payload" not in item + assert body["engine_partial"] and body["retrieval_incomplete"] + assert body["coverage"] == "bounded_candidates" and "next_cursor" not in body + assert body["sources"][0]["pending_records"] == 0 + executor._discover_federated_sources.assert_not_called() + dense.embed_many.assert_not_called() + compiled = vector._calls[0][1] + assert compiled.filter_groups[0].conditions[0].value == [str(fence.sync_id)] + + +async def test_owned_http_rejects_foreign_unavailable_and_unknown_controls( + database, indexed, http_search +): + fence, _, _ = indexed + client, vector, ctx, _, _ = http_search + base = {"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"} + for invalid in ( + {"sort": "date"}, + {"offset": 10}, + {"limit": 201}, + {"filter": []}, + {"cursor": "next"}, + ): + assert (await client.post("/sync/search", json={**base, **invalid})).status_code == 422 + assert ( + await client.post("/sync/search", json={**base, "record_types": ["unknown"]}) + ).status_code == 422 + assert ( + await client.post("/sync/search", json={**base, "sync_ids": [str(uuid4())]}) + ).status_code == 404 + ctx.organization.id = uuid4() + assert (await client.post("/sync/search", json=base)).status_code == 404 + ctx.organization.id = fence.organization_id + async with database() as db: + await db.execute(update(SourceConnection).values(is_authenticated=False)) + await db.commit() + assert (await client.post("/sync/search", json=base)).status_code == 404 + assert vector._calls == [] + + +async def test_owned_http_stale_text_removed_and_lag_explicit( + database, source, indexed, http_search +): + fence, locator, _ = indexed + client, vector, _, _, _ = http_search + vector.seed_results(SearchResults(results=[hit(fence, locator.encode())])) + service, _ = source + await capture(database, service, fence, observation(payload={"new": "private"})) + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + body = response.json() + assert response.status_code == 200 and body["items"] == [] + assert body["excluded_candidates"] == 1 and body["sources"][0]["pending_records"] == 1 + assert body["retrieval_incomplete"] + + +async def test_owned_filters_are_explicit_postfilters(indexed, http_search): + fence, locator, _ = indexed + client, vector, _, _, _ = http_search + vector.seed_results(SearchResults(results=[hit(fence, locator.encode())])) + response = await client.post( + "/sync/search", + json={ + "query": "budget", + "sync_ids": [str(fence.sync_id)], + "mode": "keyword", + "record_types": ["message"], + }, + ) + assert response.status_code == 200 + assert response.json()["items"] == [] and response.json()["postfilter_excluded"] == 1 + + +async def test_owned_search_rejects_valid_other_account_hit(database, source, indexed, http_search): + from airweave.models import Sync, SyncJob + + fence, _, _ = indexed + client, vector, _, _, _ = http_search + capture_service, _ = source + other_sync, other_job = uuid4(), uuid4() + async with database() as db: + db.add(Sync(id=other_sync, organization_id=fence.organization_id, name="Other")) + await db.flush() + db.add( + SyncJob( + id=other_job, + organization_id=fence.organization_id, + sync_id=other_sync, + status="running", + ) + ) + db.add( + SourceConnection( + name="Other", + short_name="gmail", + organization_id=fence.organization_id, + readable_collection_id="owned-search", + sync_id=other_sync, + is_authenticated=True, + ) + ) + await db.commit() + other_fence = await capture_service.activate_writer( + db, fence.organization_id, other_sync, other_job, attempt_id=uuid4(), attempt_number=1 + ) + await capture(database, capture_service, other_fence, observation()) + store = CanonicalProjectionStore() + async with database() as db: + work = (await store.pending(db, fence.organization_id, other_sync))[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, + ) + # Fake engine deliberately violates its compiled sync filter. SQL enrichment + # must still prevent this otherwise-valid same-collection account leak. + vector.seed_results(SearchResults(results=[hit(other_fence, locator.encode())])) + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + assert response.status_code == 200 + assert response.json()["items"] == [] + assert response.json()["excluded_candidates"] == 1 + + +async def test_disconnect_during_retrieval_rejects_whole_response(database, indexed, http_search): + fence, locator, _ = indexed + client, vector, _, executor, _ = http_search + vector.seed_results(SearchResults(results=[hit(fence, locator.encode())])) + execute = executor.execute + + async def disconnected(**kwargs): + result = await execute(**kwargs) + async with database() as db: + await db.execute(update(SourceConnection).values(is_authenticated=False)) + await db.commit() + return result + + executor.execute = disconnected + 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 + + +async def test_final_gate_drops_early_hit_changed_during_later_collection( + database, source, indexed, http_search +): + from airweave.models import Sync, SyncJob + + fence, locator, _ = indexed + client, vector, _, executor, _ = http_search + capture_service, _ = source + vector.seed_results(SearchResults(results=[hit(fence, locator.encode())])) + other_sync, other_job = uuid4(), uuid4() + async with database() as db: + deployment = await db.scalar(select(VectorDbDeploymentMetadata)) + db.add( + Collection( + name="Other", + readable_id="later-search", + organization_id=fence.organization_id, + vector_db_deployment_metadata_id=deployment.id, + ) + ) + db.add(Sync(id=other_sync, organization_id=fence.organization_id, name="Other")) + await db.flush() + db.add( + SyncJob( + id=other_job, + organization_id=fence.organization_id, + sync_id=other_sync, + status="running", + ) + ) + db.add( + SourceConnection( + name="Other", + short_name="gmail", + organization_id=fence.organization_id, + readable_collection_id="later-search", + sync_id=other_sync, + is_authenticated=True, + ) + ) + await db.commit() + execute = executor.execute + calls = 0 + + async def later_edit(**kwargs): + nonlocal calls + calls += 1 + if calls == 2: + await capture(database, capture_service, fence, observation(payload={"updated": True})) + return SearchResults(results=[]) + return await execute(**kwargs) + + executor.execute = later_edit + # Pin collection order only for this race proof, independently of SQL planner order. + service = next( + value + for key, value in client._transport.app.dependency_overrides.items() + if key is deps.get_container + )().owned_search + resolve = service._resolve_scopes + + async def ordered(*args): + scopes, groups = await resolve(*args) + return scopes, dict(sorted(groups.items(), key=lambda item: item[0][1] != "owned-search")) + + service._resolve_scopes = ordered + response = await client.post( + "/sync/search", + json={ + "query": "budget", + "sync_ids": [str(fence.sync_id), str(other_sync)], + "mode": "keyword", + }, + ) + assert calls == 2 and response.status_code == 200 + assert response.json()["items"] == [] and response.json()["excluded_candidates"] == 1 + assert "Private original text" not in response.text 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/vespa_client.py b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py index b8cef0ba0b..0c7322beec 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, diff --git a/backend/airweave/domains/search/executor.py b/backend/airweave/domains/search/executor.py index 4926362d65..2fd3b0416f 100644 --- a/backend/airweave/domains/search/executor.py +++ b/backend/airweave/domains/search/executor.py @@ -89,6 +89,7 @@ async def execute( ctx: ApiContext, collection_readable_id: str, user_principal: Optional[str] = None, + indexed_only: bool = False, ) -> SearchResults: """Execute the full search pipeline including federated sources.""" # 0. Resolve access control principals @@ -100,7 +101,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 @@ -131,10 +136,18 @@ async def execute( ) ) - retrieved = 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, + 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 [] @@ -142,7 +155,9 @@ async def execute( # 5. If no federated sources, vector DB already has correct limit/offset if not federated_sources: return SearchResults( - results=vector_results, retrieval_incomplete=excluded > 0, + results=vector_results, + retrieval_incomplete=excluded > 0 or engine_results.engine_partial, + **coverage, excluded_candidates=excluded, ) @@ -158,13 +173,17 @@ async def execute( merged = self._merge_with_rrf(vector_results, fed_filtered) return SearchResults( results=merged[original_offset : original_offset + original_limit], - retrieval_incomplete=excluded > 0, excluded_candidates=excluded, + 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], - retrieval_incomplete=excluded > 0, excluded_candidates=excluded, + retrieval_incomplete=excluded > 0 or engine_results.engine_partial, + **coverage, + excluded_candidates=excluded, ) async def _execute_vector_search( @@ -172,7 +191,7 @@ async def _execute_vector_search( plan: SearchPlan, collection_id: str, acl_principals: Optional[list[str]] = None, - ) -> list[SearchResult]: + ) -> SearchResults: """Embed, compile, and execute vector DB search. Adapter exceptions (EmbedderError, VectorDBError) propagate directly @@ -205,7 +224,7 @@ async def _execute_vector_search( 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 diff --git a/backend/airweave/domains/search/fakes/executor.py b/backend/airweave/domains/search/fakes/executor.py index c49226d85e..d3ad4588a5 100644 --- a/backend/airweave/domains/search/fakes/executor.py +++ b/backend/airweave/domains/search/fakes/executor.py @@ -41,6 +41,7 @@ async def execute( ctx: Any = None, collection_readable_id: str = "", user_principal: str | None = None, + indexed_only: bool = False, ) -> 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..0be9a8af2e --- /dev/null +++ b/backend/airweave/domains/search/owned.py @@ -0,0 +1,337 @@ +"""Original-record indexed retrieval using the existing executor and SQL authority.""" + +from collections import defaultdict +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy import and_, func, or_, select, tuple_ +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api.context import ApiContext +from airweave.domains.entities.canonical.projection_models import ProjectionLocator +from airweave.domains.entities.canonical.projection_store import publication_matches +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.store import content_is_available +from airweave.domains.search.owned_models import ( + OwnedSearchCoverage, + OwnedSearchHit, + 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.source_connection import SourceConnection +from airweave.models.sync import Sync + + +class OwnedSearchService: + """No provider lifecycle, planner, answer generator or reranker dependency.""" + + def __init__(self, executor: SearchPlanExecutorProtocol, registry: SourceRegistryProtocol): + """Reuse the executor and registry without another retrieval stack.""" + self._executor = executor + self._registry = registry + + async def search( + self, db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest + ) -> OwnedSearchResponse: + """Resolve exact authorized scopes before any embedding/index request.""" + scopes, groups = await self._resolve_scopes(db, ctx, request) + scope_snapshot = self._scope_identity(scopes) + hits, scores, exclusions, postfiltered = {}, {}, 0, 0 + engine_partial, full = False, False + for (collection_id, readable_id), sync_ids in groups.items(): + results = await self._executor.execute( + plan=SearchPlan( + query=SearchQuery(primary=request.query), + limit=200, + offset=0, + retrieval_strategy=request.mode, + ), + user_filter=[ + FilterGroup( + conditions=[ + FilterCondition( + field="airweave_system_metadata.sync_id", + operator="in", + value=[str(item) for item in sync_ids], + ) + ] + ) + ], + collection_id=str(collection_id), + db=db, + ctx=ctx, + collection_readable_id=readable_id, + indexed_only=True, + ) + engine_partial |= results.engine_partial + exclusions += results.excluded_candidates + full |= len(results.results) + results.excluded_candidates >= 200 + group_hits, group_scores, rejected, filtered = await self._enrich( + db, ctx, request, sync_ids, scopes, results + ) + hits.update(group_hits) + scores.update(group_scores) + exclusions += rejected + postfiltered += filtered + 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") + eligible = await self._final_publications(db, ctx, request, scores, scope_snapshot) + exclusions += len(hits.keys() - eligible) + ranked = sorted(hits.keys() & eligible, key=lambda key: (-scores[key][0], str(key))) + return OwnedSearchResponse( + items=tuple(hits[key] for key in ranked[: request.limit]), + sources=sources, + candidate_window_full=full, + engine_partial=engine_partial, + excluded_candidates=exclusions, + postfilter_excluded=postfiltered, + retrieval_incomplete=engine_partial + or full + or exclusions > 0 + or postfiltered > 0 + or len(ranked) > request.limit + or any(row.pending_records for row in sources), + ) + + async def _resolve_scopes( + self, db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest + ) -> tuple[dict[UUID, SourceConnection], dict[tuple[UUID, str], list[UUID]]]: + rows = ( + await db.execute( + select(SourceConnection, Collection.id) + .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), + ) + .execution_options(populate_existing=True) + ) + ).all() + if len(rows) != len(request.sync_ids) or {row[0].sync_id for row in rows} != set( + request.sync_ids + ): + raise HTTPException(404, "Requested indexed sources are unavailable") + scopes = {row[0].sync_id: row[0] for row in rows} + allowed_types = set() + groups = defaultdict(list) + for connection, collection_id in rows: + types = getattr( + self._registry.get(connection.short_name).source_class_ref, + "canonical_record_types", + (), + ) + if not types: + raise HTTPException(422, "Source does not support owned indexed records") + allowed_types.update(types) + groups[(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, SourceConnection], + results: SearchResults, + ) -> tuple[dict[UUID, OwnedSearchHit], dict[UUID, tuple[float, ProjectionLocator]], int, int]: + locators = [value for result in results.results if (value := self._locator(result))] + if not locators: + return {}, {}, len(results.results), 0 + hits, scores, exclusions, postfiltered = {}, {}, 0, 0 + # Recheck current publication while obtaining native identity. Never enrich + # from an unvalidated cached hit after a concurrent capture/permission change. + records = ( + await db.scalars( + select(Entity) + .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), + or_(*(publication_matches(item) for item in locators)), + ) + .execution_options(populate_existing=True) + ) + ).all() + by_id = {record.id: record for record in records} + 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 + 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, + ) + scores[row.id] = (1 / (60 + rank), locator) + 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, exclusions, postfiltered + + @staticmethod + def _scope_identity(scopes: dict[UUID, SourceConnection]) -> set[tuple]: + return { + (sync, row.id, row.short_name, row.readable_collection_id) + 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() + identities = [ + (locator.record_id, locator.revision, locator.pipeline_version, locator.generation) + for _, locator in scores.values() + ] + # 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), + ), + ) + .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, + ).in_(list(scope_snapshot)), + Entity.deleted_at.is_(None), + content_is_available(), + Entity.indexed_revision == Entity.record_revision, + Entity.indexed_pipeline_version == Sync.index_pipeline_version, + tuple_( + Entity.id, + Entity.record_revision, + Entity.indexed_pipeline_version, + Entity.indexed_generation, + ).in_(identities), + ) + ) + ) + + @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: Entity, 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), + ) + rows = await db.execute( + select(Entity.sync_id, func.count(), func.count().filter(pending)) + .join(Sync, Sync.id == Entity.sync_id) + .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 = {sync: (active, pending) for sync, active, pending in rows} + return tuple( + OwnedSearchCoverage( + sync_id=sync, + active_records=counts.get(sync, (0, 0))[0], + pending_records=counts.get(sync, (0, 0))[1], + ) + 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..e5f479582f --- /dev/null +++ b/backend/airweave/domains/search/owned_models.py @@ -0,0 +1,80 @@ +"""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.requests import RecordIdentity +from airweave.domains.search.types 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 OwnedSearchHit(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, ...] + observed_at: AwareDatetime + source_created_at: AwareDatetime | None + source_updated_at: AwareDatetime | None + completeness: Literal["complete", "partial", "metadata_only"] + + +class OwnedSearchCoverage(BaseModel): + """Counts are source-store state, never total matches for this query.""" + + sync_id: UUID + active_records: int + pending_records: int + + +class OwnedSearchResponse(BaseModel): + """Top results from bounded live candidate windows, without traversal promises.""" + + items: tuple[OwnedSearchHit, ...] + 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..ad83ae8f63 100644 --- a/backend/airweave/domains/search/protocols.py +++ b/backend/airweave/domains/search/protocols.py @@ -42,6 +42,7 @@ async def execute( ctx: ApiContext, collection_readable_id: str, user_principal: Optional[str] = None, + indexed_only: bool = False, ) -> SearchResults: """Execute a search plan and return results.""" ... diff --git a/backend/airweave/domains/search/types/results.py b/backend/airweave/domains/search/types/results.py index d950a08b70..9c6442c608 100644 --- a/backend/airweave/domains/search/types/results.py +++ b/backend/airweave/domains/search/types/results.py @@ -239,13 +239,17 @@ 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." + default=0, + ge=0, + description="Stale or unauthorized index candidates omitted from this page.", ) def __len__(self) -> int: diff --git a/backend/airweave/schemas/search_v2.py b/backend/airweave/schemas/search_v2.py index 3eae466de4..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") @@ -290,6 +290,9 @@ class SearchV2Response(BaseModel): 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." ) From 820b860b03d3431640ec2021e207e7e7f96e9e57 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:24:15 -0700 Subject: [PATCH 011/266] Verify owned source correctness on dev pull requests --- .github/workflows/owned-source-store.yml | 60 ++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 .github/workflows/owned-source-store.yml diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml new file mode 100644 index 0000000000..a65f1ea1ac --- /dev/null +++ b/.github/workflows/owned-source-store.yml @@ -0,0 +1,60 @@ +name: Owned source store + +on: + pull_request: + branches: [dev, main] + paths: + - backend/** + - Makefile + - .github/workflows/owned-source-store.yml + push: + branches: [dev] + paths: + - backend/** + - Makefile + - .github/workflows/owned-source-store.yml + +permissions: + contents: read + +concurrency: + group: owned-store-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + correctness: + runs-on: ubuntu-latest + timeout-minutes: 25 + services: + postgres: + image: postgres:16.15 + env: + POSTGRES_USER: sync_test + POSTGRES_PASSWORD: disposable-ci-only + POSTGRES_DB: sync_tests + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U sync_test -d sync_tests" + --health-interval 5s --health-timeout 5s --health-retries 10 + env: + CANONICAL_TEST_DATABASE_URL: postgresql+asyncpg://sync_test:disposable-ci-only@localhost:5432/sync_tests + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 + with: + python-version: '3.13' + - uses: snok/install-poetry@76e04a911780d5b312d89783f7b1cd627778900a + with: + version: 2.3.2 + virtualenvs-create: true + virtualenvs-in-project: true + - name: Install locked dependencies + working-directory: backend + run: poetry install --with dev,lint --no-root --no-interaction + - name: Verify durable records and publication + run: make test-store + - name: Verify source capture and transport + run: make test-capture + - name: Verify indexed retrieval and visibility + run: make test-search From 8feea1b32e814c51a5f605c67cddda6f643c9e4d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:25:47 -0700 Subject: [PATCH 012/266] Capture bounded provider calendar instances and serve owned ranges --- backend/airweave/api/v1/api.py | 4 + .../api/v1/endpoints/calendar_records.py | 65 ++++ backend/airweave/api/v1/endpoints/records.py | 12 +- .../entities/canonical/calendar_query.py | 349 ++++++++++++++++++ .../domains/entities/canonical/checkpoint.py | 13 + .../entities/canonical/projection_mappers.py | 5 +- .../entities/canonical/projection_store.py | 3 +- .../domains/entities/canonical/store.py | 12 +- .../tests/test_calendar_exclusions.py | 40 ++ .../canonical/tests/test_calendar_range.py | 256 +++++++++++++ backend/airweave/platform/configs/config.py | 36 +- .../platform/cursors/google_calendar.py | 20 +- .../platform/sources/google_calendar.py | 11 +- .../sources/records/google_calendar.py | 129 +++++-- .../sources/records/test_google_calendar.py | 105 +++++- 15 files changed, 1016 insertions(+), 44 deletions(-) create mode 100644 backend/airweave/api/v1/endpoints/calendar_records.py create mode 100644 backend/airweave/domains/entities/canonical/calendar_query.py create mode 100644 backend/airweave/domains/entities/canonical/checkpoint.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_calendar_range.py diff --git a/backend/airweave/api/v1/api.py b/backend/airweave/api/v1/api.py index 0d1f5ae7eb..14461ec3d3 100644 --- a/backend/airweave/api/v1/api.py +++ b/backend/airweave/api/v1/api.py @@ -7,6 +7,7 @@ auth_providers, billing, browse_tree, + calendar_records, collections, connect, entities, @@ -61,3 +62,6 @@ ) api_router.include_router(webhooks.router, prefix="/webhooks", tags=["webhooks"]) api_router.include_router(connect.router, prefix="/connect", tags=["connect"]) + + +api_router.include_router(calendar_records.router, prefix="/sync", tags=["records"]) diff --git a/backend/airweave/api/v1/endpoints/calendar_records.py b/backend/airweave/api/v1/endpoints/calendar_records.py new file mode 100644 index 0000000000..5f11715b8d --- /dev/null +++ b/backend/airweave/api/v1/endpoints/calendar_records.py @@ -0,0 +1,65 @@ +"""Calendar range reads use captured provider instances, never live provider requests.""" + +from uuid import UUID +from zoneinfo import ZoneInfoNotFoundError + +from fastapi import Depends, HTTPException, Query +from pydantic import AwareDatetime, ValidationError +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api import deps +from airweave.api.context import ApiContext +from airweave.api.router import TrailingSlashRouter +from airweave.core.config import settings +from airweave.db.session import get_db +from airweave.domains.entities.canonical.calendar_query import ( + CalendarRange, + CalendarRangePage, + CalendarRangeService, +) +from airweave.domains.entities.canonical.store import SourceNotFound +from airweave.models.source_connection import SourceConnection + +router = TrailingSlashRouter() + + +async def _authorized_connection(db: AsyncSession, organization_id: UUID, sync_id: UUID): + identity = await db.scalar( + select(SourceConnection.id).where( + SourceConnection.organization_id == organization_id, + SourceConnection.sync_id == sync_id, + SourceConnection.short_name == "google_calendar", + SourceConnection.is_authenticated.is_(True), + ) + ) + if identity is None: + raise SourceNotFound("Calendar source is not available") + return identity + + +@router.get("/{sync_id}/calendar/events", response_model=CalendarRangePage) +async def calendar_events( + sync_id: UUID, + calendar_id: str = Query(min_length=1), + start: AwareDatetime = Query(), + end: AwareDatetime = Query(), + timezone: str = "UTC", + limit: int = Query(100, ge=1, le=250), + cursor: str | None = None, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), +) -> CalendarRangePage: + """List observed occurrences; requests beyond coverage include an explicit sync action.""" + try: + query = CalendarRange(start=start, end=end, timezone=timezone, limit=limit, cursor=cursor) + except (ValidationError, ZoneInfoNotFoundError) as exc: + raise HTTPException(422, "Invalid timezone or window; reads are limited to31days") from exc + connection_id = await _authorized_connection(db, ctx.organization.id, sync_id) + result = await CalendarRangeService(settings.STATE_SECRET).read( + db, ctx.organization.id, sync_id, calendar_id, query + ) + + if await _authorized_connection(db, ctx.organization.id, sync_id) != connection_id: + raise SourceNotFound("Calendar source changed while reading") + return result diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index f3f20e90a7..42d1811830 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -12,6 +12,7 @@ from airweave.api.router import TrailingSlashRouter from airweave.core.container import Container from airweave.db.session import get_db +from airweave.domains.entities.canonical.calendar_query import CalendarRangeNotCaptured from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query import CanonicalQueryService from airweave.domains.entities.canonical.query_models import ( @@ -22,7 +23,6 @@ RecordPage, ) from airweave.domains.entities.canonical.store import CanonicalStoreError - from airweave.domains.search.owned_models import OwnedSearchRequest, OwnedSearchResponse router = TrailingSlashRouter() @@ -47,11 +47,13 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> "blob_not_found": 404, "stale_record_revision": 409, "blob_unavailable": 503, + "calendar_changed_restart": 409, + "calendar_range_not_captured": 409, }.get(error.code, 400) - return JSONResponse( - status_code=status, - content={"error": {"code": error.code, "message": str(error), "retryable": status == 503}}, - ) + detail = {"code": error.code, "message": str(error), "retryable": status == 503} + if isinstance(error, CalendarRangeNotCaptured): + detail["action"] = error.action + return JSONResponse(status_code=status, content={"error": detail}) @router.get("/{sync_id}/records", response_model=RecordPage) diff --git a/backend/airweave/domains/entities/canonical/calendar_query.py b/backend/airweave/domains/entities/canonical/calendar_query.py new file mode 100644 index 0000000000..abde4bd7cd --- /dev/null +++ b/backend/airweave/domains/entities/canonical/calendar_query.py @@ -0,0 +1,349 @@ +"""Bounded, store-only calendar ranges over provider-expanded original occurrences.""" + +import json +from datetime import datetime, time, timedelta, timezone +from typing import Literal +from uuid import UUID +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +from jose import JWTError, jwt +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError, model_validator +from sqlalchemy import String, and_, cast, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.query import InvalidRecordCursor, RecordNotFound +from airweave.domains.entities.canonical.store import ( + CanonicalStoreError, + SourceNotFound, + content_is_available, + source_record, +) +from airweave.models import Entity, Sync, SyncCursor, SyncJob +from airweave.platform.cursors.google_calendar import CalendarWindowCoverage, GoogleCalendarCursor + + +class CalendarRangeError(CanonicalStoreError): + """Calendar reads cannot silently fetch providers or truncate their input.""" + + code = "calendar_range_budget_exceeded" + + +class CalendarRangeNotCaptured(CanonicalStoreError): + """Ask for an explicit configured capture before reading this interval.""" + + code = "calendar_range_not_captured" + + def __init__(self, start: datetime, end: datetime): + """Carry a machine-readable explicit sync request, without executing it.""" + super().__init__( + "Requested range is not captured; configure occurrence_window and run sync" + ) + self.action = { + "operation": "sync_calendar_window", + "config": {"occurrence_window": {"start": start.isoformat(), "end": end.isoformat()}}, + } + + +class CalendarChanged(CanonicalStoreError): + """A continuation cannot mix changed local records or coverage generations.""" + + code = "calendar_changed_restart" + + +class CalendarRange(BaseModel): + """Fixed request window and ordering; only cursor position varies between pages.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + start: AwareDatetime + end: AwareDatetime + timezone: str + limit: int = Field(default=100, ge=1, le=250) + cursor: str | None = None + + @model_validator(mode="after") + def validate_window(self): + """Validate timezone and bounded half-open interval.""" + if not timedelta(0) < self.end - self.start <= timedelta(days=31): + raise ValueError("Calendar read window must be positive and at most31days") + try: + ZoneInfo(self.timezone) + except ZoneInfoNotFoundError as exc: + raise ValueError("Unknown calendar timezone") from exc + return self + + +class CalendarRangeCursor(BaseModel): + """Signed scope, exact query, checkpoint and stable position.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + purpose: Literal["canonical_calendar_range_v1"] = "canonical_calendar_range_v1" + organization_id: UUID + sync_id: UUID + calendar_id: str + start: AwareDatetime + end: AwareDatetime + timezone: str + sequence: int + coverage_scan_id: UUID + after_start: AwareDatetime + after_id: UUID + + +class CalendarRangePage(BaseModel): + """Current stored occurrences alongside their last successful coverage evidence.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + records: tuple[SourceRecord, ...] + next_cursor: str | None + has_more: bool + coverage: CalendarWindowCoverage + consistency: Literal["observed"] = "observed" + refresh_state: Literal["last_completed", "refresh_in_progress", "refresh_incomplete"] + observed_change_sequence: int + scanned_records: int + + +def occurrence_interval(payload: dict, calendar_timezone: str) -> tuple[datetime, datetime]: + """All-day dates remain native; only comparison uses their calendar timezone.""" + start, end = payload.get("start", {}), payload.get("end", {}) + zone = ZoneInfo(calendar_timezone) + if start.get("date") and end.get("date"): + from datetime import date + + return ( + datetime.combine(date.fromisoformat(start["date"]), time(), zone), + datetime.combine(date.fromisoformat(end["date"]), time(), zone), + ) + values = [] + for part in (start, end): + value = datetime.fromisoformat(part["dateTime"].replace("Z", "+00:00")) + if value.tzinfo is None: + value = value.replace(tzinfo=ZoneInfo(part.get("timeZone") or calendar_timezone)) + values.append(value) + return values[0], values[1] + + +class CalendarRangeService: + """Read current owned rows under explicit input bounds and optimistic consistency.""" + + def __init__(self, signing_key: str): + """Use the deployment cursor signing key.""" + self.signing_key = signing_key + + async def read( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + calendar_id: str, + query: CalendarRange, + ) -> CalendarRangePage: + """Return a bounded observed range, failing on changed continuation state.""" + sync = await db.scalar( + select(Sync).where(Sync.id == sync_id, Sync.organization_id == organization_id) + ) + if sync is None: + raise SourceNotFound("Source does not exist in this organization") + calendar = await db.scalar( + select(Entity).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.entity_definition_short_name == "calendar", + Entity.native_id == calendar_id, + Entity.deleted_at.is_(None), + content_is_available(), + ) + ) + if calendar is None: + raise RecordNotFound("Calendar is not currently available") + stored_cursor = await db.scalar( + select(SyncCursor).where( + SyncCursor.sync_id == sync_id, SyncCursor.organization_id == organization_id + ) + ) + state = GoogleCalendarCursor.model_validate( + stored_cursor.cursor_data if stored_cursor else {} + ) + coverage = state.occurrence_coverage.get(calendar_id) + stamp_data = ( + stored_cursor.cursor_data.get("canonical_checkpoint") if stored_cursor else None + ) + if ( + coverage is None + or stamp_data is None + or query.start < coverage.start + or query.end > coverage.end + ): + raise CalendarRangeNotCaptured(query.start, query.end) + stamp = CanonicalCheckpoint.model_validate(stamp_data) + if calendar.source_payload.get("timeZone") != coverage.timezone: + raise CalendarChanged("Calendar timezone changed; refresh captured window") + sequence = sync.observed_change_sequence + position = self._position(query, organization_id, sync_id, calendar_id, coverage, sequence) + rows = await self._bounded_rows(db, organization_id, sync_id, calendar_id) + matches = self._matching_rows(rows, coverage, query, position) + # Recheck both journal and checkpoint: no snapshot or mixed-page claim. + db.expire_all() + latest = await db.scalar( + select(Sync).where(Sync.id == sync_id, Sync.organization_id == organization_id) + ) + latest_cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == sync_id)) + latest_scan = ( + latest_cursor.cursor_data.get("occurrence_coverage", {}) + .get(calendar_id, {}) + .get("scan_id") + if latest_cursor + else None + ) + if ( + latest is None + or latest.observed_change_sequence != sequence + or latest_scan != str(coverage.scan_id) + ): + raise CalendarChanged("Calendar changed while reading; restart range query") + refresh_state = "last_completed" + if ( + latest.writer_attempt_id != stamp.writer_attempt_id + or sequence != stamp.observed_change_sequence + ): + job = await db.get(SyncJob, latest.writer_job_id) if latest.writer_job_id else None + refresh_state = ( + "refresh_in_progress" + if job + and str(getattr(job.status, "value", job.status)).lower() in {"running", "pending"} + else "refresh_incomplete" + ) + more = len(matches) > query.limit + page = matches[: query.limit] + next_cursor = None + if more: + next_cursor = jwt.encode( + CalendarRangeCursor( + organization_id=organization_id, + sync_id=sync_id, + calendar_id=calendar_id, + start=query.start, + end=query.end, + timezone=query.timezone, + sequence=sequence, + coverage_scan_id=coverage.scan_id, + after_start=page[-1][0], + after_id=page[-1][1], + ).model_dump(mode="json"), + self.signing_key, + algorithm="HS256", + ) + return CalendarRangePage( + records=tuple(item[2] for item in page), + next_cursor=next_cursor, + has_more=more, + coverage=coverage, + refresh_state=refresh_state, + observed_change_sequence=sequence, + scanned_records=len(rows), + ) + + @staticmethod + async def _bounded_rows( + db: AsyncSession, organization_id: UUID, sync_id: UUID, calendar_id: str + ) -> list[Entity]: + """Gate payload transfer and count in one PostgreSQL statement snapshot.""" + filters = ( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.record_revision > 0, + Entity.entity_definition_short_name == "event_occurrence", + Entity.container_id == calendar_id, + Entity.deleted_at.is_(None), + content_is_available(), + ) + budget = ( + select( + func.count(Entity.id).label("records"), + func.coalesce( + func.sum(func.octet_length(cast(Entity.source_payload, String))), 0 + ).label("bytes"), + ) + .where(*filters) + .cte("occurrence_budget") + ) + # Outer join always returns the budget, even for an empty or rejected scope. + # An oversized single JSON value never travels to the application. + result = await db.execute( + select(budget.c.records, budget.c.bytes, Entity) + .select_from(budget) + .outerjoin( + Entity, + and_(*filters, budget.c.records <= 10000, budget.c.bytes <= 20 * 1024 * 1024), + ) + .order_by(Entity.id) + ) + data = result.all() + if data[0][0] > 10000 or data[0][1] > 20 * 1024 * 1024: + raise CalendarRangeError( + "Stored occurrence input exceeds budget; narrow capture horizon" + ) + return [row[2] for row in data if row[2] is not None] + + @staticmethod + def _matching_rows( + rows: list[Entity], + coverage: CalendarWindowCoverage, + query: CalendarRange, + position: CalendarRangeCursor | None, + ) -> list[tuple[datetime, UUID, SourceRecord]]: + """Apply schedule semantics only after bounding the calendar's stored inputs.""" + bytes_read = 0 + matches = [] + for row in rows: + bytes_read += len(json.dumps(row.source_payload).encode()) + if bytes_read > 20 * 1024 * 1024: + raise CalendarRangeError("Stored occurrence input exceeds byte budget") + if row.source_payload.get("status") == "cancelled": + continue + try: + start, end = occurrence_interval(row.source_payload, coverage.timezone) + if end <= start: + raise ValueError("Invalid interval") + except (KeyError, ValueError, TypeError) as exc: + raise CalendarRangeError( + "Captured occurrence has an invalid schedule; refresh source" + ) from exc + start = start.astimezone(timezone.utc) + if end > query.start and start < query.end: + if position is None or (start, row.id) > (position.after_start, position.after_id): + matches.append((start, row.id, source_record(row))) + matches.sort(key=lambda item: (item[0], item[1])) + return matches + + def _position( + self, + query: CalendarRange, + organization_id: UUID, + sync_id: UUID, + calendar_id: str, + coverage: CalendarWindowCoverage, + sequence: int, + ) -> CalendarRangeCursor | None: + if query.cursor is None: + return None + try: + position = CalendarRangeCursor.model_validate( + jwt.decode(query.cursor, self.signing_key, algorithms=["HS256"]) + ) + except (JWTError, ValidationError, ValueError) as exc: + raise InvalidRecordCursor("Invalid Calendar cursor; restart range query") from exc + if position.purpose != "canonical_calendar_range_v1" or ( + position.organization_id, + position.sync_id, + position.calendar_id, + position.start, + position.end, + position.timezone, + ) != (organization_id, sync_id, calendar_id, query.start, query.end, query.timezone): + raise InvalidRecordCursor("Calendar cursor belongs to another scope or range") + if position.sequence != sequence or position.coverage_scan_id != coverage.scan_id: + raise CalendarChanged("Calendar changed since prior page; restart range query") + return position diff --git a/backend/airweave/domains/entities/canonical/checkpoint.py b/backend/airweave/domains/entities/canonical/checkpoint.py new file mode 100644 index 0000000000..d67f46f8be --- /dev/null +++ b/backend/airweave/domains/entities/canonical/checkpoint.py @@ -0,0 +1,13 @@ +"""Capture-owned checkpoint metadata, stamped only at the fenced commit boundary.""" + +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + + +class CanonicalCheckpoint(BaseModel): + """Bind successful source state to its actual committed run and journal position.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + writer_attempt_id: UUID + observed_change_sequence: int = Field(ge=0) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index d1993b05ee..686c009cf4 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -201,8 +201,9 @@ def excluded_from_search(record: SourceRecord, source_name: str) -> bool: """Retained calendar exclusions intentionally publish no searchable meeting.""" return ( source_name == "google_calendar" - and record.identity.record_type == "event" - and is_cancelled_recurring_event(record.payload) + and (record.identity.record_type == "event_occurrence" or ( + record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload) + )) ) diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index c115ee3589..f7dd2bff4d 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -183,7 +183,8 @@ async def publish( exclusion = work.record.identity.record_type == "event" and is_cancelled_recurring_event( work.record.payload ) - no_content = work.record.deleted_at is not None or exclusion + no_content = (work.record.deleted_at is not None or exclusion + or work.record.identity.record_type == "event_occurrence") if chunk_count < 0 or (not no_content and chunk_count == 0): raise ValueError("Active records require a nonempty complete projection") if no_content and chunk_count != 0: diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 6d913b9f53..93014e38f7 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import aliased +from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint from airweave.domains.entities.canonical.models import ( CaptureResult, ChangePage, @@ -493,7 +494,7 @@ async def save_checkpoint( self, db: AsyncSession, fence: WriterFence, cursor_data: dict ) -> None: """Call only after capture barrier and successful exact-scope reconciliation.""" - await self._fenced_sync(db, fence) + sync = await self._fenced_sync(db, fence) cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) if cursor is None: cursor = SyncCursor( @@ -501,7 +502,14 @@ async def save_checkpoint( sync_id=fence.sync_id, ) db.add(cursor) - cursor.cursor_data = cursor_data + # Source-supplied or previously loaded values cannot overwrite this reserved stamp. + cursor.cursor_data = { + **cursor_data, + "canonical_checkpoint": CanonicalCheckpoint( + writer_attempt_id=fence.attempt_id, + observed_change_sequence=sync.observed_change_sequence, + ).model_dump(mode="json"), + } cursor.last_updated = datetime.now(timezone.utc) await db.flush() diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index cf37578f57..302d46cde4 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -145,3 +145,43 @@ async def test_live_event_cannot_publish_empty_or_exclusion_nonempty(database, s ) with pytest.raises(ValueError, match="nonempty"): await store.publish(db, active, uuid4(), 0) + + +async def test_operational_occurrences_publish_zero_documents_without_retry(database, source): + capture, fence = source + parent = record("calendar", {"id": "cal"}) + occurrence = record( + "event_occurrence", + { + "id": "real-instance", + "status": "confirmed", + "start": {"date": "2026-10-01"}, + "end": {"date": "2026-10-02"}, + }, + "cal", + ).model_copy(update={"parent": parent.identity}) + async with database() as db: + result = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, occurrence))) + row = result.changes[-1].record + projections = CanonicalProjectionStore() + async with database() as db: + work = next( + w + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + if w.record.id == row.id + ) + processor = MagicMock() + processor.process = AsyncMock() + destination = MagicMock() + destination.collection_id = uuid4() + destination.feed_prepared = AsyncMock() + projector = CanonicalProjector(projections, database, processor, AsyncMock()) + assert await projector.project_one(work, "google_calendar", destination, MagicMock()) + processor.process.assert_not_awaited() + destination.feed_prepared.assert_not_awaited() + async with database() as db: + assert all( + w.record.id != row.id + for w in await projections.pending(db, fence.organization_id, fence.sync_id) + ) + assert (await db.get(Entity, row.id)).indexed_chunk_count == 0 diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py new file mode 100644 index 0000000000..680b554313 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py @@ -0,0 +1,256 @@ +"""Stored provider instances support bounded ranges without recurrence guessing.""" + +from datetime import datetime, timezone +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastapi import FastAPI +from httpx import ASGITransport, AsyncClient +from sqlalchemy import select + +from airweave.api import deps +from airweave.api.v1.endpoints.calendar_records import router +from airweave.api.v1.endpoints.records import record_error_response +from airweave.db.session import get_db +from airweave.domains.entities.canonical.calendar_query import ( + CalendarChanged, + CalendarRange, + CalendarRangeNotCaptured, + CalendarRangeService, +) +from airweave.domains.entities.canonical.requests import CaptureBatch +from airweave.domains.entities.canonical.store import CanonicalStoreError +from airweave.models import SyncCursor +from airweave.models.source_connection import SourceConnection +from airweave.platform.cursors.google_calendar import CalendarWindowCoverage +from airweave.platform.sources.records.google_calendar import record + + +def window(): + return CalendarRange( + start="2026-03-08T00:00:00-08:00", + end="2026-03-10T00:00:00-07:00", + timezone="America/Los_Angeles", + limit=1, + ) + + +async def seed(database, source): + capture, fence = source + parent = record("calendar", {"id": "cal", "timeZone": "America/Los_Angeles"}) + events = [ + {"id": "all-day", "start": {"date": "2026-03-08"}, "end": {"date": "2026-03-09"}}, + { + "id": "real-instance-id", + "recurringEventId": "master", + "originalStartTime": {"dateTime": "2026-03-01T09:00:00-08:00"}, + "start": {"dateTime": "2026-03-09T09:00:00-07:00"}, + "end": {"dateTime": "2026-03-09T10:00:00-07:00"}, + }, + {"id": "ends-at-boundary", "start": {"date": "2026-03-07"}, "end": {"date": "2026-03-08"}}, + ] + observations = ( + parent, + *( + record("event_occurrence", item, "cal").model_copy(update={"parent": parent.identity}) + for item in events + ), + ) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=observations)) + coverage = CalendarWindowCoverage( + start="2026-03-01T00:00:00Z", + end="2026-04-01T00:00:00Z", + timezone="America/Los_Angeles", + completed_at=datetime.now(timezone.utc), + scan_id=uuid4(), + ) + async with database() as db: + await capture.save_checkpoint( + db, + fence, + { + "occurrence_coverage": {"cal": coverage.model_dump(mode="json")}, + "canonical_checkpoint": { + "writer_attempt_id": str(uuid4()), + "observed_change_sequence": 999999, + }, + }, + ) + return parent, coverage + + +async def test_paged_range_all_day_moved_instance_identity_and_changes(database, source): + parent, coverage = await seed(database, source) + capture, fence = source + service = CalendarRangeService("secret") + async with database() as db: + first = await service.read(db, fence.organization_id, fence.sync_id, "cal", window()) + assert first.records[0].identity.native_id == "all-day" + assert first.has_more and first.refresh_state == "last_completed" + next_query = window().model_copy(update={"cursor": first.next_cursor}) + second = await service.read(db, fence.organization_id, fence.sync_id, "cal", next_query) + assert second.records[0].identity.native_id == "real-instance-id" and not second.has_more + saved = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + assert saved.cursor_data["canonical_checkpoint"]["writer_attempt_id"] == str( + fence.attempt_id + ) + assert saved.cursor_data["canonical_checkpoint"]["observed_change_sequence"] == 4 + async with database() as db: + await capture.capture( + db, + CaptureBatch( + fence=fence, + records=( + record( + "event_occurrence", + { + "id": "late", + "start": {"date": "2026-03-09"}, + "end": {"date": "2026-03-10"}, + }, + "cal", + ), + ), + ), + ) + async with database() as db: + with pytest.raises(CalendarChanged): + await service.read(db, fence.organization_id, fence.sync_id, "cal", next_query) + current = await service.read(db, fence.organization_id, fence.sync_id, "cal", window()) + assert current.coverage == coverage and current.refresh_state != "last_completed" + await capture.capture( + db, + CaptureBatch( + fence=fence, + records=( + parent.model_copy( + update={"kind": "delete", "removal_reason": "access_revoked"} + ), + ), + ), + ) + async with database() as db: + with pytest.raises(CanonicalStoreError): + await service.read(db, fence.organization_id, fence.sync_id, "cal", window()) + + +async def test_http_missing_coverage_action_and_cross_org(database, source, monkeypatch): + capture, fence = source + async with database() as db: + await capture.capture( + db, + CaptureBatch( + fence=fence, records=(record("calendar", {"id": "cal", "timeZone": "UTC"}),) + ), + ) + async with database() as db: + db.add( + SourceConnection( + organization_id=fence.organization_id, + sync_id=fence.sync_id, + name="Synthetic Calendar", + short_name="google_calendar", + is_authenticated=True, + ) + ) + await db.commit() + app = FastAPI() + app.include_router(router, prefix="/sync") + app.add_exception_handler(CanonicalStoreError, record_error_response) + owner = fence.organization_id + + 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 + params = {"calendar_id": "cal", "start": "2026-03-01T00:00:00Z", "end": "2026-03-02T00:00:00Z"} + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + path = f"/sync/{fence.sync_id}/calendar/events" + missing = await client.get(path, params=params) + assert missing.status_code == 409 + assert missing.json()["error"]["action"]["operation"] == "sync_calendar_window" + owner = uuid4() + assert (await client.get(path, params=params)).status_code == 404 + owner = fence.organization_id + assert ( + await client.get(path, params={**params, "timezone": "not/a/zone"}) + ).status_code == 422 + await seed(database, source) + original_read = CalendarRangeService.read + + async def revoke_during_read(self, *args, **kwargs): + from sqlalchemy import update + + result = await original_read(self, *args, **kwargs) + async with database() as other: + await other.execute( + update(SourceConnection) + .where(SourceConnection.sync_id == fence.sync_id) + .values(is_authenticated=False) + ) + await other.commit() + return result + + monkeypatch.setattr(CalendarRangeService, "read", revoke_during_read) + assert (await client.get(path, params=params)).status_code == 404 + + +async def test_request_outside_successful_coverage_never_returns_empty_success(database, source): + await seed(database, source) + _, fence = source + async with database() as db: + with pytest.raises(CalendarRangeNotCaptured): + await CalendarRangeService("secret").read( + db, + fence.organization_id, + fence.sync_id, + "cal", + CalendarRange( + start="2001-01-01T00:00:00Z", end="2001-01-02T00:00:00Z", timezone="UTC" + ), + ) + + +async def test_empty_observed_range_is_distinct_from_missing_coverage(database, source): + await seed(database, source) + _, fence = source + async with database() as db: + page = await CalendarRangeService("secret").read( + db, + fence.organization_id, + fence.sync_id, + "cal", + CalendarRange(start="2026-03-20T00:00:00Z", end="2026-03-21T00:00:00Z", timezone="UTC"), + ) + assert page.records == () and not page.has_more + assert page.coverage and page.refresh_state == "last_completed" + + +async def test_budget_rejects_oversized_row_before_orm_materialization(database, source): + from sqlalchemy import update + + from airweave.domains.entities.canonical.calendar_query import CalendarRangeError + from airweave.models import Entity + + await seed(database, source) + _, fence = source + async with database() as db: + await db.execute( + update(Entity) + .where(Entity.sync_id == fence.sync_id, Entity.native_id == "all-day") + .values(source_payload={"oversized": "x" * (21 * 1024 * 1024)}) + ) + await db.commit() + async with database() as db: + with pytest.raises(CalendarRangeError): + await CalendarRangeService._bounded_rows( + db, fence.organization_id, fence.sync_id, "cal" + ) + assert not any(isinstance(item, Entity) for item in db.identity_map.values()) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 6b3438d427..07614b673d 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -1,9 +1,9 @@ """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 pydantic import AwareDatetime, Field, field_validator, model_validator from airweave.platform.configs._base import BaseConfig, RequiredTemplateConfig from airweave.platform.utils.ssrf import validate_host, validate_url @@ -307,10 +307,36 @@ def validate_date_format(cls, value): return value.replace("-", "/") -class GoogleCalendarConfig(SourceConfig): - """Google Calendar configuration schema.""" +class CalendarOccurrenceWindow(BaseConfig): + """Explicit operator-requested capture window; dates must carry timezones.""" - pass + 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): + """Rolling expanded capture, or an explicit window through source configuration.""" + + 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): diff --git a/backend/airweave/platform/cursors/google_calendar.py b/backend/airweave/platform/cursors/google_calendar.py index 45e1e104ba..89f78db7ae 100644 --- a/backend/airweave/platform/cursors/google_calendar.py +++ b/backend/airweave/platform/cursors/google_calendar.py @@ -1,11 +1,25 @@ -"""Independent continuation tokens for each calendar's unexpanded event stream.""" +"""Calendar delta tokens and successfully observed expanded-window coverage.""" -from pydantic import Field +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): - """Token parameters are fixed by the canonical generator, not user date ranges.""" + """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/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index ac13bb9709..3a570cc84d 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -96,8 +96,8 @@ class GoogleCalendarSource(BaseSource): Google Calendar scheduling information for productivity and time management insights. """ - canonical_record_types = ("calendar", "event") - canonical_container_parents = {"event": "calendar"} + canonical_record_types = ("calendar", "event", "event_occurrence") + canonical_container_parents = {"event": "calendar", "event_occurrence": "calendar"} async def generate_observations( self, @@ -106,10 +106,12 @@ async def generate_observations( files: FileService | None = None, node_selections: list[NodeSelectionData] | None = None, ) -> AsyncGenerator[SourceObservation, None]: - """Capture original masters, exceptions and cancellations; no recurrence expansion.""" + """Capture originals and bounded provider-expanded instances with real IDs.""" if node_selections: raise ValueError("Calendar selection scopes are not supported by canonical capture") - async for observation in generate_calendar_observations(self._get, cursor): + async for observation in generate_calendar_observations( + self._get, cursor, self.calendar_config.resolved_window() + ): yield observation # ----------------------- @@ -127,6 +129,7 @@ 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)) diff --git a/backend/airweave/platform/sources/records/google_calendar.py b/backend/airweave/platform/sources/records/google_calendar.py index f3abb57f47..3123d150ba 100644 --- a/backend/airweave/platform/sources/records/google_calendar.py +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -1,8 +1,10 @@ """Original Calendar resources with per-calendar incremental synchronization.""" +import json from collections.abc import AsyncGenerator, Awaitable, Callable from datetime import datetime, timezone from urllib.parse import quote +from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, JsonValue @@ -20,7 +22,8 @@ SourceGoneError, ) from airweave.domains.syncs.cursors.cursor import SyncCursor -from airweave.platform.cursors.google_calendar import GoogleCalendarCursor +from airweave.platform.configs.config import CalendarOccurrenceWindow +from airweave.platform.cursors.google_calendar import CalendarWindowCoverage, GoogleCalendarCursor GetJSON = Callable[..., Awaitable[dict]] BASE = "https://www.googleapis.com/calendar/v3" @@ -113,11 +116,12 @@ async def capture_calendar( async def generate_calendar_observations( - get: GetJSON, cursor: SyncCursor | None + get: GetJSON, cursor: SyncCursor | None, window: CalendarOccurrenceWindow | None = None ) -> AsyncGenerator[SourceObservation, None]: """Refresh calendar membership; capture each accessible calendar independently.""" previous = GoogleCalendarCursor.model_validate(cursor.data if cursor else {}) tokens: dict[str, str] = {} + coverage: dict[str, CalendarWindowCoverage] = {} seen_calendars: set[str] = set() seen_pages: set[str] = set() params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"} @@ -133,25 +137,10 @@ async def generate_calendar_observations( raise ValueError("Calendar list repeated an identity across pages") seen_calendars.add(calendar_id) yield calendar_record - try: - async for observation in capture_calendar( - get, calendar_id, previous.calendar_tokens.get(calendar_id), tokens - ): - yield observation - except SourceEntityNotFoundError: - yield CaptureRecord( - identity=RecordIdentity(record_type="calendar", native_id=calendar_id), - payload={"id": calendar_id}, - kind="delete", - removal_reason="access_revoked", - observed_at=datetime.now(timezone.utc), - ) - yield RemovedScope( - record_type="event", - container_id=calendar_id, - removal_reason="access_revoked", - observed_at=datetime.now(timezone.utc), - ) + async for observation in _capture_member( + get, calendar_record, item, previous, window, tokens, coverage + ): + yield observation if not page.nextPageToken: break if page.nextPageToken in seen_pages: @@ -166,6 +155,12 @@ async def generate_calendar_observations( removal_reason="scope_removed", observed_at=datetime.now(timezone.utc), ) + yield RemovedScope( + record_type="event_occurrence", + container_id=missing, + removal_reason="scope_removed", + observed_at=datetime.now(timezone.utc), + ) yield RemovedScope( record_type="event", container_id=missing, @@ -175,3 +170,95 @@ async def generate_calendar_observations( yield CompletedScope(record_type="calendar") if cursor is not None: cursor.update(calendar_tokens=tokens) + if window is not None: + cursor.update(occurrence_coverage=coverage) + + +async def capture_occurrences( + get: GetJSON, calendar_id: str, window: CalendarOccurrenceWindow +) -> AsyncGenerator[SourceObservation, None]: + """Only a fully exhausted fixed horizon may reconcile this distinct scope.""" + yield StartedScope(record_type="event_occurrence", container_id=calendar_id) + params = { + "singleEvents": "true", + "showDeleted": "true", + "maxResults": 2500, + "timeMin": window.start.isoformat(), + "timeMax": window.end.isoformat(), + } + pages: set[str] = set() + identities: set[str] = set() + byte_count = 0 + while True: + page = CalendarPage.model_validate( + await get(f"{BASE}/calendars/{quote(calendar_id, safe='')}/events", params=params) + ) + for item in page.items: + byte_count += len(json.dumps(item).encode()) + observed = record("event_occurrence", item, calendar_id) + if observed.identity.native_id in identities: + raise ValueError("Expanded Calendar listing repeated an occurrence identity") + identities.add(observed.identity.native_id) + if len(identities) > 10000 or byte_count > 20 * 1024 * 1024: + raise ValueError("Expanded Calendar capture exceeds bounded horizon budget") + yield observed + if not page.nextPageToken: + yield CompletedScope(record_type="event_occurrence", container_id=calendar_id) + return + if page.nextPageToken in pages: + raise ValueError("Expanded Calendar listing repeated a page token") + pages.add(page.nextPageToken) + if len(pages) >= 100: + raise ValueError("Expanded Calendar capture exceeds page budget") + params["pageToken"] = page.nextPageToken + + +async def _capture_member( + get: GetJSON, + calendar_record: CaptureRecord, + item: dict[str, JsonValue], + previous: GoogleCalendarCursor, + window: CalendarOccurrenceWindow | None, + tokens: dict[str, str], + coverage: dict[str, CalendarWindowCoverage], +) -> AsyncGenerator[SourceObservation, None]: + """Capture one member or record explicit access loss without claiming its window.""" + calendar_id = calendar_record.identity.native_id + try: + async for observation in capture_calendar( + get, calendar_id, previous.calendar_tokens.get(calendar_id), tokens + ): + yield observation + if window is not None: + async for observation in capture_occurrences(get, calendar_id, window): + yield observation + calendar_timezone = item.get("timeZone") + if not isinstance(calendar_timezone, str) or not calendar_timezone: + raise ValueError("Calendar has no timezone for occurrence coverage") + coverage[calendar_id] = CalendarWindowCoverage( + start=window.start, + end=window.end, + timezone=calendar_timezone, + completed_at=datetime.now(timezone.utc), + scan_id=uuid4(), + ) + except SourceEntityNotFoundError: + yield CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id=calendar_id), + payload={"id": calendar_id}, + kind="delete", + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) + yield RemovedScope( + record_type="event_occurrence", + container_id=calendar_id, + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) + yield RemovedScope( + record_type="event", + container_id=calendar_id, + removal_reason="access_revoked", + observed_at=datetime.now(timezone.utc), + ) diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py index 8d6fdd3152..bc16e89e27 100644 --- a/backend/tests/unit/platform/sources/records/test_google_calendar.py +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -117,7 +117,8 @@ async def test_missing_calendar_removes_only_previous_calendar_children(): results = [item async for item in generate_calendar_observations(get, state)] removed = [item for item in results if type(item) is RemovedScope] assert [(item.record_type, item.container_id, item.removal_reason) for item in removed] == [ - ("event", "removed", "scope_removed") + ("event_occurrence", "removed", "scope_removed"), + ("event", "removed", "scope_removed"), ] assert state.data["calendar_tokens"] == {"kept": "new"} @@ -141,3 +142,105 @@ def test_only_recurring_cancellations_are_retained_and_reinstatement_keeps_ident record( "event", {"id": "bad", "status": "cancelled", "recurringEventId": "series"}, "calendar" ) + + +async def test_expanded_window_is_fixed_paged_and_separate_from_master_scope(): + from airweave.platform.configs.config import CalendarOccurrenceWindow + + window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") + get = ScriptedGet( + [ + {"items": [{"id": "cal", "timeZone": "America/Los_Angeles"}]}, + { + "items": [{"id": "master", "recurrence": ["RRULE:FREQ=WEEKLY"]}], + "nextSyncToken": "master-token", + }, + { + "items": [ + { + "id": "real-instance", + "recurringEventId": "master", + "originalStartTime": {"date": "2026-03-01"}, + } + ], + "nextPageToken": "p2", + }, + {"items": [{"id": "cancelled-instance", "status": "cancelled"}]}, + ] + ) + state = cursor() + observations = [item async for item in generate_calendar_observations(get, state, window)] + occurrences = [ + item + for item in observations + if isinstance(item, CaptureRecord) and item.identity.record_type == "event_occurrence" + ] + assert [item.identity.native_id for item in occurrences] == [ + "real-instance", + "cancelled-instance", + ] + assert occurrences[1].kind == "delete" + assert get.calls[2][1]["timeMin"] == get.calls[3][1]["timeMin"] + assert "syncToken" not in get.calls[2][1] and get.calls[2][1]["singleEvents"] == "true" + assert state.data["occurrence_coverage"]["cal"]["timezone"] == "America/Los_Angeles" + assert any( + type(item) is CompletedScope and item.record_type == "event_occurrence" + for item in observations + ) + + +async def test_failed_expanded_refresh_keeps_prior_coverage_and_does_not_reconcile(): + from airweave.platform.configs.config import CalendarOccurrenceWindow + + window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") + state = cursor({"cal": "old"}) + before = state.data + get = ScriptedGet( + [ + {"items": [{"id": "cal", "timeZone": "UTC"}]}, + {"items": [], "nextSyncToken": "new"}, + {"items": [{"id": "one"}], "nextPageToken": "p2"}, + SourceServerError("temporary", status_code=503), + ] + ) + observations = [] + with pytest.raises(SourceServerError): + async for item in generate_calendar_observations(get, state, window): + observations.append(item) + assert state.data == before + assert not any( + type(item) is CompletedScope and item.record_type == "event_occurrence" + for item in observations + ) + + +def test_explicit_historical_window_and_rolling_defaults_are_bounded(): + from datetime import timedelta + + from airweave.platform.configs.config import GoogleCalendarConfig + + default = GoogleCalendarConfig().resolved_window() + assert default.end - default.start == timedelta(days=120) + chosen = GoogleCalendarConfig( + occurrence_window={"start": "2001-01-01T00:00:00Z", "end": "2001-02-01T00:00:00Z"} + ) + assert chosen.resolved_window().start.year == 2001 + with pytest.raises(ValueError): + GoogleCalendarConfig( + occurrence_window={"start": "2001-01-01T00:00:00Z", "end": "2003-02-01T00:00:00Z"} + ) + + +async def test_expanded_repeated_pages_never_claim_completion(): + from airweave.platform.configs.config import CalendarOccurrenceWindow + from airweave.platform.sources.records.google_calendar import capture_occurrences + + window = CalendarOccurrenceWindow(start="2026-01-01T00:00:00Z", end="2026-02-01T00:00:00Z") + get = ScriptedGet( + [{"items": [], "nextPageToken": "repeat"}, {"items": [], "nextPageToken": "repeat"}] + ) + observed = [] + with pytest.raises(ValueError, match="repeated"): + async for item in capture_occurrences(get, "cal", window): + observed.append(item) + assert not any(type(item) is CompletedScope for item in observed) From 6555a22903d4330d0c68ec2baa980ee36ec70330 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:27:02 -0700 Subject: [PATCH 013/266] Exercise real Vespa schemas and scoped retrieval in disposable CI --- .github/workflows/owned-source-store.yml | 19 +++ backend/tests/integration/test_owned_vespa.py | 123 ++++++++++++++++++ 2 files changed, 142 insertions(+) create mode 100644 backend/tests/integration/test_owned_vespa.py diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index a65f1ea1ac..75dd2645cd 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -5,12 +5,14 @@ on: branches: [dev, main] paths: - backend/** + - vespa/app/** - Makefile - .github/workflows/owned-source-store.yml push: branches: [dev] paths: - backend/** + - vespa/app/** - Makefile - .github/workflows/owned-source-store.yml @@ -26,6 +28,15 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 25 services: + vespa: + image: vespaengine/vespa:8 + ports: + - 8081:8081 + - 19071:19071 + options: >- + --hostname vespa + --health-cmd "curl -f http://localhost:19071/state/v1/health" + --health-interval 10s --health-timeout 5s --health-retries 30 postgres: image: postgres:16.15 env: @@ -58,3 +69,11 @@ jobs: run: make test-capture - name: Verify indexed retrieval and visibility run: make test-search + - name: Record tested Vespa image + run: docker image inspect vespaengine/vespa:8 --format '{{json .RepoDigests}}' + - name: Verify real Vespa schema and retrieval + working-directory: backend + env: + OWNED_VESPA_TEST: '1' + EMBEDDING_DIMENSIONS: '384' + run: .venv/bin/pytest -q -o log_cli=false tests/integration/test_owned_vespa.py diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py new file mode 100644 index 0000000000..7617ea0bd1 --- /dev/null +++ b/backend/tests/integration/test_owned_vespa.py @@ -0,0 +1,123 @@ +"""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 asyncio +import io +import json +import os +import time +import zipfile +from pathlib import Path +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.search.adapters.vector_db.filter_translator import FilterTranslator +from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB +from airweave.domains.search.types.embeddings import QueryEmbeddings +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" +) + + +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}") + + +@pytest.mark.asyncio +async def test_real_schema_retrieval_collection_isolation_and_delete(): + """Exercise real query compilation/rank profiles; vectors are fixed fixtures.""" + package = Path(__file__).resolve().parents[3] / "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") + ) + async with httpx.AsyncClient(timeout=120) as http: + 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") + 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, + "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), + 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" + 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) From a412c886e38208f0fa308b3c90b242853d93237d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:27:33 -0700 Subject: [PATCH 014/266] Keep upstream live connector mutations opt-in on the owned fork --- .github/workflows/monke.yml | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/.github/workflows/monke.yml b/.github/workflows/monke.yml index cc08125385..d4cb0470c0 100644 --- a/.github/workflows/monke.yml +++ b/.github/workflows/monke.yml @@ -1,13 +1,9 @@ name: Monke – parallel runners on: + # The owned fork uses credential-free source-store checks on pull requests. + # Upstream's live connector mutators require separate disposable accounts. workflow_dispatch: - pull_request: - paths: - # Only trigger on platform-related changes (reduced footprint) - - "monke/**" - - "backend/airweave/platform/**" - - ".github/workflows/monke.yml" concurrency: group: monke-${{ github.ref }} From adb8d9e78263ff2377f8c2e9d8b91399115e9753 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:31:16 -0700 Subject: [PATCH 015/266] Expose full local checks and include authentication in CI --- .github/workflows/owned-source-store.yml | 5 +++-- Makefile | 14 +++++++++++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 75dd2645cd..1e23585390 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -69,11 +69,12 @@ jobs: run: make test-capture - name: Verify indexed retrieval and visibility run: make test-search + - name: Verify service authentication + run: make test-auth - name: Record tested Vespa image run: docker image inspect vespaengine/vespa:8 --format '{{json .RepoDigests}}' - name: Verify real Vespa schema and retrieval - working-directory: backend env: OWNED_VESPA_TEST: '1' EMBEDDING_DIMENSIONS: '384' - run: .venv/bin/pytest -q -o log_cli=false tests/integration/test_owned_vespa.py + run: make test-index diff --git a/Makefile b/Makefile index 9677c725ec..85525595e1 100644 --- a/Makefile +++ b/Makefile @@ -3,13 +3,16 @@ BACKEND := backend PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest -.PHONY: help setup test-store test-capture test-search validate-deploy +.PHONY: help setup check test-store test-capture test-search test-auth test-index validate-deploy help: @echo 'setup Install the locked backend and development dependencies' @echo 'test-store Verify record transactions and publication on disposable PostgreSQL' @echo 'test-capture Verify source capture, transport and blob handling' @echo 'test-search Verify search services and visibility behavior' + @echo 'test-auth Verify service-key and source authorization boundaries' + @echo 'test-index Deploy/test schemas on explicitly disposable local Vespa' + @echo 'check Run store, capture, search and authentication checks' @echo 'validate-deploy Validate the Porter application manifest locally (no deployment)' @echo '' @echo 'test-store requires CANONICAL_TEST_DATABASE_URL; it creates and removes test schemas.' @@ -30,3 +33,12 @@ test-capture: test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search + +test-auth: + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/api/tests/test_service_auth.py airweave/api/tests/test_sync_authorization.py airweave/api/tests/test_context_resolver_auth.py + +test-index: + @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py + +check: test-store test-capture test-search test-auth From 1f3353ddd72dd2c82501ab603de38e561f501504 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:32:06 -0700 Subject: [PATCH 016/266] test: verify durable Slack and Wispr capture samples --- backend/tests/live/README.md | 27 ++++ backend/tests/live/canonical_capture.py | 157 ++++++++++++------------ backend/tests/live/provider_sample.py | 108 ++++++++++++++++ 3 files changed, 215 insertions(+), 77 deletions(-) create mode 100644 backend/tests/live/provider_sample.py diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 75866f7e4a..d5e761bc1d 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -75,3 +75,30 @@ 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. + +## Slack and Wispr durable samples + +`LIVE_PROVIDERS=slack,wispr` selects the existing source adapters. Set +`LIVE_SLACK_ACCOUNT_ID`, `LIVE_WISPR_ACCOUNT_ID` and `LIVE_WISPR_USER_ID` for +previously connected, authorized accounts. Slack verifies `auth.test` identity +and the matching `users.info` email against `LIVE_EXPECTED_EMAIL`. Wispr 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. A Wispr-only +run does not require the email variable. + +Slack/Wispr stop after three distinct records containing a message/meeting, or +25 observations. No file/blob presence is required. Repeated native identities +are reduced to their latest raw observation in this bounded sample before +capture and replay; Slack history/replies can return the same parent twice. +Declared channel/container parent relationships are retained. Every record's +native JSON is compared after durable readback; identical replay must add no +journal changes. Source markers/checkpoints remain deliberately uncommitted. + +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. diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index d07b3cb690..5bc9090111 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -45,19 +45,9 @@ from airweave.domains.entities.canonical.service import CanonicalCaptureService # noqa: E402 from airweave.domains.entities.canonical.store import CanonicalRecordStore # noqa: E402 from airweave.domains.entities.canonical.tests.conftest import migrate # noqa: E402 -from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider # noqa: E402 from airweave.domains.storage.file_service import FileService # noqa: E402 from airweave.domains.storage.paths import StoragePaths # noqa: E402 from airweave.models import Organization, Sync, SyncJob # noqa: E402 -from airweave.platform.configs.config import ( # noqa: E402 - GmailConfig, - GoogleCalendarConfig, - GoogleDriveConfig, -) -from airweave.platform.http_client.airweave_client import AirweaveHttpClient # noqa: E402 -from airweave.platform.http_client.composio_transport import ComposioTransport # noqa: E402 -from airweave.platform.sources import GmailSource, GoogleDriveSource # noqa: E402 -from airweave.platform.sources.google_calendar import GoogleCalendarSource # noqa: E402 logging.disable(logging.CRITICAL) PROBE_STAGE = "initialization" @@ -88,72 +78,66 @@ def test_database_url() -> str: async def fetch_records(name, account, expected_email, key, fence, root): """Exercise actual source capture, bounded before any complete-scope claim.""" - host = "gmail.googleapis.com" if name == "gmail" else "www.googleapis.com" - auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) - transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) - transport.MAX_BINARY_BYTES = 10 * 1024 * 1024 + from provider_sample import rest_source, wispr_source + + from airweave.domains.entities.canonical.requests import RecordIdentity + from airweave.domains.entities.canonical.source import ContainerScopedSource + storage = FilesystemBackend(root / "blobs") files = FileService(fence.job_id, storage, sync_id=fence.sync_id) files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 - async with httpx.AsyncClient(transport=transport, timeout=180) as client: - wrapped = AirweaveHttpClient( - client, fence.organization_id, name, feature_flag_enabled=False + connection = ( + wispr_source(account, key, fence) + if name == "wispr" + else rest_source(name, account, expected_email, key, fence) + ) + records = {} + observations = 0 + async with asyncio.timeout(180), connection as (source, identity_verification): + parents = ( + source.canonical_container_parents if isinstance(source, ContainerScopedSource) else {} ) - if name == "gmail": - source = await GmailSource.create( - auth=auth, - logger=logging.getLogger("probe"), - http_client=wrapped, - config=GmailConfig(gmail_query="newer_than:7d smaller:5M"), - ) - profile = await source._get("https://gmail.googleapis.com/gmail/v1/users/me/profile") - email = profile.get("emailAddress", "") - elif name == "google_calendar": - source = await GoogleCalendarSource.create( - auth=auth, - logger=logging.getLogger("probe"), - http_client=wrapped, - config=GoogleCalendarConfig(), - ) - profile = await source._get( - "https://www.googleapis.com/calendar/v3/users/me/calendarList/primary" - ) - email = profile.get("id", "") - else: - source = await GoogleDriveSource.create( - auth=auth, - logger=logging.getLogger("probe"), - http_client=wrapped, - config=GoogleDriveConfig(), - ) - profile = await source._get( - "https://www.googleapis.com/drive/v3/about", params={"fields": "user(emailAddress)"} - ) - email = profile.get("user", {}).get("emailAddress", "") - if email.casefold() != expected_email.casefold(): - raise ValueError("The connected account did not match the explicitly selected identity") - records = [] - async with ( - asyncio.timeout(180), - aclosing(source.generate_observations(files=files)) as stream, - ): + async with aclosing(source.generate_observations(files=files)) as stream: async for item in stream: - # Deliberately ignore Started/CompletedScope and cursor checkpoints: this - # selected sample does not prove exhaustive source enumeration. - if isinstance(item, CaptureRecord): - records.append(item) - if len(records) >= 25 or ( - name != "google_calendar" - and len(records) >= 3 - and any(x.blobs for x in records) - ): - break - if name == "google_calendar": - if not any(x.identity.record_type == "event" for x in records): - raise ValueError("Bounded Calendar sample did not contain an event") - elif not records or not any(x.blobs for x in records): - raise ValueError("Bounded source sample did not contain a real owned blob") - return records, storage + # No scope-completion marker/checkpoint is saved by a bounded sample. + if not isinstance(item, CaptureRecord): + continue + observations += 1 + parent_type = parents.get(item.identity.record_type) + if parent_type is not None: + # Apply the declared source relationship just as the pipeline does; + # preserve the provider payload verbatim. + parent = RecordIdentity( + record_type=parent_type, native_id=item.identity.container_id + ) + if item.parent is not None and item.parent != parent: + raise ValueError("Source parent conflicts with its declared relationship") + item = item.model_copy(update={"parent": parent}) + identity = (item.identity.record_type, item.identity.entity_key) + records[identity] = item + if sample_complete(name, list(records.values()), observations): + break + values = list(records.values()) + required = {"google_calendar": "event", "slack": "message", "wispr": "meeting"}.get(name) + if required and not any(x.identity.record_type == required for x in values): + raise ValueError("Bounded source sample did not contain the required record type") + if name in {"gmail", "google_drive"} and not any(x.blobs for x in values): + raise ValueError("Bounded source sample did not contain a real owned blob") + return values, storage, identity_verification, observations + + +def sample_complete(name, records, observations): + """Bound provider work; a sample is not a completed enumeration.""" + if observations >= 25: + return True + if name == "google_calendar": + return False + if len(records) < 3: + return False + if name in {"gmail", "google_drive"}: + return any(record.blobs for record in records) + required = "message" if name == "slack" else "meeting" + return any(record.identity.record_type == required for record in records) async def verify(name, account, email, key, sessions, engine, root): @@ -177,14 +161,16 @@ async def verify(name, account, email, key, sessions, engine, root): ) global PROBE_STAGE PROBE_STAGE = name + ":provider_capture" - records, _ = await fetch_records(name, account, email, key, fence, root) + records, _, identity_verification, observations = await fetch_records( + name, account, email, key, fence, root + ) PROBE_STAGE = name + ":durable_capture" async with sessions() as db: committed = await service.capture(db, CaptureBatch(fence=fence, records=tuple(records))) await engine.dispose() # Force readback over fresh PostgreSQL connections. PROBE_STAGE = name + ":durable_readback" storage = FilesystemBackend(root / "blobs") # New filesystem storage instance. - expected = {item.identity.entity_key: item for item in records} + expected = {(item.identity.record_type, item.identity.entity_key): item for item in records} listed = [] cursor = None while True: @@ -202,7 +188,10 @@ async def verify(name, account, email, key, sessions, engine, root): assert len(listed) == len(expected) total_bytes = blob_count = 0 for stored in listed: - assert stored.payload == expected[stored.identity.entity_key].payload + assert ( + stored.payload + == expected[(stored.identity.record_type, stored.identity.entity_key)].payload + ) async with sessions() as db: reread = await query.read(db, organization_id, sync_id, stored.id) assert reread.payload == stored.payload and reread.revision == 1 @@ -246,6 +235,10 @@ async def verify(name, account, email, key, sessions, engine, root): return { "almanac_reader": consumer, "provider": name, + "identity_verification": identity_verification, + "observations": observations, + "partial_records": sum(record.completeness != "complete" for record in listed), + "record_types": sorted({record.identity.record_type for record in listed}), "records": len(listed), "journal_changes": len(changes), "blobs": blob_count, @@ -263,14 +256,24 @@ async def main(): url = test_database_url() key = os.environ["COMPOSIO_API_KEY"] selected = os.environ.get("LIVE_PROVIDERS", "gmail,google_drive").split(",") - if not selected or set(selected) - {"gmail", "google_drive", "google_calendar"}: - raise ValueError("LIVE_PROVIDERS must select gmail, google_drive and/or google_calendar") + if not selected or set(selected) - { + "gmail", + "google_drive", + "google_calendar", + "slack", + "wispr", + }: + raise ValueError( + "LIVE_PROVIDERS must select gmail, google_drive, google_calendar, slack or wispr" + ) inputs = [ - (name, os.environ[variable], os.environ["LIVE_EXPECTED_EMAIL"]) + (name, os.environ[variable], os.environ["LIVE_EXPECTED_EMAIL"] if name != "wispr" else "") for name, variable in ( ("gmail", "LIVE_GMAIL_ACCOUNT_ID"), ("google_drive", "LIVE_DRIVE_ACCOUNT_ID"), ("google_calendar", "LIVE_CALENDAR_ACCOUNT_ID"), + ("slack", "LIVE_SLACK_ACCOUNT_ID"), + ("wispr", "LIVE_WISPR_ACCOUNT_ID"), ) if name in selected ] diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py new file mode 100644 index 0000000000..b69f63de31 --- /dev/null +++ b/backend/tests/live/provider_sample.py @@ -0,0 +1,108 @@ +"""Account-checked source construction for the bounded live capture harness.""" + +import logging +import os +from contextlib import asynccontextmanager +from urllib.parse import quote + +import httpx + +from airweave.domains.sources.token_providers.protocol import ( + ManagedAuthProvider, + ManagedToolAuthProvider, +) +from airweave.platform.configs.config import ( + GmailConfig, + GoogleCalendarConfig, + GoogleDriveConfig, + SlackConfig, +) +from airweave.platform.http_client.airweave_client import AirweaveHttpClient +from airweave.platform.http_client.composio_transport import ComposioTransport +from airweave.platform.sources import GmailSource, GoogleDriveSource +from airweave.platform.sources.google_calendar import GoogleCalendarSource +from airweave.platform.sources.slack import SlackSource +from airweave.platform.sources.wispr import WisprSource + + +async def verify_rest_identity(name, source, expected_email): + """Use actual provider identity; names/connection labels are not identity.""" + if name == "gmail": + profile = await source._get("https://gmail.googleapis.com/gmail/v1/users/me/profile") + email = profile.get("emailAddress", "") + elif name == "google_calendar": + profile = await source._get( + "https://www.googleapis.com/calendar/v3/users/me/calendarList/primary" + ) + email = profile.get("id", "") + elif name == "slack": + identity = await source._get("https://slack.com/api/auth.test") + if not identity.get("team_id") or not identity.get("user_id"): + raise ValueError("Slack did not return workspace and user identity") + profile = await source._get( + "https://slack.com/api/users.info", {"user": identity["user_id"]} + ) + if profile.get("user", {}).get("id") != identity["user_id"]: + raise ValueError("Slack returned a different user") + email = profile.get("user", {}).get("profile", {}).get("email", "") + else: + profile = await source._get( + "https://www.googleapis.com/drive/v3/about", params={"fields": "user(emailAddress)"} + ) + email = profile.get("user", {}).get("emailAddress", "") + if email.casefold() != expected_email.casefold(): + raise ValueError("Provider identity does not match the explicitly selected email") + + +@asynccontextmanager +async def rest_source(name, account, expected_email, key, fence): + host = {"gmail": "gmail.googleapis.com", "slack": "slack.com"}.get(name, "www.googleapis.com") + auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) + transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) + transport.MAX_BINARY_BYTES = 10 * 1024 * 1024 + async with httpx.AsyncClient(transport=transport, timeout=180) as client: + wrapped = AirweaveHttpClient( + client, fence.organization_id, name, feature_flag_enabled=False + ) + source_type, config = { + "gmail": (GmailSource, GmailConfig(gmail_query="newer_than:7d smaller:5M")), + "google_drive": (GoogleDriveSource, GoogleDriveConfig()), + "google_calendar": (GoogleCalendarSource, GoogleCalendarConfig()), + "slack": (SlackSource, SlackConfig()), + }[name] + source = await source_type.create( + auth=auth, logger=logging.getLogger("probe"), http_client=wrapped, config=config + ) + await verify_rest_identity(name, source, expected_email) + yield source, "provider_email" + + +@asynccontextmanager +async def wispr_source(account, key, fence): + """Verify exact Composio account binding; Wispr exposes no identity profile here.""" + expected_user = os.environ["LIVE_WISPR_USER_ID"] + async with httpx.AsyncClient(timeout=180) as client: + metadata_response = await client.get( + "https://backend.composio.dev/api/v3/connected_accounts/" + quote(account, safe=""), + headers={"x-api-key": key}, + ) + metadata_response.raise_for_status() + metadata = metadata_response.json() + if ( + metadata.get("id") != account + or metadata.get("status") != "ACTIVE" + or metadata.get("toolkit", {}).get("slug") != "wispr_flow_mcp" + or metadata.get("user_id") != expected_user + ): + raise ValueError("Wispr connected-account identity/state mismatch") + source = await WisprSource.create( + auth=ManagedToolAuthProvider( + api_key=key, connected_account_id=account, user_id=expected_user + ), + logger=logging.getLogger("probe"), + http_client=AirweaveHttpClient( + client, fence.organization_id, "wispr", feature_flag_enabled=False + ), + ) + await source.validate() + yield source, "composio_account_principal_only" From 71d53fc79571df6fc83707cb77acb7ba1bf4d1f5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:33:06 -0700 Subject: [PATCH 017/266] Supply explicit query bounds in real Vespa smoke --- backend/tests/integration/test_owned_vespa.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index 7617ea0bd1..6696576c0e 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -105,7 +105,12 @@ async def test_real_schema_retrieval_collection_isolation_and_delete(): ) for mode in RetrievalStrategy: compiled = await engine.compile_query( - SearchPlan(query=SearchQuery(primary="fundraising"), retrieval_strategy=mode), + SearchPlan( + query=SearchQuery(primary="fundraising"), + retrieval_strategy=mode, + limit=20, + offset=0, + ), embeddings, collection, ) From adfc89e254369d74f76de2baa9455b193e4b9879 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:38:42 -0700 Subject: [PATCH 018/266] Pin CI Vespa to the verified engine image --- .github/workflows/owned-source-store.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 1e23585390..5132e75db2 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -29,7 +29,7 @@ jobs: timeout-minutes: 25 services: vespa: - image: vespaengine/vespa:8 + image: vespaengine/vespa@sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd ports: - 8081:8081 - 19071:19071 @@ -72,7 +72,7 @@ jobs: - name: Verify service authentication run: make test-auth - name: Record tested Vespa image - run: docker image inspect vespaengine/vespa:8 --format '{{json .RepoDigests}}' + run: docker image inspect vespaengine/vespa@sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd --format '{{json .RepoDigests}}' - name: Verify real Vespa schema and retrieval env: OWNED_VESPA_TEST: '1' From b65dea7222b6a7bdbe70fe7d5281d0938fc8755b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:39:34 -0700 Subject: [PATCH 019/266] Build and verify the staging container in CI without publishing --- .github/workflows/owned-source-store.yml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 5132e75db2..1edf3917fe 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -78,3 +78,16 @@ jobs: OWNED_VESPA_TEST: '1' EMBEDDING_DIMENSIONS: '384' run: make test-index + + image: + needs: correctness + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + - name: Build the deployable image without publishing + run: docker build --tag almanac-source-store:ci backend + - name: Verify packaged Python and entrypoint syntax + run: | + docker run --rm --entrypoint python almanac-source-store:ci -m compileall -q airweave + docker run --rm --entrypoint /bin/sh almanac-source-store:ci -n /app/entrypoint.sh From d3e0649b44f86d8f1c43b86638bcb0324d16a064 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:39:48 -0700 Subject: [PATCH 020/266] Select captured calendars and revoke withdrawn or inaccessible scopes --- .../canonical/tests/test_calendar_range.py | 81 +++++++++++ backend/airweave/platform/configs/config.py | 13 ++ .../platform/sources/google_calendar.py | 5 +- .../sources/records/google_calendar.py | 127 ++++++++++++------ .../sources/records/test_google_calendar.py | 71 ++++++++++ 5 files changed, 255 insertions(+), 42 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py index 680b554313..401888ff9f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py @@ -254,3 +254,84 @@ async def test_budget_rejects_oversized_row_before_orm_materialization(database, db, fence.organization_id, fence.sync_id, "cal" ) assert not any(isinstance(item, Entity) for item in db.identity_map.values()) + + +async def test_empty_selection_reconciles_partial_parent_without_saved_cursor(database, source): + from airweave.domains.entities.canonical.requests import ( + CompletedScope, + ReconcileScope, + StartedScope, + ) + from airweave.platform.sources.records.google_calendar import generate_calendar_observations + + capture, fence = source + parent = record("calendar", {"id": "partial", "timeZone": "UTC"}) + child = record("event", {"id": "old", "status": "confirmed"}, "partial").model_copy( + update={"parent": parent.identity} + ) + async with database() as db: + result = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, child))) + child_id = result.changes[-1].record.id + + async def never_get(*args, **kwargs): + raise AssertionError("Empty selection must not fetch provider") + + # Same-attempt retry is intentionally harder: StartedScope must clear old sightings. + async for item in generate_calendar_observations(never_get, None, calendar_ids=()): + if isinstance(item, StartedScope): + async with database() as db: + await capture.start_scope(db, fence, item) + elif type(item) is CompletedScope: + async with database() as db: + await capture.reconcile_scope( + db, + ReconcileScope(fence=fence, scope=item, observed_at=datetime.now(timezone.utc)), + ) + async with database() as db: + current = await capture.store.read(db, fence.organization_id, fence.sync_id, child_id) + assert current.content_access == "unavailable" and current.payload == {} + await db.rollback() + removed = await capture.reconcile_parents(db, fence) + assert removed.capture.changes[0].record.removal_reason != "provider_deleted" + + +async def test_missing_selected_calendar_revokes_partial_parent_without_cursor(database, source): + from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + RemovedScope, + StartedScope, + ) + from airweave.platform.sources.records.google_calendar import generate_calendar_observations + + capture, fence = source + parent = record("calendar", {"id": "missing", "timeZone": "UTC"}) + child = record("event", {"id": "retained", "status": "confirmed"}, "missing").model_copy( + update={"parent": parent.identity} + ) + async with database() as db: + initial = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, child))) + child_id = initial.changes[-1].record.id + + async def empty_membership(*args, **kwargs): + return {"items": []} + + with pytest.raises(ValueError, match="not accessible"): + async for item in generate_calendar_observations( + empty_membership, None, calendar_ids=("missing",) + ): + async with database() as db: + if isinstance(item, CaptureRecord): + await capture.capture(db, CaptureBatch(fence=fence, records=(item,))) + elif isinstance(item, StartedScope): + await capture.start_scope(db, fence, item) + elif isinstance(item, RemovedScope): + await capture.remove_scope(db, fence, item) + else: + raise AssertionError("Missing requested scope cannot complete") + async with database() as db: + current = await capture.store.read(db, fence.organization_id, fence.sync_id, child_id) + assert current.content_access == "unavailable" and current.payload == {} + assert current.removal_reason == "access_revoked" + assert ( + await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) is None + ) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 07614b673d..428478056e 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -326,6 +326,19 @@ def bounded(self): class GoogleCalendarConfig(SourceConfig): """Rolling expanded capture, or an explicit window through source configuration.""" + 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 diff --git a/backend/airweave/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index 3a570cc84d..447ae58aaa 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -110,7 +110,10 @@ async def generate_observations( if node_selections: raise ValueError("Calendar selection scopes are not supported by canonical capture") async for observation in generate_calendar_observations( - self._get, cursor, self.calendar_config.resolved_window() + self._get, + cursor, + self.calendar_config.resolved_window(), + calendar_ids=self.calendar_config.calendar_ids, ): yield observation diff --git a/backend/airweave/platform/sources/records/google_calendar.py b/backend/airweave/platform/sources/records/google_calendar.py index 3123d150ba..a3e76bb37d 100644 --- a/backend/airweave/platform/sources/records/google_calendar.py +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -1,8 +1,10 @@ """Original Calendar resources with per-calendar incremental synchronization.""" import json -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator +from contextlib import aclosing from datetime import datetime, timezone +from typing import Literal from urllib.parse import quote from uuid import uuid4 @@ -116,57 +118,35 @@ async def capture_calendar( async def generate_calendar_observations( - get: GetJSON, cursor: SyncCursor | None, window: CalendarOccurrenceWindow | None = None + get: GetJSON, + cursor: SyncCursor | None, + window: CalendarOccurrenceWindow | None = None, + *, + calendar_ids: tuple[str, ...] | None = None, ) -> AsyncGenerator[SourceObservation, None]: """Refresh calendar membership; capture each accessible calendar independently.""" previous = GoogleCalendarCursor.model_validate(cursor.data if cursor else {}) + selection = set(calendar_ids) if calendar_ids is not None else None + withdrawn = previous.calendar_tokens.keys() - selection if selection is not None else set() + for calendar_id in withdrawn: + for observation in removed_calendar(calendar_id, "scope_removed"): + yield observation tokens: dict[str, str] = {} coverage: dict[str, CalendarWindowCoverage] = {} seen_calendars: set[str] = set() - seen_pages: set[str] = set() - params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"} yield StartedScope(record_type="calendar") - while True: - page = CalendarPage.model_validate( - await get(f"{BASE}/users/me/calendarList", params=params) - ) - for item in page.items: - calendar_record = record("calendar", item) - calendar_id = calendar_record.identity.native_id - if calendar_id in seen_calendars: - raise ValueError("Calendar list repeated an identity across pages") - seen_calendars.add(calendar_id) + async with aclosing(selected_members(get, selection)) as members: + async for calendar_record in members: + seen_calendars.add(calendar_record.identity.native_id) yield calendar_record async for observation in _capture_member( - get, calendar_record, item, previous, window, tokens, coverage + get, calendar_record, calendar_record.payload, previous, window, tokens, coverage ): yield observation - if not page.nextPageToken: - break - if page.nextPageToken in seen_pages: - raise ValueError("Calendar list returned a repeated page token") - seen_pages.add(page.nextPageToken) - params["pageToken"] = page.nextPageToken - for missing in previous.calendar_tokens.keys() - seen_calendars: - yield CaptureRecord( - identity=RecordIdentity(record_type="calendar", native_id=missing), - payload={"id": missing}, - kind="delete", - removal_reason="scope_removed", - observed_at=datetime.now(timezone.utc), - ) - yield RemovedScope( - record_type="event_occurrence", - container_id=missing, - removal_reason="scope_removed", - observed_at=datetime.now(timezone.utc), - ) - yield RemovedScope( - record_type="event", - container_id=missing, - removal_reason="scope_removed", - observed_at=datetime.now(timezone.utc), - ) + for observation in missing_members( + set(previous.calendar_tokens), seen_calendars, withdrawn, selection + ): + yield observation yield CompletedScope(record_type="calendar") if cursor is not None: cursor.update(calendar_tokens=tokens) @@ -262,3 +242,68 @@ async def _capture_member( removal_reason="access_revoked", observed_at=datetime.now(timezone.utc), ) + + +def removed_calendar( + calendar_id: str, reason: Literal["scope_removed", "access_revoked"] +) -> Iterator[SourceObservation]: + """Withdraw parent visibility and exact child scopes without alleging provider deletion.""" + now = datetime.now(timezone.utc) + yield CaptureRecord( + identity=RecordIdentity(record_type="calendar", native_id=calendar_id), + payload={"id": calendar_id}, + kind="delete", + removal_reason=reason, + observed_at=now, + ) + for record_type in ("event_occurrence", "event"): + yield RemovedScope( + record_type=record_type, + container_id=calendar_id, + removal_reason=reason, + observed_at=now, + ) + + +async def selected_members( + get: GetJSON, selection: set[str] | None +) -> AsyncGenerator[CaptureRecord, None]: + """Enumerate membership completely, but fetch content only for configured native IDs.""" + seen_pages: set[str] = set() + seen_calendars: set[str] = set() + params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"} + while selection != set(): + page = CalendarPage.model_validate( + await get(f"{BASE}/users/me/calendarList", params=params) + ) + for item in page.items: + member = record("calendar", item) + calendar_id = member.identity.native_id + if calendar_id in seen_calendars: + raise ValueError("Calendar list repeated an identity across pages") + seen_calendars.add(calendar_id) + if selection is None or calendar_id in selection: + yield member + if not page.nextPageToken: + return + if page.nextPageToken in seen_pages: + raise ValueError("Calendar list returned a repeated page token") + seen_pages.add(page.nextPageToken) + params["pageToken"] = page.nextPageToken + + +def missing_members( + previous: set[str], seen_calendars: set[str], withdrawn: set[str], selection: set[str] | None +) -> Iterator[SourceObservation]: + """Revoke membership-proven absence independently of saved cursor history.""" + missing_requested = selection - seen_calendars if selection is not None else set() + for missing in previous - seen_calendars - withdrawn - missing_requested: + for observation in removed_calendar(missing, "scope_removed"): + yield observation + for missing in missing_requested: + # Full membership exhaustion proves absence even when an earlier partial + # capture never saved this calendar's token/checkpoint. + for observation in removed_calendar(missing, "access_revoked"): + yield observation + if missing_requested: + raise ValueError("A requested calendar is not accessible in current membership") diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py index bc16e89e27..fa8d60a433 100644 --- a/backend/tests/unit/platform/sources/records/test_google_calendar.py +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -244,3 +244,74 @@ async def test_expanded_repeated_pages_never_claim_completion(): async for item in capture_occurrences(get, "cal", window): observed.append(item) assert not any(type(item) is CompletedScope for item in observed) + + +async def test_selection_withdraws_old_scope_and_readd_bootstraps(): + state = cursor({"old": "old-token", "keep": "keep-token"}) + get = ScriptedGet( + [{"items": [{"id": "old"}, {"id": "keep"}]}, {"items": [], "nextSyncToken": "next"}] + ) + result = [ + item async for item in generate_calendar_observations(get, state, calendar_ids=("keep",)) + ] + assert len(get.calls) == 2 and "/keep/events" in get.calls[1][0] + assert get.calls[1][1]["syncToken"] == "keep-token" + assert state.data["calendar_tokens"] == {"keep": "next"} + assert {item.record_type for item in result if isinstance(item, RemovedScope)} == { + "event", + "event_occurrence", + } + assert all( + item.removal_reason == "scope_removed" for item in result if isinstance(item, RemovedScope) + ) + readd = ScriptedGet([{"items": [{"id": "old"}]}, {"items": [], "nextSyncToken": "fresh"}]) + _ = [item async for item in generate_calendar_observations(readd, state, calendar_ids=("old",))] + assert "syncToken" not in readd.calls[1][1] + + +async def test_empty_selection_makes_no_provider_requests_and_completes_membership(): + state = cursor({"old": "token"}) + get = ScriptedGet([]) + result = [item async for item in generate_calendar_observations(get, state, calendar_ids=())] + assert get.calls == [] and state.data["calendar_tokens"] == {} + assert any(type(item) is CompletedScope and item.record_type == "calendar" for item in result) + + +async def test_missing_requested_calendar_revokes_known_scope_without_completion(): + state = cursor({"missing": "token"}) + result = [] + with pytest.raises(ValueError, match="not accessible"): + async for item in generate_calendar_observations( + ScriptedGet([{"items": []}]), state, calendar_ids=("missing",) + ): + result.append(item) + assert state.data["calendar_tokens"] == {"missing": "token"} + assert any( + isinstance(item, RemovedScope) and item.removal_reason == "access_revoked" + for item in result + ) + assert not any(type(item) is CompletedScope for item in result) + + +def test_calendar_selection_config_distinguishes_none_empty_and_invalid(): + from airweave.platform.configs.config import GoogleCalendarConfig + + assert GoogleCalendarConfig().calendar_ids is None + assert GoogleCalendarConfig(calendar_ids=[]).calendar_ids == () + for value in (["primary"], ["a", "a"], [""], [" a"]): + with pytest.raises(ValueError): + GoogleCalendarConfig(calendar_ids=value) + + +async def test_selected_calendar_loses_access_between_membership_and_events(): + from airweave.domains.sources.exceptions import SourceEntityNotFoundError + + state = cursor({"selected": "old"}) + get = ScriptedGet([{"items": [{"id": "selected"}]}, SourceEntityNotFoundError("gone")]) + result = [ + item + async for item in generate_calendar_observations(get, state, calendar_ids=("selected",)) + ] + revoked = [item for item in result if isinstance(item, RemovedScope)] + assert len(revoked) == 2 and all(item.removal_reason == "access_revoked" for item in revoked) + assert state.data["calendar_tokens"] == {} From 9e6c9fd3751af807f9ab1351f2338c2803b2c82e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:46:16 -0700 Subject: [PATCH 021/266] test: verify owned source lifecycle across fresh processes --- .../platform/sources/gmail_capture.py | 3 + .../sources/tests/test_gmail_capture.py | 26 +- backend/tests/live/README.md | 53 +++ .../evidence/gmail-lifecycle-20260930.json | 64 +++ backend/tests/live/provider_lifecycle.py | 413 ++++++++++++++++++ backend/tests/live/provider_sample.py | 21 +- 6 files changed, 565 insertions(+), 15 deletions(-) create mode 100644 backend/tests/live/evidence/gmail-lifecycle-20260930.json create mode 100644 backend/tests/live/provider_lifecycle.py diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py index ef5bb93e7f..311234e395 100644 --- a/backend/airweave/platform/sources/gmail_capture.py +++ b/backend/airweave/platform/sources/gmail_capture.py @@ -11,6 +11,7 @@ CaptureRecord, CompletedScope, RecordIdentity, + StartedScope, ) from airweave.domains.entities.canonical.source import SourceObservation from airweave.domains.storage.exceptions import FileSkippedException @@ -185,6 +186,7 @@ async def generate(self, cursor: SyncCursor | None) -> AsyncGenerator[SourceObse successful immutable storage; missing storage or oversize bodies remain partial. """ if self.query: + yield StartedScope(record_type="message") async for record in self.enumerate(): yield record yield CompletedScope(record_type="message") @@ -196,6 +198,7 @@ async def generate(self, cursor: SyncCursor | None) -> AsyncGenerator[SourceObse if history is None: profile = await self.get(f"{BASE}/profile") boundary = str(profile["historyId"]) + yield StartedScope(record_type="message") async for record in self.enumerate(): yield record # If this boundary expires mid-crawl, fail the run. Starting another full diff --git a/backend/airweave/platform/sources/tests/test_gmail_capture.py b/backend/airweave/platform/sources/tests/test_gmail_capture.py index 4f8095ee7f..7573ca4911 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_capture.py +++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py @@ -5,7 +5,7 @@ import httpx import pytest -from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.entities.canonical.requests import CompletedScope, StartedScope from airweave.domains.syncs.cursors.cursor import SyncCursor from airweave.platform.cursors.gmail import GmailCursor from airweave.platform.sources.gmail_capture import BASE, GmailCapture @@ -65,8 +65,9 @@ async def test_bootstrap_boundary_precedes_pagination_and_replays_concurrent_edi ) state = cursor() results = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert [r.payload["id"] for r in results[:-1]] == ["a", "b", "a", "c"] - assert results[2].payload["labelIds"] == [] + assert results[0] == StartedScope(record_type="message") + assert [r.payload["id"] for r in results[1:-1]] == ["a", "b", "a", "c"] + assert results[3].payload["labelIds"] == [] assert isinstance(results[-1], CompletedScope) assert state.data["history_id"] == "12" assert provider.calls[3][1]["pageToken"] == "p2" @@ -115,7 +116,7 @@ async def test_expired_incremental_boundary_bootstraps_and_reconciles(): ) state = cursor(history_id="20", canonical_query="") results = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert results == [CompletedScope(record_type="message")] + assert results == [StartedScope(record_type="message"), CompletedScope(record_type="message")] assert state.data["history_id"] == "31" @@ -136,7 +137,8 @@ async def test_failure_after_full_crawl_does_not_finish_scope_or_advance_checkpo with pytest.raises(type(error)): async for record in GmailCapture(provider.get, None).generate(state): observed.append(record) - assert len(observed) == 1 + assert len(observed) == 2 + assert observed[0] == StartedScope(record_type="message") assert state.data == before @@ -153,9 +155,10 @@ async def test_filtered_query_always_reconciles_and_external_mime_bytes_are_part results = [ r async for r in GmailCapture(provider.get, "from:person@example.com").generate(state) ] - assert results[0].payload == raw - assert results[0].completeness == "partial" - assert results[0].blobs == () + assert results[0] == StartedScope(record_type="message") + assert results[1].payload == raw + assert results[1].completeness == "partial" + assert results[1].blobs == () assert results[-1] == CompletedScope(record_type="message") assert provider.calls[0][1]["q"] == "from:person@example.com" assert state.data["history_id"] == "" @@ -171,7 +174,7 @@ async def test_legacy_entity_cursor_cannot_skip_canonical_bootstrap(): ] ) results = [r async for r in GmailCapture(provider.get, None).generate(cursor(history_id="20"))] - assert results == [CompletedScope(record_type="message")] + assert results == [StartedScope(record_type="message"), CompletedScope(record_type="message")] @pytest.mark.asyncio @@ -216,8 +219,9 @@ async def test_filtered_page_failure_never_claims_completed_scope(): with pytest.raises(RuntimeError): async for record in GmailCapture(provider.get, "in:inbox").generate(state): results.append(record) - assert len(results) == 1 - assert not any(isinstance(r, CompletedScope) for r in results) + assert len(results) == 2 + assert results[0] == StartedScope(record_type="message") + assert not any(type(r) is CompletedScope for r in results) assert state.data["history_id"] == "20" diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index d5e761bc1d..0d11a48139 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -102,3 +102,56 @@ 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. + +## Complete 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: before:` query before either run, +with no size or category filter. Each run is a full reconciliation of this query; +this does **not** verify Gmail history-delta resume or the whole mailbox. + +Two fresh Python processes run the production `SyncOrchestrator.run`, source, +canonical pipeline, job state machine/repository and cursor read/write path. +The harness injects private database sessions and disables external telemetry; +billing, non-applicable ACL and sync-pausing collaborators are test doubles. +It does not exercise the hosted factory, Temporal worker, WorkOS, or a product +account binding. The second process loads the persisted cursor; the successful +checkpoint must identify its committed attempt. Native JSON/revision digests +compare the two reads. A real mailbox change may make the second digest differ. + +Per-run limits: 250 message observations, 600 provider requests, 10 MiB per MIME +blob, 256 MiB total blob writes, and ten minutes. The record/request/storage +limits raise instead of truncating enumeration; no completed scope or checkpoint +is claimed on an aborted enumeration. Oversized MIME parts retain the source's explicit +partial-content status. The parent reaps the child, removes its schema and private +files on success or failure, and prints only counts, status and digests. + +Verified 2026-09-30: both fresh-process seven-day runs completed, each with 239 +message observations, 287 provider requests, 239 stored records, zero partial +records, and 5,501,801 MIME blob bytes written. Each emitted one start and one +completion marker and saved a checkpoint; the second loaded the durable cursor. +Its final change sequence was 270, so the two raw-payload/revision digests were +**not identical**. This is not a zero-change second-live-run claim. A separate +bounded diagnostic repeated one attachment-bearing message GET and observed only +`payload.parts[*].body.attachmentId` differences. Gmail can rotate these original +attachment locators; the canonical store preserves them. That diagnostic does +not retrospectively attribute all 31 earlier changes. Duplicate replay of an +identical captured observation remains independently tested. + +Both private schema/file cleanup checks returned zero leftovers. Fifteen Gmail +source tests and four real-PostgreSQL pipeline tests passed; failure/recovery +fixtures are simulated provider evidence, not live induced failures. The existing +Gmail source now emits an explicit `StartedScope` before full enumeration, +including filtered reconciliation and expired-history bootstrap. + +Calendar verification via `LIVE_LIFECYCLE_PROVIDER=google_calendar`: one verified +primary calendar with a fixed seven-day occurrence window completed both fresh +processes. Initial run: 4 provider requests, 190 observations, 183 active records, +3 completed scopes. Second run: 4 requests, 22 observations, 2 completed scopes, +and one actual event listing with the saved `syncToken`. Both jobs completed and +saved checkpoints; second-run change sequence stayed 190 and the native +payload/revision digest was identical. Zero partial records or blobs. Private +schema/files were removed. No all-day or DST case was observed; those semantics +remain covered separately by synthetic tests. Calendar limits are 100 provider +requests, 10,000 observations and 180 seconds per process. diff --git a/backend/tests/live/evidence/gmail-lifecycle-20260930.json b/backend/tests/live/evidence/gmail-lifecycle-20260930.json new file mode 100644 index 0000000000..7e7ba1f242 --- /dev/null +++ b/backend/tests/live/evidence/gmail-lifecycle-20260930.json @@ -0,0 +1,64 @@ +{ + "date": "2026-09-30", + "mode": "filtered_full_reconciliation", + "scope": "fixed seven-day after/before epoch query; no size/category filter", + "provider_writes": false, + "runs": [ + { + "process": 1, + "provider_requests": 287, + "records_observed": 239, + "stored_records": 239, + "partial_records": 0, + "scope_started": 1, + "scope_completed": 1, + "blob_bytes_written": 5501801, + "checkpoint_saved": true, + "loaded_durable_checkpoint": false, + "job_status": "completed", + "payload_revision_digest": "06bec711f3e00a2101e68c1b8fd56120f91a4c18cc4ecdff54e52f015a17a556" + }, + { + "process": 2, + "provider_requests": 287, + "records_observed": 239, + "stored_records": 239, + "partial_records": 0, + "scope_started": 1, + "scope_completed": 1, + "blob_bytes_written": 5501801, + "checkpoint_saved": true, + "loaded_durable_checkpoint": true, + "observed_change_sequence": 270, + "job_status": "completed", + "payload_revision_digest": "d282ee7546026cfdc02af2bc3da16010bd94271eec5ef9bb2666a85339f79d83" + } + ], + "identical_second_read": false, + "additional_changes": 31, + "repeated_native_get_diagnostic": { + "same_message": true, + "changed_paths": [ + "/payload/parts/1/body/attachmentId", + "/payload/parts/2/body/attachmentId" + ], + "limitation": "One immediate repeated GET isolated native attachment locator volatility; it does not attribute every earlier change." + }, + "cleanup": { + "remaining_live_schemas": 0, + "remaining_private_live_directories": 0 + }, + "verification": { + "gmail_source_tests_passed": 15, + "postgres_pipeline_tests_passed": 4, + "simulated_budget_guard": "251st observation raises before completed scope" + }, + "not_verified": [ + "Gmail history-delta mode", + "whole mailbox completeness", + "Temporal worker or hosted factory", + "hosted authentication or real product binding", + "live induced failure/recovery", + "search quality" + ] +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py new file mode 100644 index 0000000000..c1ec5eff75 --- /dev/null +++ b/backend/tests/live/provider_lifecycle.py @@ -0,0 +1,413 @@ +"""Opt-in fixed Gmail reconciliation or selected Calendar incremental lifecycle. + +Two fresh processes; no Temporal worker, hosted factory, auth or product binding claim. +Parent owns private schema/files cleanup on success and failure. No provider writes. +""" + +import asyncio +import hashlib +import json +import logging +import os +import sys +import time +from contextlib import aclosing, asynccontextmanager +from datetime import datetime, timedelta, timezone +from pathlib import Path +from tempfile import TemporaryDirectory +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch +from uuid import UUID, uuid4 + +import canonical_capture as harness # settings must precede application imports +from provider_sample import rest_source +from sqlalchemy import select, text +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from airweave.adapters.event_bus.fake import FakeEventBus +from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope +from airweave.domains.entities.canonical.service import CanonicalCaptureService +from airweave.domains.entities.canonical.store import CanonicalRecordStore +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.domains.sync_pipeline.config import SyncConfig +from airweave.domains.sync_pipeline.contexts.runtime import SyncRuntime +from airweave.domains.sync_pipeline.orchestrator import SyncOrchestrator +from airweave.domains.sync_pipeline.pipeline.entity_tracker import EntityTracker +from airweave.domains.sync_pipeline.stream import AsyncSourceStream +from airweave.domains.sync_pipeline.worker_pool import AsyncWorkerPool +from airweave.domains.syncs.cursors.cursor import SyncCursor +from airweave.domains.syncs.cursors.service import SyncCursorService +from airweave.domains.syncs.jobs.repository import SyncJobRepository +from airweave.domains.syncs.jobs.state_machine import SyncJobStateMachine +from airweave.models import Entity, Organization, Sync, SyncJob +from airweave.models import SyncCursor as StoredCursor +from airweave.platform.configs.config import GoogleCalendarConfig +from airweave.platform.cursors.gmail import GmailCursor +from airweave.platform.cursors.google_calendar import GoogleCalendarCursor + + +class BudgetExceeded(RuntimeError): + """An explicit live-test bound was exceeded; source must not finish.""" + + +class BoundedStorage(harness.FilesystemBackend): + """Cap all writes, even repeated MIME reads, before retaining another blob.""" + + def __init__(self, root): + super().__init__(root) + self.written_bytes = 0 + + async def write_file(self, path, content): + if self.written_bytes + len(content) > 256 * 1024 * 1024: + raise BudgetExceeded("blob_bytes") + await super().write_file(path, content) + self.written_bytes += len(content) + + +async def bounded_observations(source, cursor, files, counters, record_limit): + """A breached test budget fails enumeration before scope completion.""" + async with aclosing(source.generate_observations(cursor=cursor, files=files)) as stream: + async for item in stream: + if isinstance(item, CaptureRecord): + counters["records_observed"] += 1 + if counters["records_observed"] > record_limit: + raise BudgetExceeded("records") + elif isinstance(item, StartedScope): + counters["started"] += 1 + elif isinstance(item, CompletedScope): + counters["completed"] += 1 + yield item + + +async def child(manifest): + engine = create_async_engine( + harness.test_database_url(), + connect_args={"server_settings": {"search_path": manifest["schema"]}}, + ) + sessions = async_sessionmaker(engine, expire_on_commit=False) + organization_id, sync_id = UUID(manifest["organization_id"]), UUID(manifest["sync_id"]) + job_id = uuid4() + counters = {"provider_requests": 0, "records_observed": 0, "started": 0, "completed": 0} + result = {"failed": True} + previous = {} + name = manifest["provider"] + is_calendar = name == "google_calendar" + counters["sync_token_requests"] = 0 + + async def request_hook(request): + counters["provider_requests"] += 1 + if request.url.params.get("syncToken"): + counters["sync_token_requests"] += 1 + if counters["provider_requests"] > manifest["request_limit"]: + raise BudgetExceeded("provider_requests") + + @asynccontextmanager + async def db_context(): + async with sessions() as db: + yield db + + try: + async with sessions() as db: + organization = await db.get(Organization, organization_id) + sync = await db.get(Sync, sync_id) + job = SyncJob( + id=job_id, organization_id=organization_id, sync_id=sync_id, status="pending" + ) + db.add(job) + await db.commit() + config = SyncConfig() + config.behavior.skip_guardrails = True + logger = logging.getLogger("private-lifecycle") + ctx = SimpleNamespace( + organization_id=organization_id, + sync_id=sync_id, + sync_job_id=job_id, + organization=organization, + sync=sync, + sync_job=job, + logger=logger, + collection_id=uuid4(), + source_connection_id=uuid4(), + source_short_name=name, + execution_config=config, + batch_size=10, + max_batch_latency_ms=20, + should_batch=True, + has_user_context=False, + tracking_email=None, + ) + cursor_service = SyncCursorService() + async with sessions() as db: + previous = await cursor_service.get_cursor_data(db, sync_id, ctx) + cursor = SyncCursor( + sync_id, GoogleCalendarCursor if is_calendar else GmailCursor, previous or None + ) + loaded = cursor.loaded_from_db + storage = BoundedStorage(Path(manifest["root"]) / "blobs") + files = FileService(job_id, storage, sync_id=sync_id) + files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 + fence = SimpleNamespace(organization_id=organization_id, sync_id=sync_id, job_id=job_id) + async with ( + asyncio.timeout(manifest["timeout"]), + rest_source( + name, + os.environ["LIVE_CALENDAR_ACCOUNT_ID" if is_calendar else "LIVE_GMAIL_ACCOUNT_ID"], + os.environ["LIVE_EXPECTED_EMAIL"], + os.environ["COMPOSIO_API_KEY"], + fence, + gmail_query=manifest.get("query", ""), + calendar_config=GoogleCalendarConfig.model_validate(manifest["calendar_config"]) + if is_calendar + else None, + request_hook=request_hook, + ) as (source, _), + ): + bus = FakeEventBus() + attempt = CaptureAttempt(id=uuid4(), number=1) + pipeline = CanonicalCapturePipeline( + CanonicalCaptureService(CanonicalRecordStore()), + sessions, + bus, + source.canonical_record_types, + attempt, + container_parents=getattr(source, "canonical_container_parents", {}), + ) + runtime = SyncRuntime( + source=source, + entity_tracker=EntityTracker(job_id, sync_id, logger), + cursor=cursor, + canonical_capture=pipeline, + ) + runner = SyncOrchestrator( + entity_pipeline=pipeline, + worker_pool=AsyncWorkerPool(logger=logger), + stream=AsyncSourceStream( + bounded_observations(source, cursor, files, counters, manifest["record_limit"]), + logger=logger, + ), + sync_context=ctx, + runtime=runtime, + access_control_pipeline=AsyncMock(), + event_bus=bus, + usage_checker=AsyncMock(), + usage_ledger=AsyncMock(), + sync_cursor_service=cursor_service, + state_machine=SyncJobStateMachine(SyncJobRepository(), bus), + lifecycle_data=None, + sync_state_machine=AsyncMock(), + ) + # Only connection ownership and external telemetry are replaced; state + # transitions, capture, reconciliation and cursor writes are production code. + with ( + patch("airweave.domains.syncs.jobs.state_machine.get_db_context", db_context), + patch("airweave.domains.sync_pipeline.orchestrator.business_events"), + ): + await runner.run() + await engine.dispose() + async with sessions() as db: + saved = await db.scalar( + select(StoredCursor.cursor_data).where(StoredCursor.sync_id == sync_id) + ) + status = await db.scalar(select(SyncJob.status).where(SyncJob.id == job_id)) + sequence = await db.scalar( + select(Sync.observed_change_sequence).where(Sync.id == sync_id) + ) + rows = list((await db.scalars(select(Entity).where(Entity.sync_id == sync_id))).all()) + assert status == "completed" + if is_calendar: + calendar_id = manifest["calendar_config"]["calendar_ids"][0] + assert set(saved["calendar_tokens"]) == {calendar_id} + assert saved["calendar_tokens"][calendar_id] + assert set(saved["occurrence_coverage"]) == {calendar_id} + expected_scopes = 2 if loaded else 3 + assert counters["started"] == counters["completed"] == expected_scopes + assert bool(counters["sync_token_requests"]) == loaded + else: + assert saved["canonical_query"] == manifest["query"] + assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 + assert saved["canonical_checkpoint"] == { + "writer_attempt_id": str(attempt.id), + "observed_change_sequence": sequence, + } + assert saved != previous + visible = [r for r in rows if r.deleted_at is None and r.source_payload is not None] + digest = hashlib.sha256( + json.dumps( + sorted( + ( + r.entity_definition_short_name, + r.container_id or "", + r.native_id, + r.record_revision, + r.source_payload, + ) + for r in visible + ), + sort_keys=True, + default=str, + ).encode() + ).hexdigest() + result = { + "failed": False, + **counters, + "stored_records": len(visible), + "partial_records": sum(r.completeness != "complete" for r in visible), + "blob_bytes_written": storage.written_bytes, + "checkpoint_saved": True, + "observed_change_sequence": sequence, + "loaded_durable_checkpoint": loaded, + "job_status": status, + "payload_revision_digest": digest, + "mode": ("calendar_incremental" if loaded else "calendar_initial") + if is_calendar + else "filtered_full_reconciliation", + "all_day_occurrences": sum( + r.entity_definition_short_name == "event_occurrence" + and "date" in r.source_payload.get("start", {}) + for r in visible + ), + } + except Exception as error: + async with sessions() as db: + saved = await db.scalar( + select(StoredCursor.cursor_data).where(StoredCursor.sync_id == sync_id) + ) + status = await db.scalar(select(SyncJob.status).where(SyncJob.id == job_id)) + result = { + "failed": True, + "error_type": type(error).__name__, + **counters, + "checkpoint_unchanged": saved == (previous or None), + "job_status": status, + } + finally: + await engine.dispose() + print(json.dumps(result), flush=True) + return 1 if result["failed"] else 0 + + +async def main(): + name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") + if name not in {"gmail", "google_calendar"}: + raise ValueError("Unsupported lifecycle provider") + url = harness.test_database_url() + schema = "canonical_live_" + uuid4().hex + admin = create_async_engine(url) + engine = create_async_engine(url, connect_args={"server_settings": {"search_path": schema}}) + created = False + old_umask = os.umask(0o077) + results = [] + try: + async with admin.begin() as connection: + await connection.execute(text(f'CREATE SCHEMA "{schema}"')) + created = True + async with engine.begin() as connection: + for migration in ( + "0000_baseline.py", + "0001_canonical_records.py", + "0002_projection_publication.py", + "0003_mail_thread_index.py", + "0004_projection_generation.py", + ): + await connection.run_sync(harness.migrate, migration) + organization_id, sync_id = uuid4(), uuid4() + sessions = async_sessionmaker(engine, expire_on_commit=False) + async with sessions() as db: + db.add(Organization(id=organization_id, name="Private provider lifecycle proof")) + await db.flush() + db.add( + Sync( + id=sync_id, + organization_id=organization_id, + name="Fixed selected provider scope", + ) + ) + await db.commit() + with TemporaryDirectory(prefix="airweave-private-live-") as directory: + root = Path(directory) + end = int(time.time()) + manifest = { + "provider": name, + "request_limit": 100 if name == "google_calendar" else 600, + "record_limit": 10000 if name == "google_calendar" else 250, + "timeout": 180 if name == "google_calendar" else 600, + "schema": schema, + "root": str(root), + "organization_id": str(organization_id), + "sync_id": str(sync_id), + "query": f"after:{end - 7 * 86400} before:{end}", + } + if name == "google_calendar": + now = datetime.now(timezone.utc).replace(microsecond=0) + manifest["calendar_config"] = { + "calendar_ids": [os.environ["LIVE_EXPECTED_EMAIL"]], + "occurrence_window": { + "start": now.isoformat(), + "end": (now + timedelta(days=7)).isoformat(), + }, + } + path = root / "manifest.json" + path.write_text(json.dumps(manifest)) + for _ in range(2): + process = await asyncio.create_subprocess_exec( + sys.executable, + str(Path(__file__).absolute()), + str(path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, _ = await asyncio.wait_for( + process.communicate(), timeout=manifest["timeout"] + 60 + ) + finally: + if process.returncode is None: + process.kill() + await process.wait() + lines = [ + json.loads(line) + for line in stdout.decode().splitlines() + if line.startswith("{") + ] + if not lines: + raise RuntimeError("Child failed without sanitized result") + results.append(lines[-1]) + print(json.dumps({"run": len(results), **lines[-1]}), flush=True) + if process.returncode: + break + finally: + os.umask(old_umask) + await engine.dispose() + if created: + async with admin.begin() as connection: + await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + await admin.dispose() + print( + json.dumps( + { + "runs": len(results), + "schema_removed": True, + "blob_directory_removed": True, + "identical_second_read": len(results) == 2 + and results[0].get("payload_revision_digest") + == results[1].get("payload_revision_digest"), + } + ) + ) + return 0 if len(results) == 2 and not any(r["failed"] for r in results) else 1 + + +if __name__ == "__main__": + try: + if len(sys.argv) == 2: + manifest = json.loads(Path(sys.argv[1]).read_text()) + harness.StoragePaths.TEMP_PROCESSING = str(Path(manifest["root"]) / "downloads") + code = asyncio.run(child(manifest)) + else: + code = asyncio.run(main()) + except Exception as error: + print(json.dumps({"failed": True, "error_type": type(error).__name__}), flush=True) + code = 1 + raise SystemExit(code) diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index b69f63de31..796e8bfdfe 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -35,6 +35,8 @@ async def verify_rest_identity(name, source, expected_email): "https://www.googleapis.com/calendar/v3/users/me/calendarList/primary" ) email = profile.get("id", "") + if profile.get("primary") is not True: + raise ValueError("Calendar identity response did not identify the primary calendar") elif name == "slack": identity = await source._get("https://slack.com/api/auth.test") if not identity.get("team_id") or not identity.get("user_id"): @@ -55,19 +57,30 @@ async def verify_rest_identity(name, source, expected_email): @asynccontextmanager -async def rest_source(name, account, expected_email, key, fence): +async def rest_source( + name, + account, + expected_email, + key, + fence, + *, + gmail_query="newer_than:7d smaller:5M", + request_hook=None, + calendar_config=None, +): host = {"gmail": "gmail.googleapis.com", "slack": "slack.com"}.get(name, "www.googleapis.com") auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) transport.MAX_BINARY_BYTES = 10 * 1024 * 1024 - async with httpx.AsyncClient(transport=transport, timeout=180) as client: + hooks = {"request": [request_hook]} if request_hook is not None else {} + async with httpx.AsyncClient(transport=transport, timeout=180, event_hooks=hooks) as client: wrapped = AirweaveHttpClient( client, fence.organization_id, name, feature_flag_enabled=False ) source_type, config = { - "gmail": (GmailSource, GmailConfig(gmail_query="newer_than:7d smaller:5M")), + "gmail": (GmailSource, GmailConfig(gmail_query=gmail_query)), "google_drive": (GoogleDriveSource, GoogleDriveConfig()), - "google_calendar": (GoogleCalendarSource, GoogleCalendarConfig()), + "google_calendar": (GoogleCalendarSource, calendar_config or GoogleCalendarConfig()), "slack": (SlackSource, SlackConfig()), }[name] source = await source_type.create( From 086577d1b50144079f5fc4a469453d25f9804827 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:49:22 -0700 Subject: [PATCH 022/266] Reuse private HTTP handoff for actual Calendar consumer proof --- backend/tests/live/README.md | 22 +++++++ backend/tests/live/almanac_handoff.py | 80 +++++++++++++++++++++++- backend/tests/live/provider_lifecycle.py | 14 +++++ 3 files changed, 114 insertions(+), 2 deletions(-) diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 0d11a48139..2b1ea34252 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -155,3 +155,25 @@ payload/revision digest was identical. Zero partial records or blobs. Private schema/files were removed. No all-day or DST case was observed; those semantics remain covered separately by synthetic tests. Calendar limits are 100 provider requests, 10,000 observations and 180 seconds per process. + +### Calendar consumer over real HTTP + +Set `LIVE_VERIFY_CALENDAR_CONSUMER=1` with the explicit `LIVE_ALMANAC_ROOT` and +`LIVE_ALMANAC_PYTHON` paths when running the selected Calendar lifecycle. The +second process uses the shared `almanac_handoff.run_consumer` loopback server and +the real `/sync` routers. A fresh Almanac process receives a private0600 manifest +and an ephemeral service key; provider credentials are excluded from its environment. +The temporary source connection exists only in the disposable PostgreSQL schema. +This does not verify hosted authentication, production bindings or Temporal workers. + +Verified September30:21 actual stored expanded events crossed this boundary in5 +pages (limit5), retaining real occurrence/calendar/account identities, title, +description, scheduling times, status, recurring master and coverage metadata. +Wrong-source reads were denied and an uncaptured2001 range returned the explicit +product error. Both capture jobs completed with4 provider requests each,190 initial +observations and22 second-run observations,183 active records, unchanged journal +sequence190 and unchanged native payload/revision digest. The second fresh process +loaded its durable checkpoint and made one actual master `syncToken` request. +No all-day events occurred in this sample; DST/all-day correctness remains fixture +coverage. The private schema and files were removed; an independent PostgreSQL +check found zero remaining `canonical_live_*` schemas. No provider writes occurred. diff --git a/backend/tests/live/almanac_handoff.py b/backend/tests/live/almanac_handoff.py index 541a5bf16f..0a590661c2 100644 --- a/backend/tests/live/almanac_handoff.py +++ b/backend/tests/live/almanac_handoff.py @@ -97,6 +97,24 @@ async def verify_almanac_reader( thread_id = candidates[0].payload["threadId"] selected = [r for r in records if r.payload.get("threadId") == thread_id] expected = [await expectation(record, storage) for record in selected] + return await run_consumer( + sessions=sessions, + organization_id=organization_id, + sync_id=sync_id, + root=root, + storage=storage, + query=query, + consumer=consumer, + repository=repository, + python=python, + extra={"thread_id": thread_id, "expected": expected}, + ) + + +async def run_consumer( + *, sessions, organization_id, sync_id, root, storage, query, consumer, repository, python, extra +): + """Shared private loopback API/process boundary for provider-specific assertions.""" probe_key = secrets.token_urlsafe(32) app = FastAPI() app.include_router(api_router) @@ -139,8 +157,7 @@ async def database(): "sync_id": str(sync_id), "account_id": str(uuid4()), "source_connection_id": str(uuid4()), - "thread_id": thread_id, - "expected": expected, + **extra, } ) ) @@ -203,3 +220,62 @@ def consumer_result(output, returncode): print(json.dumps({"consumer_failed_stage": safe_stage}), flush=True) raise RuntimeError("Almanac consumer proof failed; private diagnostics suppressed") return result + + +async def verify_calendar_reader( + *, sessions, organization_id, sync_id, root, storage, window, calendar_id +): + """Hand actual committed occurrences to the Almanac reader through the same HTTP server.""" + from airweave.domains.entities.canonical.calendar_query import ( + CalendarRange, + CalendarRangeService, + ) + from airweave.domains.entities.canonical.query import CanonicalQueryService + from airweave.domains.entities.canonical.query_store import CanonicalQueryStore + from airweave.domains.entities.canonical.store import CanonicalRecordStore + from airweave.models.source_connection import SourceConnection + + query = CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "private-test") + request = CalendarRange(start=window["start"], end=window["end"], timezone="UTC", limit=250) + async with sessions() as db: + page = await CalendarRangeService("private-test").read( + db, organization_id, sync_id, calendar_id, request + ) + if page.has_more: + raise ValueError("Consumer proof exceeds bounded expected-page limit") + connection_id = uuid4() + async with sessions() as db: + db.add( + SourceConnection( + id=connection_id, + organization_id=organization_id, + sync_id=sync_id, + name="Private calendar consumer", + short_name="google_calendar", + is_authenticated=True, + ) + ) + await db.commit() + repository = Path(os.environ["LIVE_ALMANAC_ROOT"]).resolve() + python = Path(os.environ["LIVE_ALMANAC_PYTHON"]).absolute() + consumer = repository / "backend/tests/live/read_owned_calendar.py" + if not python.is_file() or not consumer.is_file(): + raise ValueError("Explicit Almanac consumer checkout and Python required") + return await run_consumer( + sessions=sessions, + organization_id=organization_id, + sync_id=sync_id, + root=root, + storage=storage, + query=query, + consumer=consumer, + repository=repository, + python=python, + extra={ + "source_connection_id": str(connection_id), + "calendar_id": calendar_id, + "window": window, + "expected": [record.payload for record in page.records], + "coverage": page.coverage.model_dump(mode="json"), + }, + ) diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index c1ec5eff75..89416fbf21 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -249,8 +249,22 @@ async def db_context(): default=str, ).encode() ).hexdigest() + consumer_result = None + if is_calendar and loaded and os.environ.get("LIVE_VERIFY_CALENDAR_CONSUMER") == "1": + from almanac_handoff import verify_calendar_reader + + consumer_result = await verify_calendar_reader( + sessions=sessions, + organization_id=organization_id, + sync_id=sync_id, + root=Path(manifest["root"]), + storage=storage, + window=manifest["calendar_config"]["occurrence_window"], + calendar_id=manifest["calendar_config"]["calendar_ids"][0], + ) result = { "failed": False, + "almanac_consumer": consumer_result, **counters, "stored_records": len(visible), "partial_records": sum(r.completeness != "complete" for r in visible), From ad3ee373b59481450d7d5c49ecbca881a76f9603 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:52:01 -0700 Subject: [PATCH 023/266] Exercise captured blob projection and owned search against real Vespa --- Makefile | 2 +- .../tests/test_real_vespa_projection.py | 210 ++++++++++++++++++ .../entities/canonical/tests/vespa_helpers.py | 47 ++++ backend/tests/integration/test_owned_vespa.py | 41 +--- 4 files changed, 260 insertions(+), 40 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/vespa_helpers.py diff --git a/Makefile b/Makefile index 85525595e1..19e249030b 100644 --- a/Makefile +++ b/Makefile @@ -39,6 +39,6 @@ test-auth: test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py check: test-store test-capture test-search test-auth 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..0e177a5d3e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py @@ -0,0 +1,210 @@ +"""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, 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): + 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_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] + assert "fundraising meeting" in " ".join(result.json()["items"][0]["excerpts"]) + 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/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/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index 6696576c0e..cd5fe2388a 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -4,13 +4,8 @@ and is unsuitable for an existing application or customer index. """ -import asyncio -import io import json import os -import time -import zipfile -from pathlib import Path from uuid import uuid4 import httpx @@ -19,6 +14,7 @@ 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.types.embeddings import QueryEmbeddings @@ -29,44 +25,11 @@ ) -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}") - - @pytest.mark.asyncio async def test_real_schema_retrieval_collection_isolation_and_delete(): """Exercise real query compilation/rank profiles; vectors are fixed fixtures.""" - package = Path(__file__).resolve().parents[3] / "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") - ) async with httpx.AsyncClient(timeout=120) as http: - 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") + await deploy_schema(http) collection, other = str(uuid4()), str(uuid4()) sync = str(uuid4()) vector = [1.0] + [0.0] * 383 From 264609c7f2c7446b35ddb3676e47cb59048c8195 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:56:25 -0700 Subject: [PATCH 024/266] test: bound full Drive lifecycle and report incomplete capture --- backend/tests/live/README.md | 20 +++++ .../evidence/drive-lifecycle-20260930.json | 23 +++++ backend/tests/live/provider_lifecycle.py | 84 ++++++++++++++++--- 3 files changed, 116 insertions(+), 11 deletions(-) create mode 100644 backend/tests/live/evidence/drive-lifecycle-20260930.json diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 2b1ea34252..4ede48da1a 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -177,3 +177,23 @@ loaded its durable checkpoint and made one actual master `syncToken` request. No all-day events occurred in this sample; DST/all-day correctness remains fixture coverage. The private schema and files were removed; an independent PostgreSQL check found zero remaining `canonical_live_*` schemas. No provider writes occurred. + +Drive mode: `LIVE_LIFECYCLE_PROVIDER=google_drive` with `LIVE_DRIVE_ACCOUNT_ID` +uses the source's entire accessible `allDrives` corpus; no query, path selection, +or hidden sampling filter is inserted. Per process: 2,000 observations, +250 provider requests, 128 MiB total retained blob writes, 10 MiB per blob and +300 seconds. A successful initial run must exhaust enumeration plus changes +replay and commit its canonical page token. A second fresh process must actually +request `/changes` with that saved token and emit no full-rescan markers; an +expired token/reset is not reported as successful incremental proof. Failure +stops the attempt and reports whether the previous checkpoint stayed unchanged. + +Drive attempt 2026-09-30: the full accessible scope stopped at the configured +10 MiB per-file transport ceiling (`Proxy download exceeds the file size limit`) +during a binary download. The initial attempt and one sanitized diagnostic retry +each made 11 provider requests and yielded 4 observations, with one start marker +and **no** completed scope. The real job became failed; the checkpoint remained +unchanged. No second incremental process was started. This proves the bounded +abort path, not complete Drive synchronization. Both attempts removed their +private schema/files. Thirteen Drive source/content tests passed; a synthetic +storage-budget check also rejected an over-budget write before retaining it. diff --git a/backend/tests/live/evidence/drive-lifecycle-20260930.json b/backend/tests/live/evidence/drive-lifecycle-20260930.json new file mode 100644 index 0000000000..4a7ff8f85d --- /dev/null +++ b/backend/tests/live/evidence/drive-lifecycle-20260930.json @@ -0,0 +1,23 @@ +{ + "date": "2026-09-30", + "scope": "full accessible allDrives corpus; no hidden filter", + "provider_writes": false, + "limits": {"records": 2000, "provider_requests": 250, "total_blob_bytes": 134217728, "per_blob_bytes": 10485760, "seconds": 300}, + "attempts": 2, + "second_attempt_purpose": "sanitized diagnostic retry of same full scope", + "result_each_attempt": { + "failed": true, + "error_type": "ComposioProxyError", + "provider_requests": 11, + "records_observed": 4, + "started": 1, + "completed": 0, + "checkpoint_unchanged": true, + "job_status": "failed" + }, + "diagnostic_retry": {"safe_reason": "Proxy download exceeds the file size limit", "last_operation": "download"}, + "fresh_process_incremental_run": false, + "full_scope_verified": false, + "cleanup_each_attempt": {"schema_removed": true, "blob_directory_removed": true}, + "drive_source_tests_passed": 13 +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 89416fbf21..20311ce8fa 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -1,4 +1,4 @@ -"""Opt-in fixed Gmail reconciliation or selected Calendar incremental lifecycle. +"""Opt-in Gmail, selected Calendar, or full accessible Drive lifecycle proof. Two fresh processes; no Temporal worker, hosted factory, auth or product binding claim. Parent owns private schema/files cleanup on success and failure. No provider writes. @@ -9,6 +9,7 @@ import json import logging import os +import re import sys import time from contextlib import aclosing, asynccontextmanager @@ -46,6 +47,8 @@ from airweave.platform.configs.config import GoogleCalendarConfig from airweave.platform.cursors.gmail import GmailCursor from airweave.platform.cursors.google_calendar import GoogleCalendarCursor +from airweave.platform.cursors.google_drive import GoogleDriveCursor +from airweave.platform.http_client.composio_transport import ComposioProxyError class BudgetExceeded(RuntimeError): @@ -55,12 +58,13 @@ class BudgetExceeded(RuntimeError): class BoundedStorage(harness.FilesystemBackend): """Cap all writes, even repeated MIME reads, before retaining another blob.""" - def __init__(self, root): + def __init__(self, root, byte_limit): super().__init__(root) self.written_bytes = 0 + self.byte_limit = byte_limit async def write_file(self, path, content): - if self.written_bytes + len(content) > 256 * 1024 * 1024: + if self.written_bytes + len(content) > self.byte_limit: raise BudgetExceeded("blob_bytes") await super().write_file(path, content) self.written_bytes += len(content) @@ -81,6 +85,25 @@ async def bounded_observations(source, cursor, files, counters, record_limit): yield item +def safe_failure_reason(error): + """Only local fixed diagnostics; never emit provider payloads, URLs or IDs.""" + text = str(error) + if isinstance(error, BudgetExceeded) and text in {"records", "provider_requests", "blob_bytes"}: + return text + if isinstance(error, ComposioProxyError) and ( + text + in { + "Proxy download exceeds the file size limit", + "Composio response exceeds the file size limit", + "Invalid Composio response envelope", + "Temporary download request failed", + } + or re.fullmatch(r"(?:Composio proxy|Temporary download) returned HTTP [0-9]{3}", text) + ): + return text + return None + + async def child(manifest): engine = create_async_engine( harness.test_database_url(), @@ -95,11 +118,25 @@ async def child(manifest): name = manifest["provider"] is_calendar = name == "google_calendar" counters["sync_token_requests"] = 0 + counters["resumed_changes_requests"] = 0 + last_operation = "identity" async def request_hook(request): + nonlocal last_operation + last_operation = ( + "export" + if request.url.path.endswith("/export") + else ("download" if request.url.params.get("alt") == "media" else "metadata") + ) counters["provider_requests"] += 1 if request.url.params.get("syncToken"): counters["sync_token_requests"] += 1 + if ( + request.url.path.endswith("/changes") + and previous.get("canonical_page_token") + and request.url.params.get("pageToken") == previous["canonical_page_token"] + ): + counters["resumed_changes_requests"] += 1 if counters["provider_requests"] > manifest["request_limit"]: raise BudgetExceeded("provider_requests") @@ -142,10 +179,16 @@ async def db_context(): async with sessions() as db: previous = await cursor_service.get_cursor_data(db, sync_id, ctx) cursor = SyncCursor( - sync_id, GoogleCalendarCursor if is_calendar else GmailCursor, previous or None + sync_id, + { + "gmail": GmailCursor, + "google_calendar": GoogleCalendarCursor, + "google_drive": GoogleDriveCursor, + }[name], + previous or None, ) loaded = cursor.loaded_from_db - storage = BoundedStorage(Path(manifest["root"]) / "blobs") + storage = BoundedStorage(Path(manifest["root"]) / "blobs", manifest["blob_byte_limit"]) files = FileService(job_id, storage, sync_id=sync_id) files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 fence = SimpleNamespace(organization_id=organization_id, sync_id=sync_id, job_id=job_id) @@ -153,7 +196,13 @@ async def db_context(): asyncio.timeout(manifest["timeout"]), rest_source( name, - os.environ["LIVE_CALENDAR_ACCOUNT_ID" if is_calendar else "LIVE_GMAIL_ACCOUNT_ID"], + os.environ[ + { + "gmail": "LIVE_GMAIL_ACCOUNT_ID", + "google_calendar": "LIVE_CALENDAR_ACCOUNT_ID", + "google_drive": "LIVE_DRIVE_ACCOUNT_ID", + }[name] + ], os.environ["LIVE_EXPECTED_EMAIL"], os.environ["COMPOSIO_API_KEY"], fence, @@ -224,6 +273,10 @@ async def db_context(): expected_scopes = 2 if loaded else 3 assert counters["started"] == counters["completed"] == expected_scopes assert bool(counters["sync_token_requests"]) == loaded + elif name == "google_drive": + assert saved["canonical_page_token"] + assert counters["started"] == counters["completed"] == (0 if loaded else 1) + assert bool(counters["resumed_changes_requests"]) == loaded else: assert saved["canonical_query"] == manifest["query"] assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 @@ -276,7 +329,11 @@ async def db_context(): "payload_revision_digest": digest, "mode": ("calendar_incremental" if loaded else "calendar_initial") if is_calendar - else "filtered_full_reconciliation", + else ( + ("drive_incremental" if loaded else "drive_initial") + if name == "google_drive" + else "filtered_full_reconciliation" + ), "all_day_occurrences": sum( r.entity_definition_short_name == "event_occurrence" and "date" in r.source_payload.get("start", {}) @@ -292,6 +349,8 @@ async def db_context(): result = { "failed": True, "error_type": type(error).__name__, + "safe_reason": safe_failure_reason(error), + "last_operation": last_operation, **counters, "checkpoint_unchanged": saved == (previous or None), "job_status": status, @@ -304,7 +363,7 @@ async def db_context(): async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") - if name not in {"gmail", "google_calendar"}: + if name not in {"gmail", "google_calendar", "google_drive"}: raise ValueError("Unsupported lifecycle provider") url = harness.test_database_url() schema = "canonical_live_" + uuid4().hex @@ -344,9 +403,12 @@ async def main(): end = int(time.time()) manifest = { "provider": name, - "request_limit": 100 if name == "google_calendar" else 600, - "record_limit": 10000 if name == "google_calendar" else 250, - "timeout": 180 if name == "google_calendar" else 600, + "request_limit": {"gmail": 600, "google_calendar": 100, "google_drive": 250}[name], + "record_limit": {"gmail": 250, "google_calendar": 10000, "google_drive": 2000}[ + name + ], + "timeout": {"gmail": 600, "google_calendar": 180, "google_drive": 300}[name], + "blob_byte_limit": (128 if name == "google_drive" else 256) * 1024 * 1024, "schema": schema, "root": str(root), "organization_id": str(organization_id), From 9aefba3aa4d45803d51d4e5b4e304922f4023a6b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 03:56:32 -0700 Subject: [PATCH 025/266] Expose the verified container build through the maintainer Makefile --- .github/workflows/owned-source-store.yml | 2 +- Makefile | 7 ++++++- deploy/README.md | 7 +++++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 1edf3917fe..d5c2ae102a 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -86,7 +86,7 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd - name: Build the deployable image without publishing - run: docker build --tag almanac-source-store:ci backend + run: make build IMAGE=almanac-source-store:ci - name: Verify packaged Python and entrypoint syntax run: | docker run --rm --entrypoint python almanac-source-store:ci -m compileall -q airweave diff --git a/Makefile b/Makefile index 19e249030b..94a8b7a525 100644 --- a/Makefile +++ b/Makefile @@ -2,11 +2,13 @@ BACKEND := backend PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest +IMAGE ?= almanac-source-store:local -.PHONY: help setup check test-store test-capture test-search test-auth test-index validate-deploy +.PHONY: help setup check test-store test-capture test-search test-auth test-index validate-deploy build help: @echo 'setup Install the locked backend and development dependencies' + @echo 'build Build the backend image locally; set IMAGE to choose its tag' @echo 'test-store Verify record transactions and publication on disposable PostgreSQL' @echo 'test-capture Verify source capture, transport and blob handling' @echo 'test-search Verify search services and visibility behavior' @@ -18,6 +20,9 @@ help: @echo 'test-store requires CANONICAL_TEST_DATABASE_URL; it creates and removes test schemas.' @echo 'See deploy/README.md for staging prerequisites and explicit deployment steps.' +build: + docker build --tag "$(IMAGE)" $(BACKEND) + validate-deploy: porter apply validate -f porter.yaml diff --git a/deploy/README.md b/deploy/README.md index 70a5da55e4..9360a9fc27 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -50,6 +50,13 @@ Semantic quality and resource use still need measurement. Do not claim generativ search is operational without separately configured model credentials and a live search test. +## Local image build + +`make build` builds the same backend Dockerfile used by CI, without publishing or +starting services. Use `make build IMAGE=almanac-source-store:COMMIT` to choose a +tag. CI separately checks packaged Python and entrypoint syntax; a successful +build does not establish dependency startup or a deployment. + ## Database setup and deployment Validate locally from the repository root: From e6626ca8a8e3defe080f1d1c2ed11adf3b33b210 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:01:42 -0700 Subject: [PATCH 026/266] test: verify full Drive capture and fresh changes-token resume --- backend/tests/live/README.md | 22 ++++++++++++++++--- ...e-lifecycle-production-limit-20260930.json | 16 ++++++++++++++ backend/tests/live/provider_lifecycle.py | 22 ++++++++++++++++--- backend/tests/live/provider_sample.py | 3 ++- 4 files changed, 56 insertions(+), 7 deletions(-) create mode 100644 backend/tests/live/evidence/drive-lifecycle-production-limit-20260930.json diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 4ede48da1a..19609a2fb7 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -180,9 +180,12 @@ check found zero remaining `canonical_live_*` schemas. No provider writes occurr Drive mode: `LIVE_LIFECYCLE_PROVIDER=google_drive` with `LIVE_DRIVE_ACCOUNT_ID` uses the source's entire accessible `allDrives` corpus; no query, path selection, -or hidden sampling filter is inserted. Per process: 2,000 observations, -250 provider requests, 128 MiB total retained blob writes, 10 MiB per blob and -300 seconds. A successful initial run must exhaust enumeration plus changes +or hidden sampling filter is inserted. The initial verification budget was +2,000 observations, 250 provider requests, 128 MiB total retained blob writes, +10 MiB per blob and 300 seconds. The approved follow-up budget uses the source +production limit of 200 MiB per blob and 512 MiB total; record/request/time +limits stay unchanged. Drive checks for at least 2 GiB free disk before each +process starts. A successful initial run must exhaust enumeration plus changes replay and commit its canonical page token. A second fresh process must actually request `/changes` with that saved token and emit no full-rescan markers; an expired token/reset is not reported as successful incremental proof. Failure @@ -197,3 +200,16 @@ unchanged. No second incremental process was started. This proves the bounded abort path, not complete Drive synchronization. Both attempts removed their private schema/files. Thirteen Drive source/content tests passed; a synthetic storage-budget check also rejected an over-budget write before retaining it. + +The approved production-size Drive retry completed on 2026-09-30. Initial run: +85 provider requests, 55 observed/stored records, three not content-complete, +207,597,309 retained blob bytes, one completed full scope and a saved checkpoint. +A fresh second process made two requests, including `/changes` with the actual +saved page token, observed zero changes and saved its new attempt checkpoint. +Change sequence stayed 55 and payload/revision digests were identical. Both jobs +completed. Free disk before the run was 9,525,235,712 bytes; private schema/files +were removed and independently checked absent. This verifies the current +accessible Drive scope and a quiet real delta run; it does not prove the three +incomplete bodies, a hosted consumer, or live induced deletions/revocations. +The earlier smaller-budget failure evidence remains unchanged. Sanitized result: +`evidence/drive-lifecycle-production-limit-20260930.json`. diff --git a/backend/tests/live/evidence/drive-lifecycle-production-limit-20260930.json b/backend/tests/live/evidence/drive-lifecycle-production-limit-20260930.json new file mode 100644 index 0000000000..1cb5fa8e8f --- /dev/null +++ b/backend/tests/live/evidence/drive-lifecycle-production-limit-20260930.json @@ -0,0 +1,16 @@ +{ + "date": "2026-09-30", + "scope": "full accessible allDrives corpus; no hidden filter", + "provider_writes": false, + "limits": {"records": 2000, "provider_requests": 250, "total_blob_bytes": 536870912, "per_blob_bytes": 209715200, "seconds": 300, "minimum_free_disk_bytes": 2147483648}, + "free_disk_bytes_before": 9525235712, + "runs": [ + {"mode": "drive_initial", "provider_requests": 85, "records_observed": 55, "stored_records": 55, "not_content_complete_records": 3, "blob_bytes_written": 207597309, "started": 1, "completed": 1, "checkpoint_saved": true, "observed_change_sequence": 55, "loaded_durable_checkpoint": false, "job_status": "completed", "payload_revision_digest": "5e6c392143053de4d41a79b229ea1eab2b884b98be72b28c4ce6cb9cdb4aa938"}, + {"mode": "drive_incremental", "provider_requests": 2, "records_observed": 0, "stored_records": 55, "not_content_complete_records": 3, "blob_bytes_written": 0, "started": 0, "completed": 0, "resumed_changes_requests": 1, "checkpoint_saved": true, "observed_change_sequence": 55, "loaded_durable_checkpoint": true, "job_status": "completed", "payload_revision_digest": "5e6c392143053de4d41a79b229ea1eab2b884b98be72b28c4ce6cb9cdb4aa938"} + ], + "identical_second_read": true, + "cleanup": {"schema_removed": true, "blob_directory_removed": true}, + "verified": ["real local orchestrator", "real job state transitions", "durable canonical checkpoint", "fresh-process native changes-token resume", "free disk reserve boundaries"], + "not_verified": ["content completeness for three explicitly incomplete records", "hosted auth or real product binding", "Temporal worker or hosted factory", "Drive product HTTP consumer", "live induced deletion/revocation", "search quality"], + "earlier_smaller_budget_failure": "drive-lifecycle-20260930.json" +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 20311ce8fa..ccb24a4032 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -10,12 +10,13 @@ import logging import os import re +import shutil import sys import time from contextlib import aclosing, asynccontextmanager from datetime import datetime, timedelta, timezone from pathlib import Path -from tempfile import TemporaryDirectory +from tempfile import TemporaryDirectory, gettempdir from types import SimpleNamespace from unittest.mock import AsyncMock, patch from uuid import UUID, uuid4 @@ -190,7 +191,7 @@ async def db_context(): loaded = cursor.loaded_from_db storage = BoundedStorage(Path(manifest["root"]) / "blobs", manifest["blob_byte_limit"]) files = FileService(job_id, storage, sync_id=sync_id) - files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 + files.MAX_FILE_SIZE_BYTES = manifest["file_byte_limit"] fence = SimpleNamespace(organization_id=organization_id, sync_id=sync_id, job_id=job_id) async with ( asyncio.timeout(manifest["timeout"]), @@ -211,6 +212,7 @@ async def db_context(): if is_calendar else None, request_hook=request_hook, + max_file_bytes=manifest["file_byte_limit"], ) as (source, _), ): bus = FakeEventBus() @@ -361,10 +363,19 @@ async def db_context(): return 1 if result["failed"] else 0 +def check_free_disk(provider): + """Keep the explicitly approved local-disk reserve before each Drive process.""" + available = shutil.disk_usage(gettempdir()).free + if provider == "google_drive" and available < 2 * 1024**3: + raise BudgetExceeded("free_disk") + return available + + async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") if name not in {"gmail", "google_calendar", "google_drive"}: raise ValueError("Unsupported lifecycle provider") + free_disk_bytes = check_free_disk(name) url = harness.test_database_url() schema = "canonical_live_" + uuid4().hex admin = create_async_engine(url) @@ -408,7 +419,10 @@ async def main(): name ], "timeout": {"gmail": 600, "google_calendar": 180, "google_drive": 300}[name], - "blob_byte_limit": (128 if name == "google_drive" else 256) * 1024 * 1024, + "blob_byte_limit": (512 if name == "google_drive" else 256) * 1024 * 1024, + "file_byte_limit": FileService.MAX_FILE_SIZE_BYTES + if name == "google_drive" + else 10 * 1024 * 1024, "schema": schema, "root": str(root), "organization_id": str(organization_id), @@ -427,6 +441,7 @@ async def main(): path = root / "manifest.json" path.write_text(json.dumps(manifest)) for _ in range(2): + check_free_disk(name) process = await asyncio.create_subprocess_exec( sys.executable, str(Path(__file__).absolute()), @@ -464,6 +479,7 @@ async def main(): json.dumps( { "runs": len(results), + "free_disk_bytes_before": free_disk_bytes, "schema_removed": True, "blob_directory_removed": True, "identical_second_read": len(results) == 2 diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index 796e8bfdfe..c988dcc1d7 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -67,11 +67,12 @@ async def rest_source( gmail_query="newer_than:7d smaller:5M", request_hook=None, calendar_config=None, + max_file_bytes=10 * 1024 * 1024, ): host = {"gmail": "gmail.googleapis.com", "slack": "slack.com"}.get(name, "www.googleapis.com") auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) - transport.MAX_BINARY_BYTES = 10 * 1024 * 1024 + transport.MAX_BINARY_BYTES = max_file_bytes hooks = {"request": [request_hook]} if request_hook is not None else {} async with httpx.AsyncClient(transport=transport, timeout=180, event_hooks=hooks) as client: wrapped = AirweaveHttpClient( From 408b4ec58b4e209b5bd77f88a7e774f43fd5ed4b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:03:42 -0700 Subject: [PATCH 027/266] Reconcile lost container membership without retaining readable metadata --- .../domains/entities/canonical/requests.py | 1 + .../domains/entities/canonical/store.py | 2 +- .../canonical/tests/test_slack_membership.py | 79 +++++++++++++++++++ .../sync_pipeline/canonical_capture.py | 7 ++ backend/airweave/platform/sources/slack.py | 3 + .../platform/sources/test_slack_capture.py | 9 ++- 6 files changed, 96 insertions(+), 5 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_membership.py diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 3002e1b509..6cfbb49c7b 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -99,6 +99,7 @@ class ReconcileScope(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) fence: WriterFence scope: CompletedScope + removal_reason: Literal["absent", "scope_removed"] = "absent" observed_at: AwareDatetime limit: int = Field(default=250, ge=1, le=500) diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 93014e38f7..65e8f9a152 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -344,7 +344,7 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re payload=entity.source_payload, payload_schema_version=entity.payload_schema_version, kind="delete", - removal_reason="absent", + removal_reason=request.removal_reason, completeness=entity.completeness, content_hash=entity.content_hash, source_created_at=entity.source_created_at, 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..55b339b37e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py @@ -0,0 +1,79 @@ +"""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.sources.slack import SlackSource + +pytestmark = pytest.mark.integration + + +@pytest.mark.parametrize("listing_fails", [False, True]) +async def test_first_completed_membership_hides_prior_partial_channel( + database, source, listing_fails +): + 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) + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + SlackSource.canonical_record_types, + CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), + SlackSource.canonical_container_parents, + ) + connector = SlackSource( + auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock() + ) + connector._get = AsyncMock( + side_effect=ConnectionError("membership page failed") if listing_fails else None, + return_value={"channels": []}, + ) + await pipeline.start(ctx) + + async def scan(): + async for item in connector.generate_observations(cursor=runtime.cursor): + await pipeline.process([item], ctx, runtime) + + 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/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index c471866741..eb2c16b282 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -216,6 +216,13 @@ async def cleanup_orphaned_entities( 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), ), ) diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 0796304bba..d78967e710 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -217,6 +217,7 @@ async def generate_observations( if node_selections: raise ValueError("Slack capture does not yet support selected conversation scopes") seen_channels: set[str] = set() + yield StartedScope(record_type="channel") async with aclosing( self._capture_pages( "conversations.list", @@ -269,6 +270,8 @@ async def generate_observations( removal_reason="access_revoked", observed_at=observed_at, ) + # Reconcile parents even if an earlier partial run saved no channel checkpoint. + yield CompletedScope(record_type="channel") if cursor is not None: cursor.update(channel_ids=sorted(seen_channels)) diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 360e1c8e19..5cbbf3e52b 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -45,7 +45,8 @@ async def get(url, params): assert records[0].payload["unknown_native_field"] == "kept" assert any(x.identity.native_id == "1.1" and x.payload["thread_ts"] == "1" for x in records) assert records[-1].completeness == "partial" - assert observations[-1] == CompletedScope(record_type="message", container_id="C1") + assert CompletedScope(record_type="message", container_id="C1") in observations + assert observations[-1] == CompletedScope(record_type="channel") assert calls[0][1]["types"] == "public_channel,private_channel,im,mpim" @@ -98,9 +99,9 @@ async def test_prior_channel_loss_requires_explicit_provider_confirmation(monkey AsyncMock(side_effect=[{"channels": []}, SlackApiError("channel_not_found")]), ) observations = [x async for x in connector.generate_observations(cursor=cursor)] - assert observations[0].kind == "delete" - assert observations[0].removal_reason == "access_revoked" - assert isinstance(observations[1], RemovedScope) + assert observations[1].kind == "delete" + assert observations[1].removal_reason == "access_revoked" + assert isinstance(observations[2], RemovedScope) assert cursor.get()["channel_ids"] == [] From 5cf771fe90e1d02f47d198a8c006db0f72c79801 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:15:58 -0700 Subject: [PATCH 028/266] fix: honor Slack retry windows across managed transport --- .../platform/http_client/retry_helpers.py | 19 +++++--- backend/airweave/platform/sources/slack.py | 14 +++--- .../platform/sources/test_slack_capture.py | 45 +++++++++++++++++++ 3 files changed, 68 insertions(+), 10 deletions(-) diff --git a/backend/airweave/platform/http_client/retry_helpers.py b/backend/airweave/platform/http_client/retry_helpers.py index 88ca044884..eca7f3a21a 100644 --- a/backend/airweave/platform/http_client/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/slack.py b/backend/airweave/platform/sources/slack.py index d78967e710..ebba1d16db 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -5,12 +5,14 @@ import re from contextlib import aclosing from datetime import datetime, timezone +from functools import partial from typing import Any, AsyncGenerator, Dict, List, Optional -from tenacity import retry, stop_after_attempt +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.entities.canonical.requests import ( CaptureRecord, @@ -38,7 +40,7 @@ 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.http_helpers import _parse_retry_after, raise_for_status from airweave.schemas.source_connection import AuthenticationMethod, OAuthType # Pattern for Slack mrkdwn special sequences: @@ -129,8 +131,8 @@ 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: @@ -151,7 +153,9 @@ async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict: 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=60) + 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", diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 5cbbf3e52b..a09951dad6 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -119,3 +119,48 @@ async def test_omitted_but_accessible_channel_does_not_advance_cursor(monkeypatc with pytest.raises(ValueError, match="omitted"): _ = [x async for x in connector.generate_observations(cursor=cursor)] assert cursor.get()["channel_ids"] == ["C1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["json", "http", "proxy"]) +async def test_rate_limit_honors_entire_wait_and_preserves_proxy_type(kind): + import httpx + + from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError + + connector = source() + request = httpx.Request("GET", "https://slack.com/api/conversations.list") + first = ( + AuthProviderRateLimitError(provider_name="composio", retry_after=240) + if kind == "proxy" + else httpx.Response( + 429 if kind == "http" else 200, + headers={"Retry-After": "240"}, + json={"ok": False, "error": "ratelimited"}, + request=request, + ) + ) + connector.http_client.get = AsyncMock( + side_effect=[first, httpx.Response(200, json={"ok": True}, request=request)] + ) + waits = AsyncMock() + get = connector._get.retry_with(sleep=waits) + assert await get(connector, str(request.url)) == {"ok": True} + assert float(waits.call_args.args[0]) == 240 + assert connector.http_client.get.await_count == 2 + + +@pytest.mark.asyncio +async def test_proxy_rate_limit_stops_after_five_attempts(): + from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError + + connector = source() + connector.http_client.get = AsyncMock( + side_effect=AuthProviderRateLimitError(provider_name="composio", retry_after=240) + ) + waits = AsyncMock() + with pytest.raises(AuthProviderRateLimitError): + await connector._get.retry_with(sleep=waits)(connector, "https://slack.com/api/test") + assert connector.http_client.get.await_count == 5 + assert waits.await_count == 4 + assert all(float(call.args[0]) == 240 for call in waits.call_args_list) From 2ebd710ec423c3493521ea29b63d61d1a7571293 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:16:13 -0700 Subject: [PATCH 029/266] Allow source workers without optional OCR and document staging dependencies --- .../canonical/tests/test_projection_ocr.py | 76 +++++++++ .../domains/temporal/worker/__init__.py | 10 +- .../temporal/worker/tests/test_wiring.py | 8 +- .../worker/tests/test_worker_ocr_guard.py | 104 ++++++------ deploy/DEPENDENCIES.md | 152 ++++++++++++++++++ deploy/README.md | 13 ++ 6 files changed, 309 insertions(+), 54 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py create mode 100644 deploy/DEPENDENCIES.md 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..e0f8647595 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py @@ -0,0 +1,76 @@ +"""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 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 + 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(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/temporal/worker/__init__.py b/backend/airweave/domains/temporal/worker/__init__.py index a622a74416..d1c8e83e55 100644 --- a/backend/airweave/domains/temporal/worker/__init__.py +++ b/backend/airweave/domains/temporal/worker/__init__.py @@ -173,13 +173,13 @@ 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() diff --git a/backend/airweave/domains/temporal/worker/tests/test_wiring.py b/backend/airweave/domains/temporal/worker/tests/test_wiring.py index 13a59fc1d4..9d490271cc 100644 --- a/backend/airweave/domains/temporal/worker/tests/test_wiring.py +++ b/backend/airweave/domains/temporal/worker/tests/test_wiring.py @@ -32,7 +32,10 @@ def test_create_activities_returns_list(): result = create_activities() assert isinstance(result, list) - assert len(result) == 7 + assert len(result) == 9 + activity_classes = {type(method.__self__).__name__ for method in result} + assert "ProjectCanonicalRecordsActivity" in activity_classes + assert "CleanupProjectionGenerationsActivity" in activity_classes @pytest.mark.unit @@ -51,8 +54,9 @@ def test_get_workflows_returns_classes(): result = get_workflows() assert isinstance(result, list) - assert len(result) == 4 + assert len(result) == 5 class_names = [cls.__name__ for cls in result] + assert "ProjectCanonicalRecordsWorkflow" 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_ocr_guard.py b/backend/airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py index d9caf2f067..ad47e2a915 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,68 @@ -"""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 +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 +async def test_worker_composes_without_ocr_and_starts_supported_activities(monkeypatch, tmp_path): + """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(200) + if request.url.path == "/health" + 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) + monkeypatch.setattr(worker.signal, "signal", lambda *_: None) + 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/deploy/DEPENDENCIES.md b/deploy/DEPENDENCIES.md new file mode 100644 index 0000000000..f744c8b76e --- /dev/null +++ b/deploy/DEPENDENCIES.md @@ -0,0 +1,152 @@ +# Private staging dependency plan + +Reviewed against current code on 2026-09-30. No resources were provisioned by this +review. This is a proposed single-replica staging topology, not a capacity proof +or production availability design. `porter.yaml` deploys only API and worker. + +## Decision + +Keep one private dependency stack in the existing Porter cluster, with durable +PostgreSQL, Vespa and Redis volumes and canonical bytes in S3. Reuse the official +Temporal Helm chart and a supported PostgreSQL operator/add-on already available +in the cluster. Do not add a second scheduler, managed search product, GPU, +Docling service or paid OCR just to start the worker. + +Prefer an existing maintained Porter PostgreSQL add-on if it supplies backups, +restore and version pinning. Otherwise use CloudNativePG's official chart and +single-instance Cluster resource. Do not create a bespoke database controller. +Use official service images for Redis, Svix and MiniLM with small private service +manifests. A single Vespa StatefulSet with its application package and PVC is +adequate for this bounded staging trial; a new Vespa operator and multi-node +production topology are unnecessary until capacity or availability requires them. + +## Actual dependency boundary + +| Component | Why it is needed | Initial sizing proposal, not measured | +| --- | --- | --- | +| PostgreSQL | Canonical authority, auth, projection ledger, Temporal state and Svix state. Separate databases/roles; do not share the product database. | 0.5 CPU / 2 GiB; 30 GiB durable volume | +| Temporal server | Existing capture/projection/maintenance workflows. Worker cannot poll without it. | One frontend/history/matching/internal worker each; combined 1 CPU / 2 GiB initial requests | +| Redis | Auth context cache, rate limits, pub/sub; Svix delivery queue. Enable persistence for queued delivery. | 0.1 CPU / 256 MiB; 5 GiB volume | +| Vespa | Derived searchable chunks; canonical SQL remains authoritative. | 2 CPU / 8 GiB; 60 GiB volume | +| MiniLM inference | Selected dense embedder. API and worker both check `/health` during composition. | 0.5 CPU / 1.5 GiB; no GPU | +| API | Existing Porter process. | 0.5 CPU / 1 GiB | +| Application worker | Capture, converters, chunker and sparse model share this process. | 1 CPU / 3–4 GiB; current manifest is only 2 GiB and needs measurement | +| Svix | External webhook delivery. Not a startup network prerequisite, but omission loses those events after retries. | 0.1 CPU / 256 MiB; state in PostgreSQL/Redis | +| S3 | Immutable captured blobs and separate database backups. | Private bucket, workload IAM, bounded retention for backups | + +The proposed requests total about 5.7 CPU and 18–19 GiB before Kubernetes and +operator overhead. They are allocation hypotheses, not minimum vendor promises. +Use bounded capture concurrency and observe peak RSS, CPU, disk and retry lag. +Do not claim the current 3 GiB application manifest covers the dependency stack. + +OCR is optional after the worker startup correction. A scanned PDF without OCR +is captured but its projection fails and stays pending; no partial text is fed. +Text-bearing PDF extraction still works locally. Images need OCR. Generative +model keys, Cohere, frontend, Connect UI and Temporal UI are also unnecessary for +owned indexed retrieval. Do not enable their product endpoints as verified. + +Svix details matter: `SvixAdapter.__init__` constructs a client without a probe; +`InMemoryEventBus.publish` logs subscriber exceptions rather than propagating +failure to capture. This is not a durable outbox. Keep Svix for faithful existing +behavior, or separately design an explicit disabled-webhooks mode. Do not point +the service at a nonexistent Svix and describe delivery as reliable. + +## Versions and bootstrap + +Use PostgreSQL **16.15**, already exercised by integration CI. Pin the actual +operator-supported image digest before deployment. Vespa CI used: + +``` +vespaengine/vespa@sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd +``` + +Deploy the repository's 384-dimensional application schema with that image before +API/worker traffic. The current Compose file's `vespa:8`, `redis:7-alpine`, Svix, +Docling and inference image tags are floating and are not a release lockfile. +Resolve and record their architecture-specific digests before applying manifests. + +The official Temporal chart's current main source identifies chart **1.7.0** and +server **1.32.0**. This is a candidate requiring published-chart retrieval and +compatibility verification, not an already tested upgrade. The repository's +Compose `auto-setup:1.24.2` is old development configuration; do not reuse it as +production bootstrap. For a fresh isolated namespace, prefer current official +server/chart with PostgreSQL default and visibility stores, no Cassandra, +Elasticsearch, Grafana, public ingress or UI. Validate rendered chart values: +current chart layout uses `server.config.persistence.datastores`, which differs +from older copied examples. Freeze `numHistoryShards` at initial creation. + +Bootstrap in this order: + +1. Private networking, storage class, node quota and workload IAM; verify the + actual total monthly cost before allocating nodes. +2. PostgreSQL with separate application, Temporal default, Temporal visibility + and Svix databases/users. Keep migration/schema credentials out of runtime + roles. Prove a backup restoration into a disposable database. +3. Redis persistence, Svix, Temporal schema jobs and namespace. Add the `SyncId` + Keyword search attribute explicitly and verify it exists. Never retain the + Compose init command's `|| true`, which can hide failed setup. +4. Vespa application deployment, inference health, model artifact cache. The + Python sparse model and semantic chunker download public artifacts; prewarm a + writable cache or verify restricted-egress startup. This is independent of OCR. +5. Application migration/bootstrap jobs, then API and worker from the same image. + Existing system maintenance schedules remain owned by existing app startup. +6. Set critical health probes to `postgres,redis,temporal`; separately exercise + Vespa query/feed and worker polling. Private service reachability from Almanac + must be explicitly arranged; cluster-private is not cross-cloud connectivity. + +No executable Helm values are included yet: cluster storage class, database +operator availability, published chart pin and IAM secret references are not +verified. Inventing them would produce an apparently complete but unusable +manifest. The next deliverable is a small dependency version lock and rendered +values against those facts, not a new infrastructure framework. + +## Cost envelope and tradeoff + +The supplied Porter baseline is **$201.22/month**, but its component breakdown +has not been verified here. Do not assume it includes dependency compute, disks, +NAT, an external load balancer or Porter usage-based fees. + +An illustrative additional t3a.2xlarge is $0.3008/hour in AWS's published us-east-1 +Linux table: **$219.58 for 730 hours**. Adding 120 GiB gp3 at $0.08/GiB-month gives +**$430.40 combined with that baseline**, before remaining charges. Reserving $40 +for object storage, backup snapshots, logs and traffic yields **$470.40**. This +is a conditional envelope, not a quote or a $500 guarantee. Root disks, NAT, +Porter fees and autoscaling must fit the remaining allowance or replace this +proposal. Do not add this node if the supplied base already includes adequate +capacity. The account quota currently reported by the parent is 8 vCPU versus +16 required by its selected cluster configuration; no quota increase or node +creation was performed here. + +T3 Unlimited is the default and can add CPU-credit charges. Standard mode avoids +that particular overrun by throttling after credits; the 40% baseline gives only +3.2 sustained vCPU on an 8-vCPU instance. Therefore this option favors modest +staging with bounded backfills, not sustained throughput. Compare a fixed-CPU +node quote using the actual existing cluster capacity before choosing. Keep +node autoscaling bounded and an AWS budget alert, while remembering an alert is +not a spending cap. + +Separate managed PostgreSQL/Redis/Temporal/search would reduce some operations, +but fitting all of them plus the existing base under $500 has not been shown. +A separate VM running Compose is cheaper to understand initially but adds another +host, backup owner and network boundary while duplicating the existing Porter +platform. It is a fallback only if the cluster's fixed charges make the reviewed +budget impossible. Single-replica stateful services accept downtime; neither +option is production HA. + +## Evidence + +Local implementation: `porter.yaml`, `docker/docker-compose.yml`, +`backend/airweave/core/container/factory.py`, `core/config/settings.py`, +`domains/embedders/config.py`, `domains/converters/registry.py`, +`domains/temporal/worker/__init__.py`, `adapters/event_bus/in_memory.py`, +`adapters/webhooks/svix.py`, `.github/workflows/owned-source-store.yml`. + +Primary external references: + +- [Temporal official Helm chart](https://github.com/temporalio/helm-charts) +- [Current Temporal chart metadata](https://raw.githubusercontent.com/temporalio/helm-charts/main/charts/temporal/Chart.yaml) +- [CloudNativePG bootstrap](https://cloudnative-pg.io/documentation/1.27/bootstrap/) +- [Vespa Kubernetes installation](https://docs.vespa.ai/en/operations/kubernetes/operations/installation.html) +- [Svix official server](https://github.com/svix/svix-webhooks) +- [AWS T3 pricing and credit behavior](https://aws.amazon.com/ec2/instance-types/t3/) +- [AWS EBS pricing](https://aws.amazon.com/ebs/pricing/) diff --git a/deploy/README.md b/deploy/README.md index 9360a9fc27..ebf57e383d 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -45,6 +45,19 @@ for exact current fields. Do not copy the main product database from Doppler. Use a least-privilege workload IAM role for this bucket and configure the Porter AWS-role connection once that role exists. Do not add static root AWS keys. +OCR is optional. The worker starts without MISTRAL_API_KEY or DOCLING_BASE_URL. +Text extraction and source capture continue; a scanned document that needs OCR +retains its captured bytes and a failed, retryable projection instead of entering +search with missing content. Configure an OCR backend to index those documents. +This does not mean every PDF requires OCR: the PDF converter first tries local +text extraction. Image conversion requires an OCR backend. + +The selected local MiniLM service is a startup dependency for both API and worker: +composition checks its health endpoint. FastEmbed and the semantic chunker also +load public model artifacts; provide a writable persistent model cache and verify +startup with the intended network policy. OCR, generative search credentials, +Cohere reranking, frontend and Temporal UI are optional for owned indexed search. + The manifest uses local embeddings to avoid assuming third-party model credits. Semantic quality and resource use still need measurement. Do not claim generative search is operational without separately configured model credentials and a live From 71ac92f3abf468b179dfd12e7fc51580d5b9c9d2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:17:34 -0700 Subject: [PATCH 030/266] test: exercise Slack and Wispr lifecycle failure boundaries --- backend/tests/live/README.md | 29 +++ .../slack-lifecycle-incomplete-20260930.json | 25 +++ .../wispr-lifecycle-incomplete-20260930.json | 51 +++++ backend/tests/live/provider_lifecycle.py | 204 ++++++++++++------ backend/tests/live/provider_sample.py | 8 +- .../platform/sources/test_wispr_capture.py | 12 ++ 6 files changed, 267 insertions(+), 62 deletions(-) create mode 100644 backend/tests/live/evidence/slack-lifecycle-incomplete-20260930.json create mode 100644 backend/tests/live/evidence/wispr-lifecycle-incomplete-20260930.json diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 19609a2fb7..0d802a7e8b 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -213,3 +213,32 @@ accessible Drive scope and a quiet real delta run; it does not prove the three incomplete bodies, a hosted consumer, or live induced deletions/revocations. The earlier smaller-budget failure evidence remains unchanged. Sanitized result: `evidence/drive-lifecycle-production-limit-20260930.json`. + + +### Slack and Wispr full lifecycle attempts + +Both modes use the real canonical orchestrator and production `cursor=None`. +Slack reconciles only completed channel/message scopes; Wispr has no deletion +scope or durable provider checkpoint. Neither is described as delta resume. +Actual network requests are metered, including Wispr session/account operations. + +The first Slack full-accessible-history attempt stopped on a rate limit after +70 requests and 457 observations, with two started scopes and zero completed. +The job failed, checkpoint stayed unchanged, and private schema/files were removed. +No repeat process started. Retry-After was not retained in that first attempt. +The subsequent approved budget is 600 requests, 10,000 observations and 1,800 +seconds per process. Slack uses its five-attempt retry policy and honors the +full Retry-After wait; the test does not shorten waits to fit its total budget. +Only status, numeric retry delay and request count are retained as diagnostics. +The longer attempt's result is pending; do not infer full capture from the budget. + +Wispr enumeration and one diagnostic retry each failed after 18 requests and nine +partial observations. `GET_MEETING` returned an opaque tool error string; the +source stopped instead of skipping that meeting. A targeted read-only diagnostic +found no structured status/code or recognized auth/not-found/rate-limit category. +The underlying provider versus intermediary cause is unknown. No complete source, +checkpoint, repeat process or deletion guarantee is claimed. Both lifecycle +attempts cleaned private storage. This is additional failure evidence alongside +the earlier successful three-record sample, not a replacement for that sample. +See `evidence/slack-lifecycle-incomplete-20260930.json` and +`evidence/wispr-lifecycle-incomplete-20260930.json`. diff --git a/backend/tests/live/evidence/slack-lifecycle-incomplete-20260930.json b/backend/tests/live/evidence/slack-lifecycle-incomplete-20260930.json new file mode 100644 index 0000000000..98640a12e7 --- /dev/null +++ b/backend/tests/live/evidence/slack-lifecycle-incomplete-20260930.json @@ -0,0 +1,25 @@ +{ + "provider": "slack", + "date": "2026-09-30", + "scope": "full accessible history without hidden filter", + "limits": { + "requests": 150, + "observations": 5000, + "seconds": 300 + }, + "attempt": { + "provider_requests": 70, + "observations": 457, + "started_scopes": 2, + "completed_scopes": 0, + "job_status": "failed", + "safe_reason": "rate_limit", + "checkpoint_unchanged": true + }, + "second_process_started": false, + "schema_removed": true, + "blob_directory_removed": true, + "full_scope_verified": false, + "retry_after_recorded": false, + "provider_writes": false +} diff --git a/backend/tests/live/evidence/wispr-lifecycle-incomplete-20260930.json b/backend/tests/live/evidence/wispr-lifecycle-incomplete-20260930.json new file mode 100644 index 0000000000..f4507c5685 --- /dev/null +++ b/backend/tests/live/evidence/wispr-lifecycle-incomplete-20260930.json @@ -0,0 +1,51 @@ +{ + "date": "2026-09-30", + "scope": "existing exposed meeting/text enumeration; no added filter", + "limits": { + "actual_http_requests": 150, + "records": 1000, + "seconds": 300 + }, + "first_attempt": { + "failed": true, + "error_type": "ValueError", + "provider_requests": 18, + "records_observed": 9, + "started": 0, + "completed": 0, + "checkpoint_unchanged": true, + "job_status": "failed" + }, + "fresh_process_repeat": false, + "scope_completion_claimed": false, + "checkpoint_support": false, + "deletion_guarantee": false, + "provider_writes": false, + "cleanup": { + "schema_removed": true, + "blob_directory_removed": true + }, + "limitation": "Initial safe diagnostics did not identify the fixed source error; a separate diagnostic attempt is required. All captured Wispr originals remain partial.", + "diagnostic_retry": { + "provider_requests": 18, + "observations": 9, + "safe_reason": "Wispr tool execution failed; capture is incomplete", + "job_status": "failed", + "started_scopes": 0, + "completed_scopes": 0, + "checkpoint_unchanged": true, + "schema_removed": true, + "blob_directory_removed": true + }, + "read_only_targeted_diagnostic": { + "requests": 18, + "operation": "get_meeting", + "error_shape": "string", + "structured_status_or_code": false, + "recognized_categories": [ + "failed" + ], + "source_parser_failure": false, + "cause_attribution": "Upstream execution boundary reported an opaque tool error; root provider versus intermediary cause is unknown." + } +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index ccb24a4032..a5cb3f6cfc 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -22,7 +22,7 @@ from uuid import UUID, uuid4 import canonical_capture as harness # settings must precede application imports -from provider_sample import rest_source +from provider_sample import rest_source, wispr_source from sqlalchemy import select, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine @@ -50,6 +50,7 @@ from airweave.platform.cursors.google_calendar import GoogleCalendarCursor from airweave.platform.cursors.google_drive import GoogleDriveCursor from airweave.platform.http_client.composio_transport import ComposioProxyError +from airweave.platform.sources.slack import SlackApiError class BudgetExceeded(RuntimeError): @@ -89,8 +90,35 @@ async def bounded_observations(source, cursor, files, counters, record_limit): def safe_failure_reason(error): """Only local fixed diagnostics; never emit provider payloads, URLs or IDs.""" text = str(error) - if isinstance(error, BudgetExceeded) and text in {"records", "provider_requests", "blob_bytes"}: + if isinstance(error, BudgetExceeded) and text in { + "records", + "provider_requests", + "blob_bytes", + "rate_limit", + }: return text + if type(error) is ValueError and text in { + "Wispr tool execution failed; capture is incomplete", + "Wispr returned incomplete or repeated pagination", + "Wispr query cap cannot be partitioned safely", + "Wispr date partition made no progress", + "Wispr returned a meeting without identity", + "Wispr meeting text is not a string", + "Wispr returned a different meeting identity", + "Wispr meeting changed during paginated capture", + "Wispr continuation made no progress", + "Wispr meeting exceeds the bounded capture page limit", + }: + return text + if isinstance(error, SlackApiError) and error.code in { + "missing_scope", + "not_in_channel", + "channel_not_found", + "is_archived", + "method_not_supported_for_channel_type", + "not_allowed_token_type", + }: + return f"Slack request failed: {error.code}" if isinstance(error, ComposioProxyError) and ( text in { @@ -105,6 +133,37 @@ def safe_failure_reason(error): return None +def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attempt, sequence): + """Match production cursor/scope capabilities; never manufacture a resume claim.""" + if name == "google_calendar": + calendar_id = manifest["calendar_config"]["calendar_ids"][0] + assert set(saved["calendar_tokens"]) == {calendar_id} + assert saved["calendar_tokens"][calendar_id] + assert set(saved["occurrence_coverage"]) == {calendar_id} + expected_scopes = 2 if loaded else 3 + assert counters["started"] == counters["completed"] == expected_scopes + assert bool(counters["sync_token_requests"]) == loaded + elif name == "google_drive": + assert saved["canonical_page_token"] + assert counters["started"] == counters["completed"] == (0 if loaded else 1) + assert bool(counters["resumed_changes_requests"]) == loaded + elif name in {"slack", "wispr"}: + assert saved is None + if name == "wispr": + assert counters["started"] == counters["completed"] == 0 + else: + assert counters["started"] == counters["completed"] > 0 + return + else: + assert saved["canonical_query"] == manifest["query"] + assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 + assert saved["canonical_checkpoint"] == { + "writer_attempt_id": str(attempt.id), + "observed_change_sequence": sequence, + } + assert saved != previous + + async def child(manifest): engine = create_async_engine( harness.test_database_url(), @@ -121,6 +180,24 @@ async def child(manifest): counters["sync_token_requests"] = 0 counters["resumed_changes_requests"] = 0 last_operation = "identity" + rate_limits = [] + + async def observe_slack_rate_limit(response): + await response.aread() + limited = response.status_code == 429 or response.json().get("error") in { + "ratelimited", + "rate_limited", + } + if limited: + from airweave.platform.sources.http_helpers import _parse_retry_after + + rate_limits.append( + { + "status": response.status_code, + "retry_after_seconds": _parse_retry_after(response, default=60), + "request_number": counters["provider_requests"], + } + ) async def request_hook(request): nonlocal last_operation @@ -129,6 +206,8 @@ async def request_hook(request): if request.url.path.endswith("/export") else ("download" if request.url.params.get("alt") == "media" else "metadata") ) + if counters["provider_requests"] >= manifest["request_limit"]: + raise BudgetExceeded("provider_requests") counters["provider_requests"] += 1 if request.url.params.get("syncToken"): counters["sync_token_requests"] += 1 @@ -138,8 +217,6 @@ async def request_hook(request): and request.url.params.get("pageToken") == previous["canonical_page_token"] ): counters["resumed_changes_requests"] += 1 - if counters["provider_requests"] > manifest["request_limit"]: - raise BudgetExceeded("provider_requests") @asynccontextmanager async def db_context(): @@ -179,42 +256,48 @@ async def db_context(): cursor_service = SyncCursorService() async with sessions() as db: previous = await cursor_service.get_cursor_data(db, sync_id, ctx) - cursor = SyncCursor( - sync_id, - { - "gmail": GmailCursor, - "google_calendar": GoogleCalendarCursor, - "google_drive": GoogleDriveCursor, - }[name], - previous or None, - ) - loaded = cursor.loaded_from_db + cursor_schema = { + "gmail": GmailCursor, + "google_calendar": GoogleCalendarCursor, + "google_drive": GoogleDriveCursor, + }.get(name) + cursor = SyncCursor(sync_id, cursor_schema, previous or None) if cursor_schema else None + loaded = cursor.loaded_from_db if cursor else False storage = BoundedStorage(Path(manifest["root"]) / "blobs", manifest["blob_byte_limit"]) files = FileService(job_id, storage, sync_id=sync_id) files.MAX_FILE_SIZE_BYTES = manifest["file_byte_limit"] fence = SimpleNamespace(organization_id=organization_id, sync_id=sync_id, job_id=job_id) - async with ( - asyncio.timeout(manifest["timeout"]), - rest_source( + key = os.environ["COMPOSIO_API_KEY"] + account_variable = { + "gmail": "LIVE_GMAIL_ACCOUNT_ID", + "google_calendar": "LIVE_CALENDAR_ACCOUNT_ID", + "google_drive": "LIVE_DRIVE_ACCOUNT_ID", + "slack": "LIVE_SLACK_ACCOUNT_ID", + "wispr": "LIVE_WISPR_ACCOUNT_ID", + }[name] + connection = ( + wispr_source(os.environ[account_variable], key, fence, request_hook=request_hook) + if name == "wispr" + else rest_source( name, - os.environ[ - { - "gmail": "LIVE_GMAIL_ACCOUNT_ID", - "google_calendar": "LIVE_CALENDAR_ACCOUNT_ID", - "google_drive": "LIVE_DRIVE_ACCOUNT_ID", - }[name] - ], + os.environ[account_variable], os.environ["LIVE_EXPECTED_EMAIL"], - os.environ["COMPOSIO_API_KEY"], + key, fence, gmail_query=manifest.get("query", ""), calendar_config=GoogleCalendarConfig.model_validate(manifest["calendar_config"]) if is_calendar else None, request_hook=request_hook, + response_hook=observe_slack_rate_limit if name == "slack" else None, max_file_bytes=manifest["file_byte_limit"], - ) as (source, _), + ) + ) + async with ( + asyncio.timeout(manifest["timeout"]), + connection as (source, identity_verification), ): + assert source.cursor_class == cursor_schema bus = FakeEventBus() attempt = CaptureAttempt(id=uuid4(), number=1) pipeline = CanonicalCapturePipeline( @@ -267,26 +350,7 @@ async def db_context(): ) rows = list((await db.scalars(select(Entity).where(Entity.sync_id == sync_id))).all()) assert status == "completed" - if is_calendar: - calendar_id = manifest["calendar_config"]["calendar_ids"][0] - assert set(saved["calendar_tokens"]) == {calendar_id} - assert saved["calendar_tokens"][calendar_id] - assert set(saved["occurrence_coverage"]) == {calendar_id} - expected_scopes = 2 if loaded else 3 - assert counters["started"] == counters["completed"] == expected_scopes - assert bool(counters["sync_token_requests"]) == loaded - elif name == "google_drive": - assert saved["canonical_page_token"] - assert counters["started"] == counters["completed"] == (0 if loaded else 1) - assert bool(counters["resumed_changes_requests"]) == loaded - else: - assert saved["canonical_query"] == manifest["query"] - assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 - assert saved["canonical_checkpoint"] == { - "writer_attempt_id": str(attempt.id), - "observed_change_sequence": sequence, - } - assert saved != previous + validate_checkpoint(name, manifest, counters, saved, previous, loaded, attempt, sequence) visible = [r for r in rows if r.deleted_at is None and r.source_payload is not None] digest = hashlib.sha256( json.dumps( @@ -319,23 +383,26 @@ async def db_context(): ) result = { "failed": False, + "provider": name, + "identity_verification": identity_verification, "almanac_consumer": consumer_result, **counters, + "rate_limits": rate_limits, "stored_records": len(visible), "partial_records": sum(r.completeness != "complete" for r in visible), "blob_bytes_written": storage.written_bytes, - "checkpoint_saved": True, + "checkpoint_saved": saved is not None, "observed_change_sequence": sequence, "loaded_durable_checkpoint": loaded, "job_status": status, "payload_revision_digest": digest, - "mode": ("calendar_incremental" if loaded else "calendar_initial") - if is_calendar - else ( - ("drive_incremental" if loaded else "drive_initial") - if name == "google_drive" - else "filtered_full_reconciliation" - ), + "mode": { + "gmail": "filtered_full_reconciliation", + "google_calendar": "calendar_incremental" if loaded else "calendar_initial", + "google_drive": "drive_incremental" if loaded else "drive_initial", + "slack": "accessible_history_full_reconciliation", + "wispr": "exposed_meeting_enumeration_no_deletion_guarantee", + }[name], "all_day_occurrences": sum( r.entity_definition_short_name == "event_occurrence" and "date" in r.source_payload.get("start", {}) @@ -354,6 +421,7 @@ async def db_context(): "safe_reason": safe_failure_reason(error), "last_operation": last_operation, **counters, + "rate_limits": rate_limits, "checkpoint_unchanged": saved == (previous or None), "job_status": status, } @@ -373,7 +441,7 @@ def check_free_disk(provider): async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") - if name not in {"gmail", "google_calendar", "google_drive"}: + if name not in {"gmail", "google_calendar", "google_drive", "slack", "wispr"}: raise ValueError("Unsupported lifecycle provider") free_disk_bytes = check_free_disk(name) url = harness.test_database_url() @@ -414,11 +482,27 @@ async def main(): end = int(time.time()) manifest = { "provider": name, - "request_limit": {"gmail": 600, "google_calendar": 100, "google_drive": 250}[name], - "record_limit": {"gmail": 250, "google_calendar": 10000, "google_drive": 2000}[ - name - ], - "timeout": {"gmail": 600, "google_calendar": 180, "google_drive": 300}[name], + "request_limit": { + "gmail": 600, + "google_calendar": 100, + "google_drive": 250, + "slack": 600, + "wispr": 150, + }[name], + "record_limit": { + "gmail": 250, + "google_calendar": 10000, + "google_drive": 2000, + "slack": 10000, + "wispr": 1000, + }[name], + "timeout": { + "gmail": 600, + "google_calendar": 180, + "google_drive": 300, + "slack": 1800, + "wispr": 300, + }[name], "blob_byte_limit": (512 if name == "google_drive" else 256) * 1024 * 1024, "file_byte_limit": FileService.MAX_FILE_SIZE_BYTES if name == "google_drive" diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index c988dcc1d7..3eb62922fa 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -66,6 +66,7 @@ async def rest_source( *, gmail_query="newer_than:7d smaller:5M", request_hook=None, + response_hook=None, calendar_config=None, max_file_bytes=10 * 1024 * 1024, ): @@ -74,6 +75,8 @@ async def rest_source( transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) transport.MAX_BINARY_BYTES = max_file_bytes hooks = {"request": [request_hook]} if request_hook is not None else {} + if response_hook is not None: + hooks["response"] = [response_hook] async with httpx.AsyncClient(transport=transport, timeout=180, event_hooks=hooks) as client: wrapped = AirweaveHttpClient( client, fence.organization_id, name, feature_flag_enabled=False @@ -92,10 +95,11 @@ async def rest_source( @asynccontextmanager -async def wispr_source(account, key, fence): +async def wispr_source(account, key, fence, *, request_hook=None): """Verify exact Composio account binding; Wispr exposes no identity profile here.""" expected_user = os.environ["LIVE_WISPR_USER_ID"] - async with httpx.AsyncClient(timeout=180) as client: + hooks = {"request": [request_hook]} if request_hook is not None else {} + async with httpx.AsyncClient(timeout=180, event_hooks=hooks) as client: metadata_response = await client.get( "https://backend.composio.dev/api/v3/connected_accounts/" + quote(account, safe=""), headers={"x-api-key": key}, diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index fd1ec6c69f..89f253b9c8 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -103,3 +103,15 @@ async def test_unpartitionable_cap_fails(monkeypatch): ) with pytest.raises(ValueError, match="cap"): _ = [row async for row in connector._list_all()] + + +@pytest.mark.asyncio +async def test_explicit_tool_failure_never_becomes_partial_success(monkeypatch): + connector = await source() + monkeypatch.setattr( + connector, + "_post", + AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": "failed"}]), + ) + with pytest.raises(ValueError, match="tool execution failed; capture is incomplete"): + await connector._meeting("m") From 0f1922b9ae347179c4528f980eb9d33b42848672 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:20:09 -0700 Subject: [PATCH 031/266] Expose verified captured Gmail thread identity in owned search --- .../canonical/tests/test_owned_search.py | 36 +++++++++++++++++++ backend/airweave/domains/search/owned.py | 9 +++++ .../airweave/domains/search/owned_models.py | 1 + 3 files changed, 46 insertions(+) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 597c5e689e..3a11f915cc 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -362,3 +362,39 @@ async def ordered(*args): assert calls == 2 and response.status_code == 200 assert response.json()["items"] == [] and response.json()["excluded_candidates"] == 1 assert "Private original text" not in response.text + + +@pytest.mark.parametrize( + "provider,thread_id,expected", + [ + ("gmail", "thread-1", "thread-1"), + ("slack", "thread-1", None), + ("gmail", None, None), + ("gmail", "bad/id", None), + ("gmail", {"bad": "value"}, None), + ], +) +async def test_email_route_uses_visible_canonical_provider_payload( + database, indexed, http_search, provider, thread_id, expected +): + from airweave.models.entity import Entity + + fence, locator, _ = indexed + client, vector, _, _, _ = http_search + async with database() as db: + await db.execute(update(SourceConnection).values(short_name=provider)) + await db.execute( + update(Entity).values( + entity_definition_short_name="message", source_payload={"threadId": thread_id} + ) + ) + await db.commit() + candidate = hit(fence, locator.encode()) + candidate.airweave_system_metadata.source_name = provider + vector.seed_results(SearchResults(results=[candidate])) + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + assert response.status_code == 200, response.text + assert response.json()["items"][0]["email_thread_id"] == expected diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 0be9a8af2e..fba0efe002 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -1,5 +1,6 @@ """Original-record indexed retrieval using the existing executor and SQL authority.""" +import re from collections import defaultdict from uuid import UUID @@ -213,6 +214,14 @@ async def _enrich( source_created_at=row.source_created_at, source_updated_at=row.source_updated_at, completeness=row.completeness, + email_thread_id=( + thread_id + if scopes[row.sync_id].short_name == "gmail" + and row.entity_definition_short_name == "message" + and isinstance(thread_id := (row.source_payload or {}).get("threadId"), str) + and re.fullmatch(r"[A-Za-z0-9_-]{1,512}", thread_id) + else None + ), ) scores[row.id] = (1 / (60 + rank), locator) excerpt = result.textual_representation[:2000] diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py index e5f479582f..cc0f13ed0c 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -50,6 +50,7 @@ class OwnedSearchHit(BaseModel): 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 From 2e7c31162ee41e5f6bbc891914c54dc42f25253e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:21:40 -0700 Subject: [PATCH 032/266] test: verify persisted service key authorization over record HTTP routes --- .../canonical/tests/test_service_key_http.py | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_service_key_http.py 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 From 834f87e1a7c8fa4d8825463dcdd93f5e0535fe57 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:45:09 -0700 Subject: [PATCH 033/266] Persist fenced whole-scope scan pages and bounded reconciliation --- .../domains/entities/canonical/scan_models.py | 94 ++++++ .../domains/entities/canonical/scan_store.py | 190 ++++++++++++ .../domains/entities/canonical/service.py | 32 ++ .../domains/entities/canonical/store.py | 16 +- .../entities/canonical/tests/conftest.py | 1 + .../entities/canonical/tests/test_scans.py | 279 ++++++++++++++++++ backend/airweave/models/__init__.py | 2 + backend/airweave/models/capture_scan.py | 36 +++ backend/alembic/versions/0005_capture_scan.py | 56 ++++ 9 files changed, 701 insertions(+), 5 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/scan_models.py create mode 100644 backend/airweave/domains/entities/canonical/scan_store.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_scans.py create mode 100644 backend/airweave/models/capture_scan.py create mode 100644 backend/alembic/versions/0005_capture_scan.py diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py new file mode 100644 index 0000000000..3c70ad7559 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -0,0 +1,94 @@ +"""Whole-scope page commits; partial ranges must not use absence reconciliation.""" + +import json +from typing import Literal +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, field_validator + +from airweave.domains.entities.canonical.models import CaptureResult +from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, WriterFence + + +class ScanVersion(BaseModel): + """The exact durable state a caller observed before external I/O.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + sweep_id: UUID + revision: int = Field(ge=1) + + +class ScanContinuation(BaseModel): + """Connector-validated opaque progress; never original content or an unbounded queue.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + value: dict[str, JsonValue] = Field(default_factory=dict) + + @field_validator("value") + @classmethod + def bounded(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + """Bound serialized state before it enters a transaction.""" + if len(json.dumps(value, allow_nan=False, ensure_ascii=False).encode()) > 65536: + raise ValueError("Scan continuation exceeds 64 KiB") + return value + + +class BeginScan(BaseModel): + """Resume a cycle, or explicitly CAS-restart its sweep after invalid continuation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + scope: CompletedScope + cycle_id: UUID + fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") + expected: ScanVersion | None = None + restart: bool = False + continuation: ScanContinuation = Field(default_factory=ScanContinuation) + + +class ScanState(BaseModel): + """Complete means successful whole-scope enumeration and absence reconciliation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + scope: CompletedScope + cycle_id: UUID + version: ScanVersion + phase: Literal["collecting", "reconciling", "complete"] + fingerprint: str + continuation: ScanContinuation + started_at: AwareDatetime + completed_at: AwareDatetime | None + + +class CommitScanPage(BaseModel): + """Records and next state become durable together under one writer fence.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + scope: CompletedScope + cycle_id: UUID + expected: ScanVersion + records: tuple[CaptureRecord, ...] = Field(max_length=500) + continuation: ScanContinuation + final: bool = False + + +class ReconcileScan(BaseModel): + """One bounded removal batch, only after the final collection page committed.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + scope: CompletedScope + cycle_id: UUID + expected: ScanVersion + observed_at: AwareDatetime + removal_reason: Literal["absent", "scope_removed"] = "absent" + limit: int = Field(default=250, ge=1, le=500) + + +class ScanResult(BaseModel): + """Acknowledged SQL state plus changes committed by this step.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + state: ScanState + capture: CaptureResult diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py new file mode 100644 index 0000000000..8de432c4ee --- /dev/null +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -0,0 +1,190 @@ +"""Scan persistence inside the existing capture transaction; no provider I/O or commits.""" + +import json +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.requests import ( + CaptureBatch, + CompletedScope, + ReconcileScope, + WriterFence, +) +from airweave.domains.entities.canonical.scan_models import ( + BeginScan, + CommitScanPage, + ReconcileScan, + ScanContinuation, + ScanResult, + ScanState, + ScanVersion, +) +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.models.capture_scan import CaptureScan + + +class ScanConflict(CanonicalStoreError): + """Reload durable progress before performing more provider I/O.""" + + code = "scan_conflict" + + +def scope_key(scope: CompletedScope) -> str: + """Distinguish null containers from empty strings without a nullable unique key.""" + return json.dumps( + [scope.record_type, scope.container_id], ensure_ascii=False, separators=(",", ":") + ) + + +def scan_state(row: CaptureScan) -> ScanState: + """Do not leak ORM state outside its transaction.""" + return ScanState( + scope=CompletedScope(record_type=row.record_type, container_id=row.container_id), + cycle_id=row.cycle_id, + version=ScanVersion(sweep_id=row.sweep_id, revision=row.revision), + phase=row.phase, + fingerprint=row.fingerprint, + continuation=ScanContinuation(value=row.continuation), + started_at=row.started_at, + completed_at=row.completed_at, + ) + + +class CanonicalScanStore: + """Reuse record capture, lock ordering and journal semantics for recoverable scans.""" + + def __init__(self, records: CanonicalRecordStore): + """The existing record store remains the only authority for record mutations.""" + self.records = records + + async def _row( + self, db: AsyncSession, fence: WriterFence, scope: CompletedScope + ) -> CaptureScan | None: + return await db.scalar( + select(CaptureScan) + .where( + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.scope_key == scope_key(scope), + ) + .with_for_update() + .execution_options(populate_existing=True) + ) + + async def read( + self, db: AsyncSession, fence: WriterFence, scope: CompletedScope + ) -> ScanState | None: + """Read only under the active writer fence, including after an uncertain commit.""" + await self.records._fenced_sync(db, fence) + row = await self._row(db, fence, scope) + return scan_state(row) if row else None + + @staticmethod + def _expect(row: CaptureScan | None, expected: ScanVersion, cycle_id: UUID) -> CaptureScan: + if ( + row is None + or row.cycle_id != cycle_id + or row.sweep_id != expected.sweep_id + or row.revision != expected.revision + ): + raise ScanConflict("Scan changed; reload its durable state") + return row + + async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: + """Resume unchanged scans; restart and next-cycle transitions require exact CAS.""" + await self.records._fenced_sync(db, request.fence) + row = await self._row(db, request.fence, request.scope) + if row is None: + if request.expected is not None or request.restart: + raise ScanConflict("Cannot restart a missing scan") + row = CaptureScan( + organization_id=request.fence.organization_id, + sync_id=request.fence.sync_id, + scope_key=scope_key(request.scope), + record_type=request.scope.record_type, + container_id=request.scope.container_id, + revision=1, + ) + db.add(row) + else: + if request.expected is not None: + self._expect(row, request.expected, row.cycle_id) + if ( + row.cycle_id == request.cycle_id + and row.fingerprint == request.fingerprint + and not request.restart + ): + return scan_state(row) + if request.expected is None: + raise ScanConflict("Replacing a scan requires its current version") + if row.cycle_id != request.cycle_id and row.phase != "complete": + raise ScanConflict("An unfinished cycle must be completed before advancing") + if row.cycle_id == request.cycle_id and not request.restart: + raise ScanConflict("Changed scope configuration requires an explicit restart") + row.revision += 1 + row.cycle_id = request.cycle_id + row.sweep_id = uuid4() + row.fingerprint = request.fingerprint + row.phase = "collecting" + row.continuation = request.continuation.value + row.started_at = datetime.now(timezone.utc) + row.completed_at = None + await db.flush() + return scan_state(row) + + async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: + """Capture and advance the page as one transaction, never a partial acknowledgement.""" + sync = await self.records._fenced_sync(db, request.fence) + row = self._expect( + await self._row(db, request.fence, request.scope), request.expected, request.cycle_id + ) + if row.phase != "collecting": + raise ScanConflict("Scan is not collecting pages") + if any( + record.identity.record_type != row.record_type + or record.identity.container_id != row.container_id + for record in request.records + ): + raise ScanConflict("Page contains records outside its exact scope") + captured = await self.records._capture_locked( + db, + sync, + CaptureBatch(fence=request.fence, records=request.records), + seen_id=row.sweep_id, + ) + row.continuation = request.continuation.value + row.revision += 1 + if request.final: + row.phase = "reconciling" + await db.flush() + return ScanResult(state=scan_state(row), capture=captured) + + async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: + """A final page is necessary; absence completion is durable and bounded.""" + sync = await self.records._fenced_sync(db, request.fence) + row = self._expect( + await self._row(db, request.fence, request.scope), request.expected, request.cycle_id + ) + if row.phase != "reconciling": + raise ScanConflict("Only fully collected scans can reconcile absence") + result = await self.records._reconcile_scope_locked( + db, + sync, + ReconcileScope( + fence=request.fence, + scope=request.scope, + removal_reason=request.removal_reason, + observed_at=request.observed_at, + limit=request.limit, + ), + seen_id=row.sweep_id, + ) + row.revision += 1 + if not result.has_more: + row.phase = "complete" + row.completed_at = datetime.now(timezone.utc) + await db.flush() + return ScanResult(state=scan_state(row), capture=result.capture) diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index e321131355..31b243e1f0 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -9,11 +9,20 @@ from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult from airweave.domains.entities.canonical.requests import ( CaptureBatch, + CompletedScope, ReconcileScope, RemovedScope, StartedScope, WriterFence, ) +from airweave.domains.entities.canonical.scan_models import ( + BeginScan, + CommitScanPage, + ReconcileScan, + ScanResult, + ScanState, +) +from airweave.domains.entities.canonical.scan_store import CanonicalScanStore from airweave.domains.entities.canonical.store import CanonicalRecordStore @@ -23,6 +32,7 @@ class CanonicalCaptureService: def __init__(self, store: CanonicalRecordStore): """Inject the persistence owner without hidden global sessions.""" self.store = store + self.scans = CanonicalScanStore(store) async def activate_writer( self, @@ -89,3 +99,25 @@ async def save_checkpoint( """Advance only after the caller's full capture/reconciliation completion barrier.""" async with UnitOfWork(db): await self.store.save_checkpoint(db, fence, cursor_data) + + async def read_scan( + self, db: AsyncSession, fence: WriterFence, scope: CompletedScope + ) -> ScanState | None: + """Reload an exact scope after restart or uncertain acknowledgement.""" + async with UnitOfWork(db): + return await self.scans.read(db, fence, scope) + + async def begin_scan(self, db: AsyncSession, request: BeginScan) -> ScanState: + """Begin/resume under the same fence as captured records.""" + async with UnitOfWork(db): + return await self.scans.begin(db, request) + + async def commit_scan_page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: + """Commit original records and their continuation atomically.""" + async with UnitOfWork(db): + return await self.scans.page(db, request) + + async def reconcile_scan(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: + """Commit one bounded whole-scope absence reconciliation step.""" + async with UnitOfWork(db): + return await self.scans.reconcile(db, request) diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 65e8f9a152..1cea61c4bc 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -236,7 +236,7 @@ async def capture(self, db: AsyncSession, batch: CaptureBatch) -> CaptureResult: return await self._capture_locked(db, sync, batch) async def _capture_locked( - self, db: AsyncSession, sync: Sync, batch: CaptureBatch + self, db: AsyncSession, sync: Sync, batch: CaptureBatch, *, seen_id: UUID | None = None ) -> CaptureResult: changes = [] unchanged = 0 @@ -256,7 +256,7 @@ async def _capture_locked( ).scalar_one_or_none() fingerprint = capture_fingerprint(observation) if entity is not None and entity.capture_hash == fingerprint: - entity.last_seen_run_id = batch.fence.attempt_id + entity.last_seen_run_id = seen_id or batch.fence.attempt_id entity.observed_at = observation.observed_at unchanged += 1 continue @@ -293,7 +293,7 @@ async def _capture_locked( entity.removal_reason = observation.removal_reason entity.completeness = observation.completeness entity.blob_references = [blob.model_dump(mode="json") for blob in observation.blobs] - entity.last_seen_run_id = batch.fence.attempt_id + entity.last_seen_run_id = seen_id or batch.fence.attempt_id entity.projection_error = None # Flush before reading default/index state and before another observation of this ID. await db.flush() @@ -322,6 +322,12 @@ async def _capture_locked( async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> ReconcileResult: """Tombstone a bounded batch absent from one successfully enumerated scope.""" sync = await self._fenced_sync(db, request.fence) + return await self._reconcile_scope_locked(db, sync, request) + + async def _reconcile_scope_locked( + self, db: AsyncSession, sync: Sync, request: ReconcileScope, *, seen_id: UUID | None = None + ) -> ReconcileResult: + """Share exact-scope reconciliation with durable whole-scope scans.""" statement = ( select(Entity) .where( @@ -331,7 +337,7 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re Entity.container_id.is_not_distinct_from(request.scope.container_id), Entity.record_revision > 0, Entity.deleted_at.is_(None), - Entity.last_seen_run_id.is_distinct_from(request.fence.attempt_id), + Entity.last_seen_run_id.is_distinct_from(seen_id or request.fence.attempt_id), ) .order_by(Entity.id) .limit(request.limit + 1) @@ -355,7 +361,7 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re for entity in entities[: request.limit] ) result = await self._capture_locked( - db, sync, CaptureBatch(fence=request.fence, records=observations) + db, sync, CaptureBatch(fence=request.fence, records=observations), seen_id=seen_id ) return ReconcileResult(capture=result, has_more=len(entities) > request.limit) diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 267c88d6fa..a777871007 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -107,6 +107,7 @@ async def database(request): await connection.run_sync(migrate, "0002_projection_publication.py") await connection.run_sync(migrate, "0003_mail_thread_index.py") await connection.run_sync(migrate, "0004_projection_generation.py") + await connection.run_sync(migrate, "0005_capture_scan.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() 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..7dff08f693 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_scans.py @@ -0,0 +1,279 @@ +"""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.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): + 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(ScanConflict, match="unfinished"): + await begin(database, service, fence, 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 + next_cycle = await begin(database, service, fence, 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(ScanConflict, match="unfinished"): + 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="")) + async with database() as db: + other = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=CompletedScope(record_type="event", container_id=""), + cycle_id=state.cycle_id, + fingerprint=FINGERPRINT, + ), + ) + assert state.version.sweep_id != other.version.sweep_id + async with database() as db: + assert await db.scalar(select(func.count()).select_from(CaptureScan)) == 2 + 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" + + +async def test_root_refresh_does_not_restart_completed_child_in_same_cycle(database, source): + service, fence = source + child = await begin(database, service, fence) + final = await page(database, service, fence, child, final=True) + complete = (await reconcile(database, service, fence, final.state)).state + root_scope = CompletedScope(record_type="channel") + async with database() as db: + root = await service.begin_scan( + db, + BeginScan( + fence=fence, scope=root_scope, cycle_id=child.cycle_id, fingerprint=FINGERPRINT + ), + ) + async with database() as db: + refreshed = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=root_scope, + cycle_id=child.cycle_id, + fingerprint=FINGERPRINT, + expected=root.version, + restart=True, + ), + ) + assert refreshed.version.sweep_id != root.version.sweep_id + assert await begin(database, service, fence, cycle=child.cycle_id) == complete diff --git a/backend/airweave/models/__init__.py b/backend/airweave/models/__init__.py index 43f5ccae23..fb137f2072 100644 --- a/backend/airweave/models/__init__.py +++ b/backend/airweave/models/__init__.py @@ -4,6 +4,7 @@ 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 @@ -39,6 +40,7 @@ "Collection", "Entity", "EntityChange", + "CaptureScan", "ProjectionGeneration", "EntityCount", "Connection", diff --git a/backend/airweave/models/capture_scan.py b/backend/airweave/models/capture_scan.py new file mode 100644 index 0000000000..630c7308a1 --- /dev/null +++ b/backend/airweave/models/capture_scan.py @@ -0,0 +1,36 @@ +"""One recoverable whole-scope scan, independent of worker attempts.""" + +from datetime import datetime +from uuid import UUID + +from sqlalchemy import BigInteger, CheckConstraint, DateTime, ForeignKey, String, UniqueConstraint +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) + 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) + started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + __table_args__ = ( + 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/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") From c4422e44ed1fb00df528e683efcd9f6ad8e66b14 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:48:50 -0700 Subject: [PATCH 034/266] test: retain Slack backoff deadline evidence and fix cancellation fixture --- backend/tests/live/README.md | 10 +- ...lifecycle-backoff-incomplete-20260930.json | 178 ++++++++++++++++++ backend/tests/live/provider_lifecycle.py | 1 + 3 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 backend/tests/live/evidence/slack-lifecycle-backoff-incomplete-20260930.json diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 0d802a7e8b..2928bcd78d 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -230,7 +230,15 @@ The subsequent approved budget is 600 requests, 10,000 observations and 1,800 seconds per process. Slack uses its five-attempt retry policy and honors the full Retry-After wait; the test does not shorten waits to fit its total budget. Only status, numeric retry delay and request count are retained as diagnostics. -The longer attempt's result is pending; do not infer full capture from the budget. +The longer attempt reached its 1,800-second deadline after 74 requests and 328 +observations. Thirty rate limits each supplied a 60-second Retry-After. The job +became cancelled; two scopes started, none completed, and the checkpoint stayed +unchanged. Cancellation telemetry then accessed missing fixture connection +metadata, masking the timeout with AttributeError. The fixture is corrected; +this was not rerun. No second process started, and private schema/files were +removed. Full Slack coverage remains unverified. A durable per-scope sweep and +continuation design is needed before repeatedly rescanning this history. +See `evidence/slack-lifecycle-backoff-incomplete-20260930.json`. Wispr enumeration and one diagnostic retry each failed after 18 requests and nine partial observations. `GET_MEETING` returned an opaque tool error string; the diff --git a/backend/tests/live/evidence/slack-lifecycle-backoff-incomplete-20260930.json b/backend/tests/live/evidence/slack-lifecycle-backoff-incomplete-20260930.json new file mode 100644 index 0000000000..438eee1329 --- /dev/null +++ b/backend/tests/live/evidence/slack-lifecycle-backoff-incomplete-20260930.json @@ -0,0 +1,178 @@ +{ + "provider": "slack", + "date": "2026-09-30", + "scope": "full accessible history without hidden filter", + "limits": { + "requests": 600, + "observations": 10000, + "seconds": 1800 + }, + "attempt": { + "provider_requests": 74, + "observations": 328, + "started_scopes": 2, + "completed_scopes": 0, + "job_status": "cancelled", + "reported_error_type": "AttributeError", + "checkpoint_unchanged": true, + "rate_limits": [ + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 12 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 14 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 17 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 19 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 21 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 23 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 25 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 27 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 30 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 33 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 35 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 37 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 39 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 41 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 43 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 45 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 47 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 49 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 51 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 54 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 56 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 58 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 60 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 62 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 64 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 66 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 68 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 70 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 72 + }, + { + "status": 429, + "retry_after_seconds": 60.0, + "request_number": 74 + } + ] + }, + "diagnosis": "Deadline cancellation transitioned job to cancelled; fixture lacked connection metadata accessed by cancellation telemetry, masking timeout with AttributeError. Fixture corrected after run; no rerun.", + "free_disk_bytes_before": 9485815808, + "second_process_started": false, + "schema_removed": true, + "blob_directory_removed": true, + "full_scope_verified": false, + "provider_writes": false +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index a5cb3f6cfc..21f8281d54 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -246,6 +246,7 @@ async def db_context(): collection_id=uuid4(), source_connection_id=uuid4(), source_short_name=name, + connection=SimpleNamespace(id=uuid4(), short_name=name), execution_config=config, batch_size=10, max_batch_latency_ms=20, From 2ce5196e4db8f2f18d77b3880c8ac1ec4a10d5f8 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:51:31 -0700 Subject: [PATCH 035/266] fix: require positive critical dependency health for staging readiness --- .github/workflows/owned-source-store.yml | 2 ++ Makefile | 10 ++++--- backend/airweave/adapters/health/temporal.py | 10 +++++-- .../health/tests/test_health_probes.py | 13 +++++++-- backend/airweave/api/v1/endpoints/health.py | 5 ++-- backend/airweave/core/health/service.py | 9 ++++--- .../airweave/core/health/tests/test_health.py | 27 ++++++++++++++++++- deploy/README.md | 6 +++-- porter.yaml | 1 + 9 files changed, 66 insertions(+), 17 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index d5c2ae102a..110eeff482 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -71,6 +71,8 @@ jobs: run: make test-search - name: Verify service authentication run: make test-auth + - name: Verify dependency readiness + run: make test-health - name: Record tested Vespa image run: docker image inspect vespaengine/vespa@sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd --format '{{json .RepoDigests}}' - name: Verify real Vespa schema and retrieval diff --git a/Makefile b/Makefile index 94a8b7a525..722a384bbc 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest IMAGE ?= almanac-source-store:local -.PHONY: help setup check test-store test-capture test-search test-auth test-index validate-deploy build +.PHONY: help setup check test-store test-capture test-search test-auth test-health test-index validate-deploy build help: @echo 'setup Install the locked backend and development dependencies' @@ -13,8 +13,9 @@ help: @echo 'test-capture Verify source capture, transport and blob handling' @echo 'test-search Verify search services and visibility behavior' @echo 'test-auth Verify service-key and source authorization boundaries' + @echo 'test-health Verify dependency probes and readiness behavior' @echo 'test-index Deploy/test schemas on explicitly disposable local Vespa' - @echo 'check Run store, capture, search and authentication checks' + @echo 'check Run store, capture, search, authentication and health checks' @echo 'validate-deploy Validate the Porter application manifest locally (no deployment)' @echo '' @echo 'test-store requires CANONICAL_TEST_DATABASE_URL; it creates and removes test schemas.' @@ -42,8 +43,11 @@ test-search: test-auth: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/api/tests/test_service_auth.py airweave/api/tests/test_sync_authorization.py airweave/api/tests/test_context_resolver_auth.py +test-health: + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/core/health/tests airweave/adapters/health/tests airweave/core/container/tests/test_health_wiring.py + test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py -check: test-store test-capture test-search test-auth +check: test-store test-capture test-search test-auth test-health diff --git a/backend/airweave/adapters/health/temporal.py b/backend/airweave/adapters/health/temporal.py index 965f6a299d..62e0b01307 100644 --- a/backend/airweave/adapters/health/temporal.py +++ b/backend/airweave/adapters/health/temporal.py @@ -15,17 +15,23 @@ class TemporalHealthProbe(HealthProbe): """Probes Temporal via the gRPC health check on its service client.""" def __init__(self, get_client: Callable[[], TemporalClientType | None]) -> None: + """Use the application's current connection without creating a probe-owned client.""" self._get_client = get_client @property def name(self) -> str: + """Identify this dependency in readiness configuration and responses.""" return "temporal" async def check(self) -> DependencyCheck: + """Report the actual service health, or skipped until a client is available.""" client = self._get_client() if client is None: return DependencyCheck(status=CheckStatus.skipped) start = time.perf_counter() - await client.service_client.check_health() + serving = await client.service_client.check_health() latency = (time.perf_counter() - start) * 1000 - return DependencyCheck(status=CheckStatus.up, latency_ms=round(latency, 2)) + return DependencyCheck( + status=CheckStatus.up if serving else CheckStatus.down, + latency_ms=round(latency, 2), + ) diff --git a/backend/airweave/adapters/health/tests/test_health_probes.py b/backend/airweave/adapters/health/tests/test_health_probes.py index e22550b508..9cacf0f31c 100644 --- a/backend/airweave/adapters/health/tests/test_health_probes.py +++ b/backend/airweave/adapters/health/tests/test_health_probes.py @@ -6,7 +6,6 @@ from airweave.schemas.health import CheckStatus - # --------------------------------------------------------------------------- # PostgresHealthProbe # --------------------------------------------------------------------------- @@ -104,7 +103,7 @@ async def test_client_available_returns_up(self): from airweave.adapters.health.temporal import TemporalHealthProbe client = MagicMock() - client.service_client.check_health = AsyncMock() + client.service_client.check_health = AsyncMock(return_value=True) probe = TemporalHealthProbe(lambda: client) result = await probe.check() @@ -114,6 +113,16 @@ async def test_client_available_returns_up(self): assert result.latency_ms >= 0 client.service_client.check_health.assert_awaited_once() + @pytest.mark.asyncio + async def test_server_not_serving_returns_down(self): + from airweave.adapters.health.temporal import TemporalHealthProbe + + client = MagicMock() + client.service_client.check_health = AsyncMock(return_value=False) + result = await TemporalHealthProbe(lambda: client).check() + + assert result.status == CheckStatus.down + @pytest.mark.asyncio async def test_client_available_but_unhealthy(self): from airweave.adapters.health.temporal import TemporalHealthProbe diff --git a/backend/airweave/api/v1/endpoints/health.py b/backend/airweave/api/v1/endpoints/health.py index d65953d46a..261f01d193 100644 --- a/backend/airweave/api/v1/endpoints/health.py +++ b/backend/airweave/api/v1/endpoints/health.py @@ -39,9 +39,8 @@ async def readiness( ) -> ReadinessResponse: """Readiness probe — checks critical dependencies. - Only critical probes (Postgres) gate the HTTP status code. Informational - probes (Redis, Temporal) are reported for observability but do not cause - a 503. + Configured critical probes must report up. Informational probes are reported + for observability but do not cause a 503. """ result = await health.check_readiness(debug=settings.DEBUG) diff --git a/backend/airweave/core/health/service.py b/backend/airweave/core/health/service.py index f75610baf6..52864a384d 100644 --- a/backend/airweave/core/health/service.py +++ b/backend/airweave/core/health/service.py @@ -70,8 +70,6 @@ async def check_readiness(self, *, debug: bool) -> ReadinessResponse: ) checks: dict[str, DependencyCheck] = {} - ready = True - for entry in results: if isinstance(entry, BaseException): continue @@ -83,11 +81,14 @@ async def check_readiness(self, *, debug: bool) -> ReadinessResponse: status=CheckStatus.down, error=self._sanitize_error(outcome, debug=debug), ) - if name in critical_names: - ready = False else: checks[name] = outcome + # Critical means positively verified up, including when a probe returns + # a status rather than raising or is cancelled before reporting one. + ready = all( + name in checks and checks[name].status == CheckStatus.up for name in critical_names + ) return ReadinessResponse( status="ready" if ready else "not_ready", checks=checks, diff --git a/backend/airweave/core/health/tests/test_health.py b/backend/airweave/core/health/tests/test_health.py index b3a49feddc..6a83158b9c 100644 --- a/backend/airweave/core/health/tests/test_health.py +++ b/backend/airweave/core/health/tests/test_health.py @@ -2,6 +2,7 @@ import asyncio import errno +from unittest.mock import AsyncMock import pytest @@ -12,7 +13,7 @@ FakeSlowProbe, ) from airweave.core.health.service import HealthService -from airweave.schemas.health import CheckStatus +from airweave.schemas.health import CheckStatus, DependencyCheck # --------------------------------------------------------------------------- # check_readiness @@ -22,6 +23,30 @@ class TestCheckReadiness: """Tests for the readiness orchestrator.""" + @pytest.mark.asyncio + @pytest.mark.parametrize("status", [CheckStatus.down, CheckStatus.skipped]) + @pytest.mark.parametrize("critical", [True, False]) + async def test_returned_non_up_status_gates_only_critical_dependencies(self, status, critical): + probe = FakeProbe("temporal") + probe.check = AsyncMock(return_value=DependencyCheck(status=status)) + svc = HealthService( + critical=[probe] if critical else [], + informational=[] if critical else [probe], + ) + + result = await svc.check_readiness(debug=False) + + assert result.status == ("not_ready" if critical else "ready") + assert result.checks["temporal"].status == status + + @pytest.mark.asyncio + async def test_cancelled_critical_probe_cannot_report_ready(self): + probe = FakeProbe("temporal") + probe.check = AsyncMock(side_effect=asyncio.CancelledError()) + svc = HealthService(critical=[probe], informational=[]) + + assert (await svc.check_readiness(debug=False)).status == "not_ready" + @pytest.mark.asyncio async def test_all_up(self): svc = HealthService( diff --git a/deploy/README.md b/deploy/README.md index ebf57e383d..f17ac5a848 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -33,8 +33,10 @@ package; local MiniLM inference; and a private S3 bucket for canonical bytes. Use the repository's existing dependency configurations as inputs, not its local Compose file as a production deployment: that file enables local auth and exposes ports. Dependency image pins, persistent volumes, backups and a restore exercise -remain deployment work. The stock readiness route gates PostgreSQL only; verify -index, worker and capture readiness separately. +remain deployment work. This manifest gates API readiness on PostgreSQL, Redis +and Temporal reporting up. Missing/skipped critical probes do not pass. Verify +index, worker polling and capture readiness separately; API dependency health +does not prove those workflows work. Create the Porter environment group `almanac-source-store-staging` with only the owned service's values. Supply database connection settings, Redis endpoint, diff --git a/porter.yaml b/porter.yaml index dbfd4eeb43..7c2f52627d 100644 --- a/porter.yaml +++ b/porter.yaml @@ -13,6 +13,7 @@ env: ENVIRONMENT: dev LOCAL_DEVELOPMENT: "false" AUTH_MODE: api_key + HEALTH_CRITICAL_PROBES: postgres,redis,temporal ANALYTICS_ENABLED: "false" STRIPE_ENABLED: "false" STORAGE_BACKEND: aws From 272980cdd7b5651386b9dc2e47a28a3ed6529d3a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:53:35 -0700 Subject: [PATCH 036/266] fix: retry critical Temporal initialization through process startup --- backend/airweave/core/health/protocols.py | 8 +-- .../core/health/tests/test_startup.py | 50 +++++++++++++++++++ backend/airweave/main.py | 4 ++ deploy/README.md | 5 ++ 4 files changed, 63 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/core/health/tests/test_startup.py diff --git a/backend/airweave/core/health/protocols.py b/backend/airweave/core/health/protocols.py index 1e5fdf6f47..06c0ff377d 100644 --- a/backend/airweave/core/health/protocols.py +++ b/backend/airweave/core/health/protocols.py @@ -13,9 +13,9 @@ class HealthProbe(Protocol): """Protocol for a single infrastructure health check. - Implementations return a ``DependencyCheck`` on success and **raise** - on failure. The orchestrator handles timeouts and error sanitization, - so probes can stay simple. + Implementations return a ``DependencyCheck`` with the observed status or + raise when a check fails. The orchestrator handles timeouts and error + sanitization, so probes can stay simple. The critical-vs-informational distinction is a deployment concern and lives in the wiring layer, not in the probe itself. @@ -30,7 +30,7 @@ async def check(self) -> DependencyCheck: """Probe the dependency and return its status. Returns: - A ``DependencyCheck`` with ``status=up`` and measured latency. + A ``DependencyCheck`` with the observed status and measured latency. Raises: Any exception on failure — the orchestrator will catch it. diff --git a/backend/airweave/core/health/tests/test_startup.py b/backend/airweave/core/health/tests/test_startup.py new file mode 100644 index 0000000000..cdada5e756 --- /dev/null +++ b/backend/airweave/core/health/tests/test_startup.py @@ -0,0 +1,50 @@ +"""Critical Temporal startup failures must be recoverable through process restart.""" + +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI + +from airweave import main +from airweave.core import container, metrics_service +from airweave.db import session + + +@pytest.mark.asyncio +@pytest.mark.parametrize("critical", [True, False]) +async def test_schedule_startup_failure_respects_critical_dependency(monkeypatch, critical): + """A required dependency cannot leave an API alive but permanently skipped.""" + services = MagicMock() + services.temporal_schedule_service.ensure_system_schedules = AsyncMock( + side_effect=[ConnectionError("Temporal unavailable"), None] + ) + monkeypatch.setattr(container, "container", services) + monkeypatch.setattr(container, "initialize_container", MagicMock()) + monkeypatch.setattr(main, "validate_embedding_config", AsyncMock()) + monkeypatch.setattr(main, "AsyncSessionLocal", MagicMock(return_value=AsyncMock())) + monkeypatch.setattr( + main.settings, + "HEALTH_CRITICAL_PROBES", + "postgres,temporal" if critical else "postgres", + ) + monkeypatch.setattr(session, "health_check_engine", AsyncMock()) + + @asynccontextmanager + async def metrics(*args): + yield + + monkeypatch.setattr(metrics_service, "metrics_lifespan", metrics) + + if critical: + with pytest.raises(ConnectionError, match="Temporal unavailable"): + async with main.lifespan(FastAPI()): + pytest.fail("Critical startup failure must not enter serving lifespan") + else: + async with main.lifespan(FastAPI()): + pass + + # A subsequent startup can succeed; no sticky health flag hides recovery. + async with main.lifespan(FastAPI()): + pass + assert services.temporal_schedule_service.ensure_system_schedules.await_count == 2 diff --git a/backend/airweave/main.py b/backend/airweave/main.py index 271adebfa3..2064278cb6 100644 --- a/backend/airweave/main.py +++ b/backend/airweave/main.py @@ -71,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}" ) diff --git a/deploy/README.md b/deploy/README.md index f17ac5a848..84c782e6f1 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -38,6 +38,11 @@ and Temporal reporting up. Missing/skipped critical probes do not pass. Verify index, worker polling and capture readiness separately; API dependency health does not prove those workflows work. +When Temporal is critical, failure to initialize its system schedules aborts API +startup so the existing process supervisor can retry. Optional Temporal setups +retain their permissive startup behavior. This avoids an API stuck unready with +no initialized Temporal client after a transient startup outage. + Create the Porter environment group `almanac-source-store-staging` with only the owned service's values. Supply database connection settings, Redis endpoint, Temporal endpoint/namespace, Vespa endpoint, TEXT2VEC_INFERENCE_URL, From 13e5b46088f8fdad762de36e0b174900ec99a488 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:56:47 -0700 Subject: [PATCH 037/266] docs: distinguish actual staging pool limits from cost estimates --- deploy/DEPENDENCIES.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/deploy/DEPENDENCIES.md b/deploy/DEPENDENCIES.md index f744c8b76e..b0d1881806 100644 --- a/deploy/DEPENDENCIES.md +++ b/deploy/DEPENDENCIES.md @@ -133,6 +133,23 @@ platform. It is a fallback only if the cluster's fixed charges make the reviewed budget impossible. Single-replica stateful services accept downtime; neither option is production HA. +### Existing Porter configuration readback + +On 2026-09-30, the existing cluster6108 form showed Karpenter cost optimization +enabled, max8 application vCPU, on-demand capacity,50GiB node disks, no instance +family restriction and no per-instance size restriction. No settings were saved. +This does **not** establish a dollar ceiling or the proposed t3a.2xlarge price: +the configured pool can select other instance families and sizes. The Porter +[node-group documentation](https://docs.porter.run/cloud-accounts/node-groups) +also describes fixed system and monitoring groups alongside the autoscaled +application group, and private-node egress through NAT by default. Their actual +sizes, node count and NAT charges have not been verified for this cluster. + +Before retrying provisioning, reconcile the original$201.22 estimate against all +three groups, control plane, disks and networking. Select an explicitly priced +bounded instance set or fixed node shape that fits the full$500 cap. Max CPU8 +alone is insufficient evidence. The quota preflight remains a separate gate. + ## Evidence Local implementation: `porter.yaml`, `docker/docker-compose.yml`, From 4856e25561ff4ede69ee73a4012eb9b8b9648ac5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 04:58:34 -0700 Subject: [PATCH 038/266] Fence scoped capture with durable cycle ownership and finalization --- backend/airweave/crud/crud_sync_cursor.py | 31 +- .../entities/canonical/cycle_models.py | 64 ++++ .../domains/entities/canonical/cycle_store.py | 201 ++++++++++ .../domains/entities/canonical/scan_store.py | 62 ++- .../domains/entities/canonical/service.py | 25 ++ .../domains/entities/canonical/store.py | 3 + .../entities/canonical/tests/test_cycles.py | 358 ++++++++++++++++++ .../entities/canonical/tests/test_scans.py | 86 ++--- 8 files changed, 776 insertions(+), 54 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/cycle_models.py create mode 100644 backend/airweave/domains/entities/canonical/cycle_store.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_cycles.py diff --git a/backend/airweave/crud/crud_sync_cursor.py b/backend/airweave/crud/crud_sync_cursor.py index fd3f0fbba5..bdd5568eb8 100644 --- a/backend/airweave/crud/crud_sync_cursor.py +++ b/backend/airweave/crud/crud_sync_cursor.py @@ -9,6 +9,7 @@ from airweave import models, schemas from airweave.core.context import BaseContext from airweave.crud._base_organization import CRUDBaseOrganization +from airweave.domains.entities.canonical.cycle_models import CYCLE_KEY class CRUDSyncCursor( @@ -16,6 +17,21 @@ class CRUDSyncCursor( ): """CRUD operations for sync cursor.""" + async def _require_legacy_cursor( + self, db: AsyncSession, sync_id: UUID, ctx: BaseContext + ) -> None: + """Serialize with canonical writers before testing ownership, never erase active cycles.""" + sync = await db.scalar( + select(models.Sync) + .where(models.Sync.id == sync_id, models.Sync.organization_id == ctx.organization.id) + .with_for_update() + ) + if sync is None: + raise ValueError("Source does not exist in this organization") + cursor = await self.get_by_sync_id(db, sync_id=sync_id, ctx=ctx) + if cursor is not None and CYCLE_KEY in cursor.cursor_data: + raise ValueError("Cycle-owned progress cannot be changed by legacy cursor writes") + async def get_by_sync_id( self, db: AsyncSession, *, sync_id: UUID, ctx: BaseContext ) -> Optional[models.SyncCursor]: @@ -33,7 +49,7 @@ async def get_by_sync_id( models.SyncCursor.sync_id == sync_id, models.SyncCursor.organization_id == ctx.organization.id, ) - result = await db.execute(stmt) + result = await db.execute(stmt.execution_options(populate_existing=True)) return result.scalar_one_or_none() async def create_or_update( @@ -55,6 +71,9 @@ async def create_or_update( Returns: Created or updated sync cursor """ + await self._require_legacy_cursor(db, sync_id, ctx) + if CYCLE_KEY in obj_in.cursor_data: + raise ValueError("Reserved cycle state cannot be supplied by a source") # Check if cursor already exists for this sync existing_cursor = await self.get_by_sync_id(db, sync_id=sync_id, ctx=ctx) @@ -63,11 +82,7 @@ async def create_or_update( return await self.update(db, db_obj=existing_cursor, obj_in=obj_in, ctx=ctx) else: # Create new cursor - # Handle both dict and Pydantic model inputs - if isinstance(obj_in, dict): - obj_in["sync_id"] = sync_id - else: - obj_in.sync_id = sync_id + obj_in.sync_id = sync_id return await self.create(db, obj_in=obj_in, ctx=ctx) async def update_cursor_data( @@ -89,6 +104,9 @@ async def update_cursor_data( Returns: Updated sync cursor if found, None otherwise """ + await self._require_legacy_cursor(db, sync_id, ctx) + if CYCLE_KEY in cursor_data: + raise ValueError("Reserved cycle state cannot be supplied by a source") cursor = await self.get_by_sync_id(db, sync_id=sync_id, ctx=ctx) if cursor: @@ -108,6 +126,7 @@ async def delete_by_sync_id(self, db: AsyncSession, *, sync_id: UUID, ctx: BaseC Returns: True if deleted, False if not found """ + await self._require_legacy_cursor(db, sync_id, ctx) cursor = await self.get_by_sync_id(db, sync_id=sync_id, ctx=ctx) if cursor: diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py new file mode 100644 index 0000000000..12920d42ed --- /dev/null +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -0,0 +1,64 @@ +"""Typed engine-owned cycle state inside the existing SyncCursor.""" + +from typing import Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from airweave.domains.entities.canonical.requests import WriterFence + +CYCLE_KEY = "canonical_cycle" + + +class CycleConfiguration(BaseModel): + """Whole-scope topology; child containers are visible root native IDs.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") + root_record_type: str = Field(min_length=1, max_length=200) + child_record_types: tuple[str, ...] = Field(default=(), max_length=20) + + @model_validator(mode="after") + def distinct_types(self) -> "CycleConfiguration": + """Keep scope discovery bounded and unambiguous.""" + if len(set(self.child_record_types)) != len(self.child_record_types): + raise ValueError("Child record types must be unique") + if any(not item or len(item) > 200 for item in self.child_record_types): + raise ValueError("Invalid child record type") + return self + + +class CycleVersion(BaseModel): + """Exact checkpoint version observed by the caller.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + revision: int = Field(ge=1) + + +class CaptureCycle(BaseModel): + """Completion stays durable until an explicit next-cycle CAS.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + schema_version: Literal[1] = 1 + version: CycleVersion + configuration: CycleConfiguration + phase: Literal["active", "complete"] = "active" + root_writer_attempt_id: UUID | None = None + + +class BeginCycle(BaseModel): + """Resume an active cycle or CAS-advance its completed predecessor.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + configuration: CycleConfiguration + expected: CycleVersion | None = None + + +class CompleteCycle(BaseModel): + """Publish only an unchanged active cycle with all required scans complete.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + expected: CycleVersion diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py new file mode 100644 index 0000000000..71c2a18a64 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -0,0 +1,201 @@ +"""Cycle ownership on the existing sync cursor; callers already hold the Sync fence.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from sqlalchemy import exists, select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint +from airweave.domains.entities.canonical.cycle_models import ( + CYCLE_KEY, + BeginCycle, + CaptureCycle, + CompleteCycle, + CycleVersion, +) +from airweave.domains.entities.canonical.requests import CompletedScope, WriterFence +from airweave.domains.entities.canonical.store import CanonicalStoreError, content_is_available +from airweave.models.capture_scan import CaptureScan +from airweave.models.entity import Entity +from airweave.models.sync import Sync +from airweave.models.sync_cursor import SyncCursor + + +class CycleConflict(CanonicalStoreError): + """Reload the active cycle; do not guess recovery or erase engine state.""" + + code = "cycle_conflict" + + +async def cursor_row(db: AsyncSession, fence: WriterFence) -> SyncCursor | None: + """The caller holds the same sync lock as all capture mutations.""" + return await db.scalar( + select(SyncCursor) + .where( + SyncCursor.sync_id == fence.sync_id, + SyncCursor.organization_id == fence.organization_id, + ) + .execution_options(populate_existing=True) + ) + + +def cycle_state(cursor: SyncCursor | None) -> CaptureCycle | None: + """Reject malformed state rather than falling back to a new cycle.""" + if cursor is None or CYCLE_KEY not in cursor.cursor_data: + return None + return CaptureCycle.model_validate(cursor.cursor_data[CYCLE_KEY]) + + +def persist(cursor: SyncCursor, state: CaptureCycle) -> None: + """Preserve provider checkpoint fields until finalization.""" + cursor.cursor_data = {**cursor.cursor_data, CYCLE_KEY: state.model_dump(mode="json")} + cursor.last_updated = datetime.now(timezone.utc) + + +async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: + """An explicit completed-version CAS starts the next cycle.""" + cursor = await cursor_row(db, request.fence) + state = cycle_state(cursor) + if state: + if request.expected is not None and request.expected != state.version: + raise CycleConflict("Cycle changed; reload durable progress") + if state.phase == "active": + if state.configuration != request.configuration: + raise CycleConflict( + "Active cycle configuration changed; explicit abandonment required" + ) + return state + if request.expected is None: + return state + elif request.expected is not None: + raise CycleConflict("Expected cycle does not exist") + if cursor is None: + cursor = SyncCursor( + organization_id=request.fence.organization_id, + sync_id=request.fence.sync_id, + cursor_data={}, + ) + db.add(cursor) + state = CaptureCycle( + version=CycleVersion( + cycle_id=uuid4(), revision=1 if state is None else state.version.revision + 1 + ), + configuration=request.configuration, + ) + persist(cursor, state) + await db.flush() + return state + + +async def attest_cycle( + db: AsyncSession, fence: WriterFence, cycle_id: UUID +) -> tuple[SyncCursor, CaptureCycle]: + """Reject obsolete or foreign cycle identities before any scope mutation.""" + cursor = await cursor_row(db, fence) + state = cycle_state(cursor) + if state is None or state.phase != "active" or state.version.cycle_id != cycle_id: + raise CycleConflict("Scan does not belong to the active cycle") + return cursor, state + + +async def root_ready(db: AsyncSession, fence: WriterFence, state: CaptureCycle) -> bool: + """A retry must refresh membership before touching or finishing child scopes.""" + root = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.record_type == state.configuration.root_record_type, + CaptureScan.container_id.is_(None), + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.phase == "complete", + ) + ) + return root is not None and ( + not state.configuration.child_record_types + or state.root_writer_attempt_id == fence.attempt_id + ) + + +async def attest_scope( + db: AsyncSession, fence: WriterFence, state: CaptureCycle, scope: CompletedScope +) -> bool: + """Only whole root or visible declared child scopes belong to this cycle.""" + if scope.record_type == state.configuration.root_record_type and scope.container_id is None: + return True + if ( + scope.record_type not in state.configuration.child_record_types + or scope.container_id is None + ): + raise CycleConflict("Scope is not declared by the active cycle") + if not await root_ready(db, fence, state): + raise CycleConflict("Refresh and complete root membership before child work") + parent = await db.scalar( + select(Entity).where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == state.configuration.root_record_type, + Entity.container_id.is_(None), + Entity.native_id == scope.container_id, + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + ) + ) + if parent is None: + raise CycleConflict("Child scope no longer has a visible parent") + return False + + +async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) -> CaptureCycle: + """Publish the checkpoint only when SQL proves all currently visible scopes complete.""" + cursor, state = await attest_cycle(db, request.fence, request.expected.cycle_id) + if state.version != request.expected: + raise CycleConflict("Cycle changed before completion") + if not await root_ready(db, request.fence, state): + raise CycleConflict("Root enumeration is incomplete or stale") + for child_type in state.configuration.child_record_types: + complete_child = exists( + select(CaptureScan.id).where( + CaptureScan.organization_id == request.fence.organization_id, + CaptureScan.sync_id == request.fence.sync_id, + CaptureScan.record_type == child_type, + CaptureScan.container_id == Entity.native_id, + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.phase == "complete", + ) + ) + missing = await db.scalar( + select(Entity.id) + .where( + Entity.organization_id == request.fence.organization_id, + Entity.sync_id == request.fence.sync_id, + Entity.entity_definition_short_name == state.configuration.root_record_type, + Entity.container_id.is_(None), + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + ~complete_child, + ) + .limit(1) + ) + if missing is not None: + raise CycleConflict("A visible parent still has an incomplete child scope") + state = state.model_copy( + update={ + "phase": "complete", + "version": CycleVersion( + cycle_id=state.version.cycle_id, revision=state.version.revision + 1 + ), + } + ) + persist(cursor, state) + cursor.cursor_data = { + **cursor.cursor_data, + "canonical_checkpoint": CanonicalCheckpoint( + writer_attempt_id=request.fence.attempt_id, + observed_change_sequence=sync.observed_change_sequence, + ).model_dump(mode="json"), + } + await db.flush() + return state diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 8de432c4ee..7d2e811d29 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -7,10 +7,18 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from airweave.domains.entities.canonical.cycle_models import CycleVersion +from airweave.domains.entities.canonical.cycle_store import ( + CycleConflict, + attest_cycle, + attest_scope, + persist, +) from airweave.domains.entities.canonical.requests import ( CaptureBatch, CompletedScope, ReconcileScope, + RecordIdentity, WriterFence, ) from airweave.domains.entities.canonical.scan_models import ( @@ -96,7 +104,20 @@ def _expect(row: CaptureScan | None, expected: ScanVersion, cycle_id: UUID) -> C async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: """Resume unchanged scans; restart and next-cycle transitions require exact CAS.""" await self.records._fenced_sync(db, request.fence) + cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) + is_root = await attest_scope(db, request.fence, cycle, request.scope) + if request.fingerprint != cycle.configuration.fingerprint: + raise CycleConflict("Scan configuration differs from the active cycle") row = await self._row(db, request.fence, request.scope) + if ( + row is not None + and row.cycle_id == request.cycle_id + and is_root + and cycle.configuration.child_record_types + and cycle.root_writer_attempt_id != request.fence.attempt_id + and not request.restart + ): + raise CycleConflict("Restart root membership for this writer attempt") if row is None: if request.expected is not None or request.restart: raise ScanConflict("Cannot restart a missing scan") @@ -120,8 +141,6 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: return scan_state(row) if request.expected is None: raise ScanConflict("Replacing a scan requires its current version") - if row.cycle_id != request.cycle_id and row.phase != "complete": - raise ScanConflict("An unfinished cycle must be completed before advancing") if row.cycle_id == request.cycle_id and not request.restart: raise ScanConflict("Changed scope configuration requires an explicit restart") row.revision += 1 @@ -132,12 +151,30 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: row.continuation = request.continuation.value row.started_at = datetime.now(timezone.utc) row.completed_at = None + if is_root: + cycle = cycle.model_copy( + update={ + "root_writer_attempt_id": request.fence.attempt_id, + "version": CycleVersion( + cycle_id=cycle.version.cycle_id, revision=cycle.version.revision + 1 + ), + } + ) + persist(cursor, cycle) await db.flush() return scan_state(row) async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: """Capture and advance the page as one transaction, never a partial acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) + _, cycle = await attest_cycle(db, request.fence, request.cycle_id) + is_root = await attest_scope(db, request.fence, cycle, request.scope) + if ( + is_root + and cycle.configuration.child_record_types + and cycle.root_writer_attempt_id != request.fence.attempt_id + ): + raise CycleConflict("Root membership belongs to an earlier writer attempt") row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) @@ -149,6 +186,13 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: for record in request.records ): raise ScanConflict("Page contains records outside its exact scope") + if not is_root: + parent = RecordIdentity( + record_type=cycle.configuration.root_record_type, + native_id=request.scope.container_id, + ) + if any(record.parent != parent for record in request.records): + raise ScanConflict("Child records must retain their declared parent identity") captured = await self.records._capture_locked( db, sync, @@ -165,6 +209,14 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: """A final page is necessary; absence completion is durable and bounded.""" sync = await self.records._fenced_sync(db, request.fence) + _, cycle = await attest_cycle(db, request.fence, request.cycle_id) + is_root = await attest_scope(db, request.fence, cycle, request.scope) + if ( + is_root + and cycle.configuration.child_record_types + and cycle.root_writer_attempt_id != request.fence.attempt_id + ): + raise CycleConflict("Root membership belongs to an earlier writer attempt") row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) @@ -176,7 +228,11 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul ReconcileScope( fence=request.fence, scope=request.scope, - removal_reason=request.removal_reason, + removal_reason=( + "scope_removed" + if is_root and cycle.configuration.child_record_types + else request.removal_reason + ), observed_at=request.observed_at, limit=request.limit, ), diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index 31b243e1f0..87e0ede261 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -6,6 +6,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork +from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle, CompleteCycle +from airweave.domains.entities.canonical.cycle_store import ( + begin_cycle, + complete_cycle, + cursor_row, + cycle_state, +) from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult from airweave.domains.entities.canonical.requests import ( CaptureBatch, @@ -121,3 +128,21 @@ async def reconcile_scan(self, db: AsyncSession, request: ReconcileScan) -> Scan """Commit one bounded whole-scope absence reconciliation step.""" async with UnitOfWork(db): return await self.scans.reconcile(db, request) + + async def read_cycle(self, db: AsyncSession, fence: WriterFence) -> CaptureCycle | None: + """Read the current durable boundary without creating a new cycle.""" + async with UnitOfWork(db): + await self.store._fenced_sync(db, fence) + return cycle_state(await cursor_row(db, fence)) + + async def begin_cycle(self, db: AsyncSession, request: BeginCycle) -> CaptureCycle: + """Resume or explicitly advance the existing checkpoint's cycle.""" + async with UnitOfWork(db): + await self.store._fenced_sync(db, request.fence) + return await begin_cycle(db, request) + + async def complete_cycle(self, db: AsyncSession, request: CompleteCycle) -> CaptureCycle: + """Atomically verify scope completion and publish the source checkpoint.""" + async with UnitOfWork(db): + sync = await self.store._fenced_sync(db, request.fence) + return await complete_cycle(db, sync, request) diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 1cea61c4bc..ca41549852 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -10,6 +10,7 @@ from sqlalchemy.orm import aliased from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint +from airweave.domains.entities.canonical.cycle_models import CYCLE_KEY from airweave.domains.entities.canonical.models import ( CaptureResult, ChangePage, @@ -502,6 +503,8 @@ async def save_checkpoint( """Call only after capture barrier and successful exact-scope reconciliation.""" sync = await self._fenced_sync(db, fence) cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + if CYCLE_KEY in cursor_data or (cursor is not None and CYCLE_KEY in cursor.cursor_data): + raise CanonicalStoreError("Cycle-owned progress must use explicit cycle finalization") if cursor is None: cursor = SyncCursor( organization_id=fence.organization_id, diff --git a/backend/airweave/domains/entities/canonical/tests/test_cycles.py b/backend/airweave/domains/entities/canonical/tests/test_cycles.py new file mode 100644 index 0000000000..e371eb8336 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_cycles.py @@ -0,0 +1,358 @@ +"""Real PostgreSQL proofs for the single durable cycle boundary and scoped recovery.""" + +import asyncio +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +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, complete_cycle +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, +) +from airweave.domains.entities.canonical.scan_models import ( + BeginScan, + CommitScanPage, + ReconcileScan, + ScanContinuation, +) +from airweave.domains.entities.canonical.store import CanonicalStoreError, StaleWriter +from airweave.models.entity import Entity +from airweave.models.sync_cursor import SyncCursor + +CONFIG = CycleConfiguration( + fingerprint="a" * 64, root_record_type="channel", child_record_types=("message",) +) +ROOT = CompletedScope(record_type="channel") +CHILD = CompletedScope(record_type="message", container_id="C1") + + +def record(kind="channel", native_id="C1"): + parent = RecordIdentity(record_type="channel", native_id="C1") if kind == "message" else None + return CaptureRecord( + identity=RecordIdentity( + record_type=kind, native_id=native_id, container_id="C1" if parent else None + ), + parent=parent, + payload={"id": native_id}, + observed_at=datetime.now(timezone.utc), + ) + + +async def cycle(database, service, fence, **kwargs): + async with database() as db: + return await service.begin_cycle( + db, BeginCycle(fence=fence, configuration=CONFIG, **kwargs) + ) + + +async def scan(database, service, fence, active, scope=ROOT, **kwargs): + async with database() as db: + return await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=scope, + cycle_id=active.version.cycle_id, + fingerprint=CONFIG.fingerprint, + **kwargs, + ), + ) + + +async def page(database, service, fence, state, *records, final=False): + 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=records, + continuation=ScanContinuation(value={"next": "synthetic"}), + final=final, + ), + ) + return result.state + + +async def finish(database, service, fence, state): + async with database() as db: + result = 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), + removal_reason="scope_removed" if state.scope == ROOT else "absent", + ), + ) + assert result.state.phase == "complete" + return result.state + + +async def refresh_cycle(database, service, fence): + async with database() as db: + return await service.read_cycle(db, fence) + + +async def complete(database, service, fence): + state = await refresh_cycle(database, service, fence) + async with database() as db: + return await service.complete_cycle(db, CompleteCycle(fence=fence, expected=state.version)) + + +async def fixture_scopes(database, service, fence, *, child_complete=True): + active = await cycle(database, service, fence) + root = await scan(database, service, fence, active) + root = await finish( + database, service, fence, await page(database, service, fence, root, record(), final=True) + ) + child = await scan(database, service, fence, active, CHILD) + child = await page( + database, service, fence, child, record("message", "one"), final=child_complete + ) + if child_complete: + child = await finish(database, service, fence, child) + return active, root, child + + +async def test_lost_final_ack_boundary_and_explicit_next_cycle(database, source): + service, fence = source + active, root, child = await fixture_scopes(database, service, fence) + async with database() as db: + assert "canonical_checkpoint" not in (await db.scalar(select(SyncCursor))).cursor_data + assert (await cycle(database, service, fence)).version.cycle_id == active.version.cycle_id + finished = await complete(database, service, fence) + assert finished.phase == "complete" + assert await cycle(database, service, fence) == finished + async with database() as db: + stored = (await db.scalar(select(SyncCursor))).cursor_data + assert stored["canonical_checkpoint"]["observed_change_sequence"] == 2 + following = await cycle(database, service, fence, expected=finished.version) + assert following.version.cycle_id != finished.version.cycle_id + with pytest.raises(CycleConflict, match="Root enumeration"): + await complete(database, service, fence) + refreshed = await scan(database, service, fence, following, expected=root.version) + assert refreshed.phase == "collecting" and refreshed.version.sweep_id != root.version.sweep_id + with pytest.raises(CycleConflict): + await page(database, service, fence, child, final=True) + + +async def test_retry_refreshes_root_without_resetting_child_pages(database, source): + service, fence = source + active, root, child = await fixture_scopes(database, service, fence, child_complete=False) + 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, + ) + with pytest.raises(StaleWriter): + await page(database, service, fence, child, record("message", "late")) + with pytest.raises(CycleConflict, match="Refresh"): + await page(database, service, newer, child, record("message", "early")) + with pytest.raises(CycleConflict, match="Restart root"): + await scan(database, service, newer, active) + root = await scan(database, service, newer, active, expected=root.version, restart=True) + await finish( + database, service, newer, await page(database, service, newer, root, record(), final=True) + ) + resumed = await scan(database, service, newer, active, CHILD) + assert resumed == child + await finish( + database, + service, + newer, + await page(database, service, newer, resumed, record("message", "two"), final=True), + ) + assert (await complete(database, service, newer)).phase == "complete" + async with database() as db: + messages = ( + await db.scalars(select(Entity).where(Entity.entity_definition_short_name == "message")) + ).all() + assert {row.native_id for row in messages} == {"one", "two"} + assert all(row.deleted_at is None for row in messages) + + +async def test_removed_parent_is_not_pending_and_reappearance_gets_new_scan(database, source): + service, fence = source + active, root, child = await fixture_scopes(database, service, fence, child_complete=False) + with pytest.raises(CycleConflict, match="incomplete child"): + await complete(database, service, fence) + root = await scan(database, service, fence, active, expected=root.version, restart=True) + root = await finish( + database, service, fence, await page(database, service, fence, root, final=True) + ) + with pytest.raises(CycleConflict, match="visible parent"): + await page(database, service, fence, child, final=True) + finished = await complete(database, service, fence) + following = await cycle(database, service, fence, expected=finished.version) + root = await scan(database, service, fence, following, expected=root.version) + await finish( + database, service, fence, await page(database, service, fence, root, record(), final=True) + ) + replacement = await scan(database, service, fence, following, CHILD, expected=child.version) + assert replacement.phase == "collecting" + assert replacement.version.sweep_id != child.version.sweep_id + with pytest.raises(CycleConflict): + await scan( + database, service, fence, active, CHILD, expected=replacement.version, restart=True + ) + + +async def test_configuration_unknown_scope_and_foreign_cycle_are_rejected(database, source): + service, fence = source + active = await cycle(database, service, fence) + async with database() as db: + with pytest.raises(CycleConflict, match="configuration changed"): + await service.begin_cycle( + db, + BeginCycle( + fence=fence, configuration=CONFIG.model_copy(update={"fingerprint": "b" * 64}) + ), + ) + with pytest.raises(CycleConflict, match="not declared"): + await scan(database, service, fence, active, CompletedScope(record_type="arbitrary")) + foreign = active.model_copy( + update={"version": active.version.model_copy(update={"cycle_id": uuid4()})} + ) + with pytest.raises(CycleConflict, match="active cycle"): + await scan(database, service, fence, foreign) + + +async def test_finalization_rollback_and_concurrent_cas(database, source): + service, fence = source + await fixture_scopes(database, service, fence) + active = await refresh_cycle(database, service, fence) + request = CompleteCycle(fence=fence, expected=active.version) + async with database() as db: + with pytest.raises(RuntimeError, match="after flush"): + async with UnitOfWork(db): + sync = await service.store._fenced_sync(db, fence) + await complete_cycle(db, sync, request) + raise RuntimeError("synthetic failure after flush") + assert await refresh_cycle(database, service, fence) == active + async with database() as db: + assert "canonical_checkpoint" not in (await db.scalar(select(SyncCursor))).cursor_data + + async def publish(): + async with database() as db: + return await service.complete_cycle(db, request) + + results = await asyncio.gather(publish(), publish(), return_exceptions=True) + assert sum(isinstance(result, CycleConflict) for result in results) == 1 + assert (await refresh_cycle(database, service, fence)).phase == "complete" + + +@pytest.mark.parametrize("finished", [False, True]) +async def test_legacy_checkpoint_cannot_erase_cycle(database, source, finished): + service, fence = source + await fixture_scopes(database, service, fence) + if finished: + await complete(database, service, fence) + before = await refresh_cycle(database, service, fence) + async with database() as db: + with pytest.raises(CanonicalStoreError, match="explicit cycle finalization"): + await service.save_checkpoint(db, fence, {"channel_ids": []}) + assert await refresh_cycle(database, service, fence) == before + + +@pytest.mark.parametrize("finished", [False, True]) +@pytest.mark.parametrize("operation", ["create", "update", "delete"]) +async def test_legacy_crud_cannot_erase_active_or_completed_cycle( + database, source, finished, operation +): + from airweave.core.context import BaseContext + from airweave.crud.crud_sync_cursor import sync_cursor + from airweave.schemas.organization import Organization + from airweave.schemas.sync_cursor import SyncCursorCreate + + service, fence = source + await fixture_scopes(database, service, fence) + if finished: + await complete(database, service, fence) + before = await refresh_cycle(database, service, fence) + now = datetime.now(timezone.utc) + ctx = BaseContext( + organization=Organization( + id=fence.organization_id, name="Synthetic owner", created_at=now, modified_at=now + ) + ) + async with database() as db: + with pytest.raises(ValueError, match="Cycle-owned"): + if operation == "create": + await sync_cursor.create_or_update( + db, sync_id=fence.sync_id, ctx=ctx, obj_in=SyncCursorCreate(cursor_data={}) + ) + elif operation == "update": + await sync_cursor.update_cursor_data( + db, sync_id=fence.sync_id, ctx=ctx, cursor_data={} + ) + else: + await sync_cursor.delete_by_sync_id(db, sync_id=fence.sync_id, ctx=ctx) + assert await refresh_cycle(database, service, fence) == before + + +async def test_child_pages_cannot_omit_visibility_parent_and_root_removal_is_private( + database, source +): + from airweave.domains.entities.canonical.scan_store import ScanConflict + + service, fence = source + active, root, child = await fixture_scopes(database, service, fence, child_complete=False) + with pytest.raises(ScanConflict, match="parent identity"): + await page( + database, + service, + fence, + child, + record("message", "unscoped").model_copy(update={"parent": None}), + ) + root = await scan(database, service, fence, active, expected=root.version, restart=True) + root = await page(database, service, fence, root, final=True) + async with database() as db: + result = await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=ROOT, + cycle_id=root.cycle_id, + expected=root.version, + observed_at=datetime.now(timezone.utc), + ), + ) + assert result.capture.changes[0].record.removal_reason == "scope_removed" + + +@pytest.mark.parametrize("completeness", ["partial", "metadata_only"]) +async def test_visible_metadata_parent_still_requires_child_scan(database, source, completeness): + service, fence = source + active = await cycle(database, service, fence) + root = await scan(database, service, fence, active) + root = await page( + database, + service, + fence, + root, + record().model_copy(update={"completeness": completeness}), + final=True, + ) + await finish(database, service, fence, root) + with pytest.raises(CycleConflict, match="incomplete child"): + await complete(database, service, fence) diff --git a/backend/airweave/domains/entities/canonical/tests/test_scans.py b/backend/airweave/domains/entities/canonical/tests/test_scans.py index 7dff08f693..4ed96de94c 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_scans.py +++ b/backend/airweave/domains/entities/canonical/tests/test_scans.py @@ -7,6 +7,12 @@ 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, @@ -27,6 +33,18 @@ 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, @@ -157,16 +175,31 @@ async def test_partial_cannot_reconcile_and_expired_cursor_restarts_sweep(databa 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(ScanConflict, match="unfinished"): - await begin(database, service, fence, expected=state.version) + 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 - next_cycle = await begin(database, service, fence, expected=complete.version) + 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(ScanConflict, match="unfinished"): + with pytest.raises(CycleConflict): await begin( database, service, @@ -199,19 +232,11 @@ async def test_scope_identity_null_and_empty_are_distinct_and_pages_cannot_cross 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: - other = await service.begin_scan( - db, - BeginScan( - fence=fence, - scope=CompletedScope(record_type="event", container_id=""), - cycle_id=state.cycle_id, - fingerprint=FINGERPRINT, - ), - ) - assert state.version.sweep_id != other.version.sweep_id - async with database() as db: - assert await db.scalar(select(func.count()).select_from(CaptureScan)) == 2 + assert await db.scalar(select(func.count()).select_from(CaptureScan)) == 1 assert await db.scalar(select(func.count()).select_from(Entity)) == 0 @@ -248,32 +273,3 @@ async def test_scan_migration_preserves_existing_legacy_records(database, source 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" - - -async def test_root_refresh_does_not_restart_completed_child_in_same_cycle(database, source): - service, fence = source - child = await begin(database, service, fence) - final = await page(database, service, fence, child, final=True) - complete = (await reconcile(database, service, fence, final.state)).state - root_scope = CompletedScope(record_type="channel") - async with database() as db: - root = await service.begin_scan( - db, - BeginScan( - fence=fence, scope=root_scope, cycle_id=child.cycle_id, fingerprint=FINGERPRINT - ), - ) - async with database() as db: - refreshed = await service.begin_scan( - db, - BeginScan( - fence=fence, - scope=root_scope, - cycle_id=child.cycle_id, - fingerprint=FINGERPRINT, - expected=root.version, - restart=True, - ), - ) - assert refreshed.version.sweep_id != root.version.sweep_id - assert await begin(database, service, fence, cycle=child.cycle_id) == complete From b19e6553923bfe3528d41885f2cb881e23a765ad Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:01:27 -0700 Subject: [PATCH 039/266] Allow explicit fenced CAS restart of unfinished capture cycles --- .../entities/canonical/cycle_models.py | 9 +++ .../domains/entities/canonical/cycle_store.py | 15 +++++ .../domains/entities/canonical/service.py | 14 ++++- .../entities/canonical/tests/test_cycles.py | 60 +++++++++++++++++++ 4 files changed, 97 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 12920d42ed..201a546d1c 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -62,3 +62,12 @@ class CompleteCycle(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) fence: WriterFence expected: CycleVersion + + +class RestartCycle(BaseModel): + """Explicitly abandon the exact active cycle without certifying its partial capture.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + expected: CycleVersion + configuration: CycleConfiguration diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 71c2a18a64..333738655b 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -13,6 +13,7 @@ CaptureCycle, CompleteCycle, CycleVersion, + RestartCycle, ) from airweave.domains.entities.canonical.requests import CompletedScope, WriterFence from airweave.domains.entities.canonical.store import CanonicalStoreError, content_is_available @@ -199,3 +200,17 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - } await db.flush() return state + + +async def restart_cycle(db: AsyncSession, request: RestartCycle) -> CaptureCycle: + """Retain captured records/checkpoint; invalidate old scans with a fresh cycle UUID.""" + cursor, previous = await attest_cycle(db, request.fence, request.expected.cycle_id) + if previous.version != request.expected: + raise CycleConflict("Cycle changed before explicit restart") + state = CaptureCycle( + version=CycleVersion(cycle_id=uuid4(), revision=previous.version.revision + 1), + configuration=request.configuration, + ) + persist(cursor, state) + await db.flush() + return state diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index 87e0ede261..bfeebd3af4 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -6,12 +6,18 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork -from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle, CompleteCycle +from airweave.domains.entities.canonical.cycle_models import ( + BeginCycle, + CaptureCycle, + CompleteCycle, + RestartCycle, +) from airweave.domains.entities.canonical.cycle_store import ( begin_cycle, complete_cycle, cursor_row, cycle_state, + restart_cycle, ) from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult from airweave.domains.entities.canonical.requests import ( @@ -146,3 +152,9 @@ async def complete_cycle(self, db: AsyncSession, request: CompleteCycle) -> Capt async with UnitOfWork(db): sync = await self.store._fenced_sync(db, request.fence) return await complete_cycle(db, sync, request) + + async def restart_cycle(self, db: AsyncSession, request: RestartCycle) -> CaptureCycle: + """Explicitly replace active cycle state under the writer fence and exact CAS.""" + async with UnitOfWork(db): + await self.store._fenced_sync(db, request.fence) + return await restart_cycle(db, request) diff --git a/backend/airweave/domains/entities/canonical/tests/test_cycles.py b/backend/airweave/domains/entities/canonical/tests/test_cycles.py index e371eb8336..3b744a8b8f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_cycles.py +++ b/backend/airweave/domains/entities/canonical/tests/test_cycles.py @@ -356,3 +356,63 @@ async def test_visible_metadata_parent_still_requires_child_scan(database, sourc await finish(database, service, fence, root) with pytest.raises(CycleConflict, match="incomplete child"): await complete(database, service, fence) + + +async def test_explicit_restart_preserves_originals_and_last_checkpoint(database, source): + from airweave.domains.entities.canonical.cycle_models import RestartCycle + + service, fence = source + # A previous successful checkpoint must not be advanced by abandoning partial work. + async with database() as db: + await service.save_checkpoint(db, fence, {"previous_success": "synthetic"}) + active, root, child = await fixture_scopes(database, service, fence, child_complete=False) + active = await refresh_cycle(database, service, fence) + replacement_config = CONFIG.model_copy(update={"fingerprint": "b" * 64}) + request = RestartCycle(fence=fence, expected=active.version, configuration=replacement_config) + async with database() as db: + replacement = await service.restart_cycle(db, request) + assert replacement.version.cycle_id != active.version.cycle_id + assert replacement.phase == "active" and replacement.root_writer_attempt_id is None + async with database() as db: + with pytest.raises(CycleConflict): + await service.restart_cycle(db, request) + with pytest.raises(CycleConflict): + await page(database, service, fence, child, final=True) + with pytest.raises(CycleConflict, match="Root enumeration"): + await complete(database, service, fence) + async with database() as db: + saved = (await db.scalar(select(SyncCursor))).cursor_data + assert saved["previous_success"] == "synthetic" + assert saved["canonical_checkpoint"]["observed_change_sequence"] == 0 + original = await db.scalar(select(Entity).where(Entity.native_id == "one")) + assert original.deleted_at is None and original.source_payload == {"id": "one"} + async with database() as db: + next_root = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=ROOT, + cycle_id=replacement.version.cycle_id, + fingerprint=replacement_config.fingerprint, + expected=root.version, + ), + ) + assert next_root.phase == "collecting" + + +async def test_explicit_restart_rejects_completed_boundary_and_stale_revision(database, source): + from airweave.domains.entities.canonical.cycle_models import RestartCycle + + service, fence = source + active, _, _ = await fixture_scopes(database, service, fence) + async with database() as db: + with pytest.raises(CycleConflict, match="changed"): + await service.restart_cycle( + db, RestartCycle(fence=fence, expected=active.version, configuration=CONFIG) + ) + finished = await complete(database, service, fence) + async with database() as db: + with pytest.raises(CycleConflict, match="active cycle"): + await service.restart_cycle( + db, RestartCycle(fence=fence, expected=finished.version, configuration=CONFIG) + ) From 63f4718db4fdb915ee7d08cb2f03776ed7a86c91 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:16:19 -0700 Subject: [PATCH 040/266] test: preserve Gmail locator observation provenance --- .../canonical/tests/test_gmail_observation.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py new file mode 100644 index 0000000000..f12ff32f40 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py @@ -0,0 +1,75 @@ +"""Raw Gmail locator observations retain provenance even when blob bytes match.""" + +from copy import deepcopy +from datetime import datetime, timezone +from hashlib import sha256 +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore +from airweave.domains.entities.canonical.requests import ( + BlobReference, + CaptureRecord, + RecordIdentity, +) +from airweave.domains.entities.canonical.store import capture_fingerprint +from airweave.domains.entities.canonical.tests.helpers import capture, publish_prepared + +pytestmark = pytest.mark.integration + + +async def test_locator_rotation_preserves_originals_and_requires_current_publication( + database, source +): + service, fence = source + digest = sha256(b"identical message body").hexdigest() + payload = { + "id": "synthetic-message", + "threadId": "synthetic-thread", + "internalDate": "1700000000000", + "labelIds": ["INBOX"], + "payload": { + "mimeType": "text/plain", + "partId": "0", + "body": {"size": 22, "attachmentId": "locator-A"}, + }, + } + original = CaptureRecord( + identity=RecordIdentity(record_type="message", native_id=payload["id"]), + payload=payload, + observed_at=datetime.now(timezone.utc), + blobs=( + BlobReference( + key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=22, + media_type="text/plain", + source_path="/payload/body", + ), + ), + ) + rotated_payload = deepcopy(payload) + rotated_payload["payload"]["body"]["attachmentId"] = "locator-B" + rotated = original.model_copy(update={"payload": rotated_payload}) + assert capture_fingerprint(original) != capture_fingerprint(rotated) + first = await capture(database, service, fence, original) + projections = CanonicalProjectionStore() + async with database() as db: + work = (await projections.pending(db, fence.organization_id, fence.sync_id))[0] + assert await publish_prepared(projections, db, work, uuid4(), 1) + async with database() as db: + assert not await projections.pending(db, fence.organization_id, fence.sync_id) + second = await capture(database, service, fence, rotated) + async with database() as db: + pending = await projections.pending(db, fence.organization_id, fence.sync_id) + history = await service.store.changes(db, fence.organization_id, fence.sync_id) + assert first.changes[0].record.revision == 1 + assert second.changes[0].record.revision == 2 + assert len(pending) == 1 and pending[0].record.revision == 2 + assert [c.record.payload["payload"]["body"]["attachmentId"] for c in history.changes] == [ + "locator-A", + "locator-B", + ] + replay = await capture(database, service, fence, rotated) + assert replay.unchanged == 1 and not replay.changes From b6826861280f91320633231643958ce560391a6d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:22:10 -0700 Subject: [PATCH 041/266] Resume Slack pages through fenced canonical capture cycles --- .../entities/canonical/cycle_models.py | 9 + .../domains/entities/canonical/cycle_store.py | 47 ++ .../domains/entities/canonical/page_source.py | 45 ++ .../domains/entities/canonical/scan_store.py | 62 ++- .../domains/entities/canonical/service.py | 16 + .../canonical/tests/test_slack_membership.py | 23 +- .../canonical/tests/test_slack_recovery.py | 431 ++++++++++++++++++ .../sync_pipeline/canonical_capture.py | 60 ++- .../domains/sync_pipeline/canonical_scan.py | 237 ++++++++++ .../airweave/domains/sync_pipeline/factory.py | 18 +- .../domains/sync_pipeline/orchestrator.py | 26 +- .../sync_pipeline/tests/test_factory.py | 19 +- .../activities/project_canonical_records.py | 3 +- backend/airweave/platform/sources/slack.py | 248 ++++++---- backend/tests/live/canonical_capture.py | 4 + backend/tests/live/provider_lifecycle.py | 50 +- .../platform/sources/test_slack_capture.py | 122 ++--- 17 files changed, 1207 insertions(+), 213 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/page_source.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py create mode 100644 backend/airweave/domains/sync_pipeline/canonical_scan.py diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 201a546d1c..b2916f7f93 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -45,6 +45,7 @@ class CaptureCycle(BaseModel): configuration: CycleConfiguration phase: Literal["active", "complete"] = "active" root_writer_attempt_id: UUID | None = None + completed_job_id: UUID | None = None class BeginCycle(BaseModel): @@ -71,3 +72,11 @@ class RestartCycle(BaseModel): fence: WriterFence expected: CycleVersion configuration: CycleConfiguration + + +class CycleRoot(BaseModel): + """Bounded operational scope discovery, without copying provider payloads.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + id: UUID + native_id: str diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 333738655b..472adb483f 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -12,6 +12,7 @@ BeginCycle, CaptureCycle, CompleteCycle, + CycleRoot, CycleVersion, RestartCycle, ) @@ -185,6 +186,7 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - state = state.model_copy( update={ "phase": "complete", + "completed_job_id": request.fence.job_id, "version": CycleVersion( cycle_id=state.version.cycle_id, revision=state.version.revision + 1 ), @@ -214,3 +216,48 @@ async def restart_cycle(db: AsyncSession, request: RestartCycle) -> CaptureCycle persist(cursor, state) await db.flush() return state + + +async def list_cycle_roots( + db: AsyncSession, + fence: WriterFence, + cycle_id: UUID, + *, + after: UUID | None = None, + missing: bool = False, +) -> tuple[CycleRoot, ...]: + """Read up to 100 roots under the fence; absence checks require the final root page.""" + _, state = await attest_cycle(db, fence, cycle_id) + predicates = [ + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == state.configuration.root_record_type, + Entity.container_id.is_(None), + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + ] + if after is not None: + predicates.append(Entity.id > after) + if missing: + root = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.record_type == state.configuration.root_record_type, + CaptureScan.container_id.is_(None), + CaptureScan.cycle_id == cycle_id, + CaptureScan.phase == "reconciling", + ) + ) + if root is None or state.root_writer_attempt_id != fence.attempt_id: + raise CycleConflict("Root absence checks require completed current-attempt collection") + predicates.append(Entity.last_seen_run_id.is_distinct_from(root.sweep_id)) + elif not await root_ready(db, fence, state): + raise CycleConflict("Root membership must complete before child discovery") + rows = ( + await db.execute( + select(Entity.id, Entity.native_id).where(*predicates).order_by(Entity.id).limit(100) + ) + ).all() + return tuple(CycleRoot(id=row.id, native_id=row.native_id) for row in rows) 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..b59ebc825e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -0,0 +1,45 @@ +"""Opted-in whole-scope pages; the existing capture pipeline owns durable progress.""" + +from typing import Protocol, runtime_checkable + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.domains.entities.canonical.cycle_models import CycleConfiguration +from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope +from airweave.domains.entities.canonical.scan_models import ScanContinuation + + +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) + continuation: ScanContinuation + final: bool = False + + +class InvalidScanContinuation(Exception): + """The provider rejected a saved cursor; restart the whole scope, never reconcile it.""" + + +class ScopeAccessLost(Exception): + """Provider explicitly confirmed access loss for the requested child container.""" + + +@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] + capture_cycle_configuration: CycleConfiguration + + async def capture_page( + self, scope: CompletedScope, continuation: ScanContinuation + ) -> CapturePage: + """Fetch one page without advancing durable state or performing writes.""" + ... + + async def confirm_root_absent(self, native_id: str) -> None: + """Raise unless an omitted previously captured root is confirmed inaccessible.""" + ... diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 7d2e811d29..a6d3b3d00d 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -4,10 +4,10 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import select +from sqlalchemy import or_, select from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.cycle_models import CycleVersion +from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleVersion from airweave.domains.entities.canonical.cycle_store import ( CycleConflict, attest_cycle, @@ -30,8 +30,13 @@ ScanState, ScanVersion, ) -from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.store import ( + CanonicalRecordStore, + CanonicalStoreError, + content_is_available, +) from airweave.models.capture_scan import CaptureScan +from airweave.models.entity import Entity class ScanConflict(CanonicalStoreError): @@ -193,6 +198,8 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: ) if any(record.parent != parent for record in request.records): raise ScanConflict("Child records must retain their declared parent identity") + if is_root and cycle.configuration.child_record_types: + await self._invalidate_revived_children(db, request, cycle) captured = await self.records._capture_locked( db, sync, @@ -206,6 +213,55 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: await db.flush() return ScanResult(state=scan_state(row), capture=captured) + async def _invalidate_revived_children( + self, + db: AsyncSession, + request: CommitScanPage, + cycle: CaptureCycle, + ) -> None: + """Reviving an unavailable parent requires a fresh child scan in this same transaction.""" + upserts = tuple( + record.identity.native_id for record in request.records if record.kind == "upsert" + ) + if not upserts: + return + revived = tuple( + ( + await db.scalars( + select(Entity.native_id).where( + Entity.organization_id == request.fence.organization_id, + Entity.sync_id == request.fence.sync_id, + Entity.entity_definition_short_name == cycle.configuration.root_record_type, + Entity.container_id.is_(None), + Entity.native_id.in_(upserts), + or_(Entity.deleted_at.is_not(None), ~content_is_available()), + ) + ) + ).all() + ) + if not revived: + return + rows = ( + await db.scalars( + select(CaptureScan) + .where( + CaptureScan.organization_id == request.fence.organization_id, + CaptureScan.sync_id == request.fence.sync_id, + CaptureScan.cycle_id == request.cycle_id, + CaptureScan.record_type.in_(cycle.configuration.child_record_types), + CaptureScan.container_id.in_(revived), + ) + .with_for_update() + ) + ).all() + for child in rows: + child.sweep_id = uuid4() + child.revision += 1 + child.phase = "collecting" + child.continuation = {} + child.started_at = datetime.now(timezone.utc) + child.completed_at = None + async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: """A final page is necessary; absence completion is durable and bounded.""" sync = await self.records._fenced_sync(db, request.fence) diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index bfeebd3af4..b83ef65946 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -10,6 +10,7 @@ BeginCycle, CaptureCycle, CompleteCycle, + CycleRoot, RestartCycle, ) from airweave.domains.entities.canonical.cycle_store import ( @@ -17,6 +18,7 @@ complete_cycle, cursor_row, cycle_state, + list_cycle_roots, restart_cycle, ) from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult @@ -158,3 +160,17 @@ async def restart_cycle(self, db: AsyncSession, request: RestartCycle) -> Captur async with UnitOfWork(db): await self.store._fenced_sync(db, request.fence) return await restart_cycle(db, request) + + async def list_cycle_roots( + self, + db: AsyncSession, + fence: WriterFence, + cycle_id: UUID, + *, + after: UUID | None = None, + missing: bool = False, + ) -> tuple[CycleRoot, ...]: + """Read a bounded scope-discovery page after verifying current writer ownership.""" + async with UnitOfWork(db): + await self.store._fenced_sync(db, fence) + return await list_cycle_roots(db, fence, cycle_id, after=after, missing=missing) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py index 55b339b37e..30c6d287d0 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py @@ -10,7 +10,7 @@ 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.sources.slack import SlackSource +from airweave.platform.sources.slack import SlackApiError, SlackSource pytestmark = pytest.mark.integration @@ -36,6 +36,16 @@ async def test_first_completed_membership_hides_prior_partial_channel( 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._get = AsyncMock( + side_effect=( + ConnectionError("membership page failed") + if listing_fails + else [{"channels": []}, SlackApiError("channel_not_found")] + ), + ) pipeline = CanonicalCapturePipeline( service, database, @@ -43,19 +53,12 @@ async def test_first_completed_membership_hides_prior_partial_channel( SlackSource.canonical_record_types, CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), SlackSource.canonical_container_parents, - ) - connector = SlackSource( - auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock() - ) - connector._get = AsyncMock( - side_effect=ConnectionError("membership page failed") if listing_fails else None, - return_value={"channels": []}, + page_source=connector, ) await pipeline.start(ctx) async def scan(): - async for item in connector.generate_observations(cursor=runtime.cursor): - await pipeline.process([item], ctx, runtime) + await pipeline.run_scans(ctx, runtime, AsyncMock()) if listing_fails: with pytest.raises(ConnectionError, match="membership page failed"): 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..718ee516c0 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -0,0 +1,431 @@ +"""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.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"}]} + + +def runner(database, source, responses, attempt=1, service=None): + 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._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, + ) + 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 + + +async def test_cancel_after_history_resumes_pending_reply_without_repeating_history( + database, source +): + first, _, _ = runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()]) + 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, _ = runner(database, source, [ROOT, REPLIES], attempt=2) + await run(second) + assert [call.args[0].split("/")[-1] for call in connector._get.call_args_list] == [ + "conversations.list", + "conversations.replies", + ] + 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, _ = 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, _ = 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, _, _ = 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, _ = 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 + + +async def test_repeated_invalid_cursor_fails_without_checkpoint_or_absence(database, source): + first, _, _ = runner( + database, + source, + [ROOT, HISTORY, SlackApiError("invalid_cursor"), HISTORY, SlackApiError("invalid_cursor")], + ) + 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, _, _ = 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 = 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, _, _ = 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, _ = runner(database, source, [ROOT], attempt=2) + await run(second) + assert connector._get.await_count == 1 + third, connector, _ = 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, _, _ = 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, _, _ = 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, _ = 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, _, _ = 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, _, _ = 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, _, _ = 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, _ = 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, _, _ = 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, _ = 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, _, _ = runner(database, source, first_responses) + await run(first, checkpoint=False) + restore_attempt = 2 + if completed_child: + removed, _, _ = 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, _ = 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, _, _ = runner(database, source, [ROOT, HISTORY, REPLIES]) + await run(first, checkpoint=False) + removed, _, _ = 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, _, _ = 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, _, _ = 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, _ = 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) diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index eb2c16b282..cc26a74720 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -1,6 +1,6 @@ """Capture pipeline: native records commit before independent search projection.""" -from collections.abc import Callable +from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timezone @@ -8,7 +8,9 @@ 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 from airweave.domains.entities.canonical.models import CaptureResult +from airweave.domains.entities.canonical.page_source import CanonicalPageSource from airweave.domains.entities.canonical.requests import ( CaptureBatch, CaptureRecord, @@ -21,6 +23,7 @@ ) from airweave.domains.entities.canonical.service import CanonicalCaptureService from airweave.domains.entities.canonical.source import SourceObservation +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 @@ -39,6 +42,7 @@ def __init__( record_types: tuple[str, ...], attempt: CaptureAttempt, container_parents: dict[str, str] | None = None, + page_source: CanonicalPageSource | None = None, ): """Inject transactional persistence and existing progress-event infrastructure.""" if not record_types: @@ -55,6 +59,15 @@ def __init__( ) if parent in self._container_parents: raise ValueError("Only root-container visibility relationships are supported") + self.page_source = page_source + self._cycle: CaptureCycle | None = None + if page_source is not None: + config = page_source.capture_cycle_configuration + if config.root_record_type not in self._record_types or any( + self._container_parents.get(child) != config.root_record_type + for child in config.child_record_types + ): + raise ValueError("Page source cycle must match its declared container topology") self._attempt = attempt self._fence: WriterFence | None = None self._completed_scopes: set[CompletedScope] = set() @@ -73,6 +86,35 @@ async def start(self, sync_context: SyncContext) -> None: 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.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, + self._with_parent, + check_limits, + ).run() + def _writer(self) -> WriterFence: if self._fence is None: raise SyncFailureError("Canonical capture was not activated") @@ -243,6 +285,22 @@ async def _reconcile_parents(self, sync_context: SyncContext, runtime: SyncRunti 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 + async with self._sessions() as db: + self._cycle = await self._service.complete_cycle( + db, CompleteCycle(fence=self._writer(), expected=current.version) + ) + 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: 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..0ab512ed0a --- /dev/null +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -0,0 +1,237 @@ +"""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 +from airweave.domains.entities.canonical.cycle_store import CycleConflict +from airweave.domains.entities.canonical.models import CaptureResult +from airweave.domains.entities.canonical.page_source import ( + CanonicalPageSource, + InvalidScanContinuation, + ScopeAccessLost, +) +from airweave.domains.entities.canonical.requests import ( + CaptureBatch, + CaptureRecord, + CompletedScope, + RecordIdentity, + RemovedScope, + WriterFence, +) +from airweave.domains.entities.canonical.scan_models import ( + BeginScan, + CommitScanPage, + ReconcileScan, + ScanState, +) +from airweave.domains.entities.canonical.service import CanonicalCaptureService + + +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]], + with_parent: Callable[[CaptureRecord], CaptureRecord], + check_limits: Callable[[], Awaitable[None]], + ): + """Reuse the capture service and pipeline progress/guard callbacks.""" + self.service, self.sessions, self.fence, self.source = service, sessions, fence, source + self.progress, self.with_parent, self.check_limits = progress, with_parent, check_limits + + 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 + async with self.sessions() as db: + cycle = await self.service.begin_cycle( + db, BeginCycle(fence=self.fence, configuration=configuration, expected=expected) + ) + root = CompletedScope(record_type=configuration.root_record_type) + await self.scan(cycle, root, refresh_membership=True) + after: UUID | None = None + while True: + async with self.sessions() as db: + roots = await self.service.list_cycle_roots( + db, self.fence, cycle.version.cycle_id, after=after + ) + if not roots: + break + for item in roots: + for child_type in configuration.child_record_types: + scope = CompletedScope(record_type=child_type, container_id=item.native_id) + try: + await self.scan(cycle, scope) + except ScopeAccessLost: + await self.withdraw_root( + configuration.root_record_type, + item.native_id, + configuration.child_record_types, + ) + break + after = roots[-1].id + 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( + self, + cycle: CaptureCycle, + scope: CompletedScope, + *, + refresh_membership: bool = False, + ) -> None: + """Only a provider invalid-cursor response permits one explicit sweep restart.""" + 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 refresh_membership + and cycle.root_writer_attempt_id != self.fence.attempt_id + ) + async with self.sessions() as db: + state = await self.service.begin_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, + ), + ) + restarts = 0 + while state.phase == "collecting": + await self.check_limits() + try: + page = await self.source.capture_page(scope, state.continuation) + except InvalidScanContinuation: + if restarts >= 1: + raise + restarts += 1 + async with self.sessions() as db: + state = 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, + ), + ) + 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.with_parent(r) for r in page.records), + continuation=page.continuation, + final=page.final, + ), + ) + state = result.state + await self.progress(result.capture, page.records) + if state.phase == "reconciling" and refresh_membership: + await self.confirm_omissions(cycle) + await self.reconcile(state) + + async def confirm_omissions(self, cycle: CaptureCycle) -> None: + """A provider error or accessible omission fails before any absence removal.""" + after: UUID | None = None + while True: + async with self.sessions() as db: + roots = await self.service.list_cycle_roots( + db, self.fence, cycle.version.cycle_id, after=after, missing=True + ) + if not roots: + return + for item in roots: + await self.check_limits() + await self.source.confirm_root_absent(item.native_id) + after = roots[-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_root( + self, root_type: str, native_id: str, child_types: tuple[str, ...] + ) -> None: + """Explicit provider access loss hides the parent before bounded child cleanup.""" + observed = datetime.now(timezone.utc) + async with self.sessions() as db: + result = await self.service.capture( + db, + CaptureBatch( + fence=self.fence, + records=( + CaptureRecord( + identity=RecordIdentity(record_type=root_type, native_id=native_id), + payload={"id": native_id}, + kind="delete", + removal_reason="access_revoked", + observed_at=observed, + ), + ), + ), + ) + await self.progress(result, ()) + for child_type in child_types: + while True: + await self.check_limits() + async with self.sessions() as db: + result = await self.service.remove_scope( + db, + self.fence, + RemovedScope( + record_type=child_type, + container_id=native_id, + removal_reason="access_revoked", + observed_at=observed, + ), + ) + await self.progress(result.capture, ()) + if not result.has_more: + break diff --git a/backend/airweave/domains/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index e25c87c509..b9757b832c 100644 --- a/backend/airweave/domains/sync_pipeline/factory.py +++ b/backend/airweave/domains/sync_pipeline/factory.py @@ -32,6 +32,7 @@ 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 @@ -193,7 +194,9 @@ async def create_orchestrator( ) source_entry = self._source_registry.get(sc.short_name) canonical_source = ( - source_result.source if isinstance(source_result.source, CanonicalSource) else None + source_result.source + if isinstance(source_result.source, (CanonicalSource, CanonicalPageSource)) + else None ) destinations = ( [] @@ -245,6 +248,9 @@ async def create_orchestrator( event_bus=self._event_bus, record_types=canonical_source.canonical_record_types, attempt=resolve_capture_attempt(capture_attempt), + page_source=( + canonical_source if isinstance(canonical_source, CanonicalPageSource) else None + ), container_parents=( canonical_source.canonical_container_parents if isinstance(canonical_source, ContainerScopedSource) @@ -343,8 +349,12 @@ def _build_stream( runtime: SyncRuntime, source_result: SourceBuildResult, sync_context: SyncContext, - ) -> AsyncSourceStream: - """Build the async source stream from the source generator.""" + ) -> AsyncSourceStream | None: + """Page sources are driven sequentially; legacy generators retain their stream.""" + if isinstance(runtime.source, CanonicalPageSource): + 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, @@ -382,7 +392,7 @@ async def _build_source( """Build source instance, cursor, file service, and node selections.""" if execution_config and execution_config.behavior.replay_from_arf: source_class = self._source_registry.get(source_connection.short_name).source_class_ref - if isinstance(source_class, CanonicalSource): + if isinstance(source_class, (CanonicalSource, CanonicalPageSource)): raise ValueError( "Canonical records must be reindexed from Postgres, not mutable ARF" ) diff --git a/backend/airweave/domains/sync_pipeline/orchestrator.py b/backend/airweave/domains/sync_pipeline/orchestrator.py index f9371e5faf..888b4064c4 100644 --- a/backend/airweave/domains/sync_pipeline/orchestrator.py +++ b/backend/airweave/domains/sync_pipeline/orchestrator.py @@ -56,7 +56,7 @@ def __init__( self, entity_pipeline: EntityPipeline | CanonicalCapturePipeline, worker_pool: AsyncWorkerPool, - stream: AsyncSourceStream, + stream: AsyncSourceStream | None, sync_context: SyncContext, runtime: SyncRuntime, access_control_pipeline: AccessControlPipeline, @@ -205,10 +205,21 @@ async def _start_sync(self) -> None: ) if self.runtime.canonical_capture is not None: await self.runtime.canonical_capture.start(self.sync_context) - await self.stream.start() + 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} " @@ -299,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, @@ -779,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/tests/test_factory.py b/backend/airweave/domains/sync_pipeline/tests/test_factory.py index f7b3a1cfa0..2b0bd6d395 100644 --- a/backend/airweave/domains/sync_pipeline/tests/test_factory.py +++ b/backend/airweave/domains/sync_pipeline/tests/test_factory.py @@ -134,7 +134,7 @@ async def test_create_orchestrator_raises_when_source_connection_missing(): @pytest.mark.asyncio -@pytest.mark.parametrize("canonical", [False, True]) +@pytest.mark.parametrize("canonical", [False, True, "page"]) async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): """entity_repo is forwarded to EntityActionResolver and EntityPipeline.""" entity_repo = MagicMock() @@ -180,7 +180,14 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): ): mock_source = MagicMock() mock_source.generate_entities = MagicMock(return_value=AsyncMock()) - if canonical: + if canonical == "page": + from airweave.domains.sources.token_providers.static import StaticTokenProvider + from airweave.platform.sources.slack import SlackSource + + mock_source = SlackSource( + auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=MagicMock() + ) + elif canonical: mock_source.canonical_record_types = ("event",) mock_source.generate_observations = MagicMock(return_value=AsyncMock()) mock_build_source.return_value = SourceBuildResult( @@ -208,8 +215,12 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): assert isinstance(orchestrator.entity_pipeline, CanonicalCapturePipeline) mock_build_destinations.assert_not_called() mock_disp_builder.assert_not_called() - mock_source.generate_entities.assert_not_called() - mock_source.generate_observations.assert_called_once() + if canonical == "page": + assert orchestrator.stream is None + assert orchestrator.entity_pipeline.page_source is mock_source + 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 diff --git a/backend/airweave/domains/temporal/activities/project_canonical_records.py b/backend/airweave/domains/temporal/activities/project_canonical_records.py index 9d5756f6eb..422f17ac86 100644 --- a/backend/airweave/domains/temporal/activities/project_canonical_records.py +++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py @@ -8,6 +8,7 @@ from airweave.core.logging import logger from airweave.db.session import get_db_context +from airweave.domains.entities.canonical.page_source import CanonicalPageSource 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 @@ -49,7 +50,7 @@ async def run(self, organization_id: str, sync_id: str, after_id: str | None = N return ProjectionBatchResult().model_dump(mode="json") source_name, collection_id = row entry = self.source_registry.get(source_name) - if not isinstance(entry.source_class_ref, CanonicalSource): + if not isinstance(entry.source_class_ref, (CanonicalSource, CanonicalPageSource)): return ProjectionBatchResult().model_dump(mode="json") # Do not construct an index client for sources with no pending records. pending = await CanonicalProjectionStore().pending( diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index ebba1d16db..35283e102b 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -2,25 +2,31 @@ from __future__ import annotations +import hashlib import re -from contextlib import aclosing from datetime import datetime, timezone from functools import partial from typing import Any, AsyncGenerator, Dict, List, Optional +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError 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.entities.canonical.cycle_models import CycleConfiguration +from airweave.domains.entities.canonical.page_source import ( + CapturePage, + InvalidScanContinuation, + ScopeAccessLost, +) from airweave.domains.entities.canonical.requests import ( CaptureRecord, CompletedScope, RecordIdentity, - RemovedScope, - StartedScope, ) +from airweave.domains.entities.canonical.scan_models import ScanContinuation from airweave.domains.sources.exceptions import SourceAuthError, SourceRateLimitError from airweave.domains.sources.token_providers.protocol import ( SourceAuthProvider, @@ -91,6 +97,32 @@ def __init__(self, code: str): super().__init__(f"Slack request failed: {code}") +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) + + @source( name="Slack", short_name="slack", @@ -112,6 +144,11 @@ class SlackSource(BaseSource): canonical_record_types = ("channel", "message") canonical_container_parents = {"message": "channel"} + capture_cycle_configuration = CycleConfiguration( + fingerprint=hashlib.sha256(b"slack:whole-conversations:history-and-replies:v1").hexdigest(), + root_record_type="channel", + child_record_types=("message",), + ) @classmethod async def create( @@ -167,33 +204,6 @@ async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict: raise SlackApiError(error) return payload - async def _capture_pages( - self, operation: str, item_key: str, params: dict[str, Any] - ) -> AsyncGenerator[dict, None]: - """Drain a cursor endpoint; incomplete pagination is a failed capture.""" - page_cursor = "" - seen_cursors: set[str] = set() - while True: - page_params = {**params, "limit": 100} - if page_cursor: - page_params["cursor"] = page_cursor - page = await self._get(f"https://slack.com/api/{operation}", page_params) - items = page.get(item_key) - if not isinstance(items, list): - raise ValueError("Slack capture response lacks its record list") - for item in items: - if not isinstance(item, dict): - raise ValueError("Slack capture returned an invalid record") - yield item - page_cursor = page.get("response_metadata", {}).get("next_cursor", "").strip() - if not page_cursor: - if page.get("has_more"): - raise ValueError("Slack capture has more records without a cursor") - return - if page_cursor in seen_cursors: - raise ValueError("Slack capture repeated a pagination cursor") - seen_cursors.add(page_cursor) - @staticmethod def _capture_message(message: dict, channel_id: str) -> CaptureRecord: """Keep the complete native response; file bytes are explicitly not captured yet.""" @@ -206,83 +216,129 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: completeness="partial" if message.get("files") else "complete", ) - async def generate_observations( - self, - *, - cursor: SyncCursor | None = None, - files: FileService | None = None, - node_selections: list[NodeSelectionData] | None = None, - ) -> AsyncGenerator[CaptureRecord | CompletedScope | StartedScope | RemovedScope, None]: - """Reconcile complete accessible histories, preserving threads and channel identity. - - Slack offers no durable deletion cursor here. Every run re-enumerates history; - a failure aborts reconciliation. Lost channel access is never called deletion. - """ - if node_selections: - raise ValueError("Slack capture does not yet support selected conversation scopes") - seen_channels: set[str] = set() - yield StartedScope(record_type="channel") - async with aclosing( - self._capture_pages( + async def capture_page( + self, scope: CompletedScope, continuation: ScanContinuation + ) -> CapturePage: + """Fetch one page; record and nested reply progress are committed by the pipeline.""" + try: + return await self._validated_capture_page(scope, continuation) + except ValidationError: + # Validation diagnostics may contain private provider values. + raise ValueError("Slack capture returned invalid page or continuation data") from None + + 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, ) - ) as channels: - async for channel in channels: - channel_id = channel["id"] - seen_channels.add(channel_id) - yield CaptureRecord( - identity=RecordIdentity(record_type="channel", native_id=channel_id), - payload=channel, + records = tuple( + CaptureRecord( + identity=RecordIdentity(record_type="channel", native_id=item["id"]), + payload=item, observed_at=datetime.now(timezone.utc), ) - yield StartedScope(record_type="message", container_id=channel_id) - async with aclosing( - self._capture_pages( - "conversations.history", "messages", {"channel": channel_id} + for item in items + ) + return CapturePage( + records=records, + final=not following, + continuation=ScanContinuation( + value=SlackRootContinuation(cursor=following, recent=recent).model_dump( + mode="json" ) - ) as history: - async for message in history: - yield self._capture_message(message, channel_id) - if message.get("reply_count", 0): - async with aclosing( - self._capture_pages( - "conversations.replies", - "messages", - {"channel": channel_id, "ts": message["ts"]}, - ) - ) as replies: - async for reply in replies: - yield self._capture_message(reply, channel_id) - yield CompletedScope(record_type="message", container_id=channel_id) - - previous_channels = set(cursor.get().get("channel_ids", [])) if cursor else set() - for channel_id in previous_channels - seen_channels: - await self._confirm_channel_access_lost(channel_id) - observed_at = datetime.now(timezone.utc) - yield CaptureRecord( - identity=RecordIdentity(record_type="channel", native_id=channel_id), - payload={"id": channel_id}, - kind="delete", - removal_reason="access_revoked", - observed_at=observed_at, + ), ) - yield RemovedScope( - record_type="message", - container_id=channel_id, - removal_reason="access_revoked", - observed_at=observed_at, + 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 {})}, ) - # Reconcile parents even if an earlier partial run saved no channel checkpoint. - yield CompletedScope(record_type="channel") - if cursor is not None: - cursor.update(channel_ids=sorted(seen_channels)) + except SlackApiError as exc: + if exc.code == "invalid_cursor": + raise InvalidScanContinuation("Slack page cursor expired") from exc + raise + 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_channel_access_lost(self, channel_id: str) -> None: - """Require provider confirmation before hiding an omitted prior channel.""" + async def confirm_root_absent(self, native_id: str) -> None: + """Require provider confirmation before hiding an omitted prior conversation.""" try: - await self._get("https://slack.com/api/conversations.info", {"channel": channel_id}) + 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 diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index 5bc9090111..e906b5c67b 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -80,6 +80,7 @@ async def fetch_records(name, account, expected_email, key, fence, root): """Exercise actual source capture, bounded before any complete-scope claim.""" from provider_sample import rest_source, wispr_source + from airweave.domains.entities.canonical.page_source import CanonicalPageSource from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.entities.canonical.source import ContainerScopedSource @@ -94,6 +95,8 @@ async def fetch_records(name, account, expected_email, key, fence, root): records = {} observations = 0 async with asyncio.timeout(180), connection as (source, identity_verification): + if isinstance(source, CanonicalPageSource): + raise ValueError("Page sources require the production provider_lifecycle.py harness") parents = ( source.canonical_container_parents if isinstance(source, ContainerScopedSource) else {} ) @@ -295,6 +298,7 @@ async def main(): "0002_projection_publication.py", "0003_mail_thread_index.py", "0004_projection_generation.py", + "0005_capture_scan.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 21f8281d54..1a652e8cca 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -27,6 +27,7 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from airweave.adapters.event_bus.fake import FakeEventBus +from airweave.domains.entities.canonical.page_source import CanonicalPageSource from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope from airweave.domains.entities.canonical.service import CanonicalCaptureService from airweave.domains.entities.canonical.store import CanonicalRecordStore @@ -72,6 +73,26 @@ async def write_file(self, path, content): self.written_bytes += len(content) +class BoundedPageSource: + """Test budget around the real page adapter; production driver still owns all progress.""" + + def __init__(self, source, counters, record_limit): + self.source, self.counters, self.record_limit = source, counters, record_limit + self.canonical_record_types = source.canonical_record_types + self.canonical_container_parents = source.canonical_container_parents + self.capture_cycle_configuration = source.capture_cycle_configuration + + async def capture_page(self, scope, continuation): + page = await self.source.capture_page(scope, continuation) + self.counters["records_observed"] += len(page.records) + if self.counters["records_observed"] > self.record_limit: + raise BudgetExceeded("records") + return page + + async def confirm_root_absent(self, native_id): + await self.source.confirm_root_absent(native_id) + + async def bounded_observations(source, cursor, files, counters, record_limit): """A breached test budget fails enumeration before scope completion.""" async with aclosing(source.generate_observations(cursor=cursor, files=files)) as stream: @@ -147,12 +168,12 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert saved["canonical_page_token"] assert counters["started"] == counters["completed"] == (0 if loaded else 1) assert bool(counters["resumed_changes_requests"]) == loaded - elif name in {"slack", "wispr"}: + elif name == "slack": + assert saved["canonical_cycle"]["phase"] == "complete" + assert saved["canonical_cycle"]["completed_job_id"] + elif name == "wispr": assert saved is None - if name == "wispr": - assert counters["started"] == counters["completed"] == 0 - else: - assert counters["started"] == counters["completed"] > 0 + assert counters["started"] == counters["completed"] == 0 return else: assert saved["canonical_query"] == manifest["query"] @@ -301,6 +322,11 @@ async def db_context(): assert source.cursor_class == cursor_schema bus = FakeEventBus() attempt = CaptureAttempt(id=uuid4(), number=1) + page_source = ( + BoundedPageSource(source, counters, manifest["record_limit"]) + if isinstance(source, CanonicalPageSource) + else None + ) pipeline = CanonicalCapturePipeline( CanonicalCaptureService(CanonicalRecordStore()), sessions, @@ -308,6 +334,7 @@ async def db_context(): source.canonical_record_types, attempt, container_parents=getattr(source, "canonical_container_parents", {}), + page_source=page_source, ) runtime = SyncRuntime( source=source, @@ -318,9 +345,15 @@ async def db_context(): runner = SyncOrchestrator( entity_pipeline=pipeline, worker_pool=AsyncWorkerPool(logger=logger), - stream=AsyncSourceStream( - bounded_observations(source, cursor, files, counters, manifest["record_limit"]), - logger=logger, + stream=( + None + if page_source is not None + else AsyncSourceStream( + bounded_observations( + source, cursor, files, counters, manifest["record_limit"] + ), + logger=logger, + ) ), sync_context=ctx, runtime=runtime, @@ -463,6 +496,7 @@ async def main(): "0002_projection_publication.py", "0003_mail_thread_index.py", "0004_projection_generation.py", + "0005_capture_scan.py", ): await connection.run_sync(harness.migrate, migration) organization_id, sync_id = uuid4(), uuid4() diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index a09951dad6..84270d8692 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -1,12 +1,13 @@ -"""Thread preservation and failure-safe Slack capture.""" +"""Slack native page contracts; durable recovery is tested separately against PostgreSQL.""" from unittest.mock import AsyncMock, MagicMock import pytest -from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope +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.sources.slack import SlackSource +from airweave.platform.sources.slack import SlackApiError, SlackSource def source(): @@ -16,56 +17,34 @@ def source(): @pytest.mark.asyncio -async def test_channel_pagination_threads_and_native_payload(monkeypatch): +async def test_page_commits_history_before_threads_and_preserves_native_fields(): connector = source() - calls = [] - - async def get(url, params): - calls.append((url, params)) - if url.endswith("conversations.list"): - return {"channels": [{"id": "C1", "unknown_native_field": "kept"}]} - if url.endswith("conversations.history"): - if params.get("cursor"): - return {"messages": [{"ts": "2", "text": "second", "files": [{"id": "F1"}]}]} - return { + connector._get = AsyncMock( + side_effect=[ + {"channels": [{"id": "C1", "unknown_native_field": "kept"}]}, + { "messages": [{"ts": "1", "reply_count": 1, "blocks": [{"type": "rich_text"}]}], "response_metadata": {"next_cursor": "next"}, "has_more": True, - } - return { - "messages": [ - {"ts": "1", "reply_count": 1}, - {"ts": "1.1", "thread_ts": "1", "text": "reply"}, - ] - } - - monkeypatch.setattr(connector, "_get", get) - observations = [x async for x in connector.generate_observations()] - records = [x for x in observations if isinstance(x, CaptureRecord)] - assert records[0].payload["unknown_native_field"] == "kept" - assert any(x.identity.native_id == "1.1" and x.payload["thread_ts"] == "1" for x in records) - assert records[-1].completeness == "partial" - assert CompletedScope(record_type="message", container_id="C1") in observations - assert observations[-1] == CompletedScope(record_type="channel") - assert calls[0][1]["types"] == "public_channel,private_channel,im,mpim" - - -@pytest.mark.asyncio -async def test_failed_thread_never_completes_scope(monkeypatch): - connector = source() - get = AsyncMock( - side_effect=[ - {"channels": [{"id": "C1"}]}, - {"messages": [{"ts": "1", "reply_count": 1}]}, - RuntimeError("permission lost"), + }, + {"messages": [{"ts": "1.1", "thread_ts": "1", "text": "reply"}]}, + {"messages": [{"ts": "2", "files": [{"id": "F1"}]}]}, ] ) - monkeypatch.setattr(connector, "_get", get) - emitted = [] - with pytest.raises(RuntimeError, match="permission"): - async for observation in connector.generate_observations(): - emitted.append(observation) - assert not any(type(x) is CompletedScope for x in emitted) + root = await connector.capture_page(CompletedScope(record_type="channel"), ScanContinuation()) + assert root.final and root.records[0].payload["unknown_native_field"] == "kept" + scope = CompletedScope(record_type="message", container_id="C1") + history = await connector.capture_page(scope, ScanContinuation()) + assert history.records[0].payload["blocks"] == [{"type": "rich_text"}] + assert history.continuation.value["pending_threads"] == ["1"] and not history.final + replies = await connector.capture_page(scope, history.continuation) + assert replies.records[0].payload["thread_ts"] == "1" and not replies.final + last = await connector.capture_page(scope, replies.continuation) + assert last.final and last.records[0].completeness == "partial" + calls = connector._get.call_args_list + assert calls[0].args[1]["types"] == "public_channel,private_channel,im,mpim" + assert calls[2].args[0].endswith("conversations.replies") + assert calls[3].args[1]["cursor"] == "next" @pytest.mark.asyncio @@ -76,49 +55,26 @@ async def test_failed_thread_never_completes_scope(monkeypatch): [{"messages": [], "response_metadata": {"next_cursor": "repeat"}}] * 2, ], ) -async def test_incomplete_pagination_fails(monkeypatch, pages): +async def test_incomplete_pagination_fails(pages): connector = source() - monkeypatch.setattr(connector, "_get", AsyncMock(side_effect=pages)) + connector._get = AsyncMock(side_effect=pages) + progress = ScanContinuation() with pytest.raises(ValueError): - _ = [x async for x in connector._capture_pages("conversations.history", "messages", {})] - - -@pytest.mark.asyncio -async def test_prior_channel_loss_requires_explicit_provider_confirmation(monkeypatch): - from uuid import uuid4 - - from airweave.domains.entities.canonical.requests import RemovedScope - from airweave.domains.syncs.cursors.cursor import SyncCursor - from airweave.platform.sources.slack import SlackApiError - - connector = source() - cursor = SyncCursor(uuid4(), cursor_data={"channel_ids": ["C1"]}) - monkeypatch.setattr( - connector, - "_get", - AsyncMock(side_effect=[{"channels": []}, SlackApiError("channel_not_found")]), - ) - observations = [x async for x in connector.generate_observations(cursor=cursor)] - assert observations[1].kind == "delete" - assert observations[1].removal_reason == "access_revoked" - assert isinstance(observations[2], RemovedScope) - assert cursor.get()["channel_ids"] == [] + for _ in pages: + result = await connector.capture_page( + CompletedScope(record_type="message", container_id="C1"), progress + ) + progress = result.continuation @pytest.mark.asyncio -async def test_omitted_but_accessible_channel_does_not_advance_cursor(monkeypatch): - from uuid import uuid4 - - from airweave.domains.syncs.cursors.cursor import SyncCursor - +async def test_prior_channel_loss_requires_explicit_provider_confirmation(): connector = source() - cursor = SyncCursor(uuid4(), cursor_data={"channel_ids": ["C1"]}) - monkeypatch.setattr( - connector, "_get", AsyncMock(side_effect=[{"channels": []}, {"channel": {"id": "C1"}}]) - ) + connector._get = AsyncMock(side_effect=SlackApiError("channel_not_found")) + await connector.confirm_root_absent("C1") + connector._get = AsyncMock(return_value={"channel": {"id": "C1"}}) with pytest.raises(ValueError, match="omitted"): - _ = [x async for x in connector.generate_observations(cursor=cursor)] - assert cursor.get()["channel_ids"] == ["C1"] + await connector.confirm_root_absent("C1") @pytest.mark.asyncio From 088944842c4ab51bd166306b535491d72e3d8690 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:30:35 -0700 Subject: [PATCH 042/266] Preserve source-audited reasons when withdrawing unavailable scopes --- .../domains/entities/canonical/page_source.py | 20 +++++++++++--- .../domains/entities/canonical/requests.py | 4 ++- .../canonical/tests/test_slack_recovery.py | 27 +++++++++++++++++++ .../domains/sync_pipeline/canonical_scan.py | 16 +++++++---- 4 files changed, 58 insertions(+), 9 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index b59ebc825e..c241add9ed 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -2,10 +2,14 @@ from typing import Protocol, runtime_checkable -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from airweave.domains.entities.canonical.cycle_models import CycleConfiguration -from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + ScopeRemovalReason, +) from airweave.domains.entities.canonical.scan_models import ScanContinuation @@ -23,7 +27,17 @@ class InvalidScanContinuation(Exception): class ScopeAccessLost(Exception): - """Provider explicitly confirmed access loss for the requested child container.""" + """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) @runtime_checkable diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 6cfbb49c7b..4fe249be39 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -6,6 +6,8 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator +ScopeRemovalReason = Literal["scope_removed", "access_revoked"] + class RecordIdentity(BaseModel): """Native identity inside one sync, including provider-required container.""" @@ -111,5 +113,5 @@ class StartedScope(CompletedScope): class RemovedScope(CompletedScope): """Provider-confirmed scope loss; applies to children even if seen in this run.""" - removal_reason: Literal["scope_removed", "access_revoked"] + removal_reason: ScopeRemovalReason observed_at: AwareDatetime diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index 718ee516c0..0b2b9876dc 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -429,3 +429,30 @@ async def test_unavailable_parent_without_tombstone_also_invalidates_child_progr 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 = runner(database, source, [ROOT, HISTORY]) + original_page = pipeline.page_source.capture_page + + async def unavailable(scope, continuation): + 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) + + 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 == {} diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 0ab512ed0a..b249340c68 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -14,6 +14,7 @@ CanonicalPageSource, InvalidScanContinuation, ScopeAccessLost, + ScopeRemovalReason, ) from airweave.domains.entities.canonical.requests import ( CaptureBatch, @@ -79,11 +80,12 @@ async def run(self) -> CaptureCycle: scope = CompletedScope(record_type=child_type, container_id=item.native_id) try: await self.scan(cycle, scope) - except ScopeAccessLost: + except ScopeAccessLost as error: await self.withdraw_root( configuration.root_record_type, item.native_id, configuration.child_record_types, + error.removal_reason, ) break after = roots[-1].id @@ -197,9 +199,13 @@ async def reconcile(self, state: ScanState) -> None: await self.progress(result.capture, ()) async def withdraw_root( - self, root_type: str, native_id: str, child_types: tuple[str, ...] + self, + root_type: str, + native_id: str, + child_types: tuple[str, ...], + removal_reason: ScopeRemovalReason, ) -> None: - """Explicit provider access loss hides the parent before bounded child cleanup.""" + """Withdraw the parent before bounded children, retaining the source-audited reason.""" observed = datetime.now(timezone.utc) async with self.sessions() as db: result = await self.service.capture( @@ -211,7 +217,7 @@ async def withdraw_root( identity=RecordIdentity(record_type=root_type, native_id=native_id), payload={"id": native_id}, kind="delete", - removal_reason="access_revoked", + removal_reason=removal_reason, observed_at=observed, ), ), @@ -228,7 +234,7 @@ async def withdraw_root( RemovedScope( record_type=child_type, container_id=native_id, - removal_reason="access_revoked", + removal_reason=removal_reason, observed_at=observed, ), ) From 6a5b35b97dd6315598248da9b7d648e6dec40ca2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:33:36 -0700 Subject: [PATCH 043/266] Capture selected Linear issue scopes through canonical pages --- .../auth_provider/providers/composio.py | 1 + backend/airweave/platform/configs/config.py | 18 +- backend/airweave/platform/sources/linear.py | 825 ++++++------------ .../linear_queries/attachments.graphql | 9 + .../sources/linear_queries/comments.graphql | 14 + .../sources/linear_queries/identity.graphql | 4 + .../linear_queries/issue-membership.graphql | 3 + .../sources/linear_queries/issues.graphql | 15 + .../sources/linear_queries/teams.graphql | 6 + .../unit/platform/sources/test_linear.py | 268 +----- .../platform/sources/test_linear_capture.py | 307 +++++++ 11 files changed, 621 insertions(+), 849 deletions(-) create mode 100644 backend/airweave/platform/sources/linear_queries/attachments.graphql create mode 100644 backend/airweave/platform/sources/linear_queries/comments.graphql create mode 100644 backend/airweave/platform/sources/linear_queries/identity.graphql create mode 100644 backend/airweave/platform/sources/linear_queries/issue-membership.graphql create mode 100644 backend/airweave/platform/sources/linear_queries/issues.graphql create mode 100644 backend/airweave/platform/sources/linear_queries/teams.graphql create mode 100644 backend/tests/unit/platform/sources/test_linear_capture.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index d3daafc57d..cda796285a 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -116,6 +116,7 @@ async def get_auth_result( "google_calendar": frozenset({"www.googleapis.com"}), "google_drive": frozenset({"www.googleapis.com"}), "slack": frozenset({"slack.com"}), + "linear": frozenset({"api.linear.app"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 428478056e..16dc016873 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -2,6 +2,7 @@ from datetime import datetime, timedelta, timezone from typing import Literal, Optional +from uuid import UUID from pydantic import AwareDatetime, Field, field_validator, model_validator @@ -467,9 +468,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): diff --git a/backend/airweave/platform/sources/linear.py b/backend/airweave/platform/sources/linear.py index 35d40bf5a9..6771e82a02 100644 --- a/backend/airweave/platform/sources/linear.py +++ b/backend/airweave/platform/sources/linear.py @@ -1,42 +1,42 @@ -"""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 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 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.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.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.http_client.retry_helpers import ( retry_if_rate_limit_or_timeout, @@ -46,27 +46,61 @@ 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 = 1 + + +class _Model(BaseModel): + model_config = ConfigDict(extra="ignore") + + +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) -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 _Envelope(_Model): + data: dict[str, JsonValue] | None = None + errors: list[_GraphError] = Field(default_factory=list) -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 _Progress(BaseModel): + model_config = ConfigDict(extra="forbid") + version: Literal[1] = 1 + after: str | None = Field(default=None, min_length=1) @source( @@ -81,534 +115,200 @@ 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( + fingerprint=hashlib.sha256( + json.dumps(fingerprint, sort_keys=True).encode() + ).hexdigest(), + root_record_type="issue", + child_record_types=("comment", "attachment"), + ) + 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} + _URL, headers=headers, json={"query": QUERIES[name], "variables": variables} ) - - 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", - ) - + # 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 - - # ------------------------------------------------------------------ - # 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 - }} - }} - }} - """ - ) - - 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( - 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, - ) - - async def _download_description_attachment( - self, - entity: LinearAttachmentEntity, - 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, - ) - 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}") + 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 - 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}" + 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 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}" + return issue + + async def capture_page( + self, scope: CompletedScope, continuation: ScanContinuation + ) -> 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 = data.get("issue") + # Null without documented error evidence is not a deletion or access-loss proof. + 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(self._capture(scope, node) for node in page.nodes), + continuation=ScanContinuation(value=_Progress(after=following).model_dump()), + final=following is None, ) - 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, - ) - - 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 - }} - }} - }} - """ + @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 - # ------------------------------------------------------------------ + async def confirm_root_absent(self, native_id: str) -> None: + """Fail closed on ambiguous absence; readable out-of-scope roots may be removed.""" + UUID(native_id) + data = await self._query("issue-membership", {"issueId": native_id}) + issue = _Issue.model_validate(data.get("issue")) + if str(issue.id) != native_id: + raise ValueError("Linear returned the wrong issue identity") + 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 +317,11 @@ async def generate_entities( files: FileService | None = None, node_selections: list[NodeSelectionData] | None = None, ) -> AsyncGenerator[BaseEntity, None]: - """Generate all entities from Linear. + """Legacy crawling is replaced by canonical page capture and offline projection.""" + raise NotImplementedError("Linear requires canonical page capture") + yield # pragma: no cover - 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)") - - 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..8702f9a525 --- /dev/null +++ b/backend/airweave/platform/sources/linear_queries/attachments.graphql @@ -0,0 +1,9 @@ +query OwnedLinearIssueAttachments($issueId: String!, $first: Int!, $after: String) { + issue(id: $issueId) { + id team { id } + attachments(first: $first, after: $after, includeArchived: true, orderBy: createdAt) { + nodes { id title url metadata createdAt updatedAt archivedAt issue { id } } + 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..0656009ffb --- /dev/null +++ b/backend/airweave/platform/sources/linear_queries/comments.graphql @@ -0,0 +1,14 @@ +query OwnedLinearIssueComments($issueId: String!, $first: Int!, $after: String) { + issue(id: $issueId) { + 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 } + } + 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..af2825ab67 --- /dev/null +++ b/backend/airweave/platform/sources/linear_queries/issue-membership.graphql @@ -0,0 +1,3 @@ +query OwnedLinearIssueMembership($issueId: String!) { + issue(id: $issueId) { id team { id } archivedAt trashed } +} 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/tests/unit/platform/sources/test_linear.py b/backend/tests/unit/platform/sources/test_linear.py index febefbf2be..6e9f6fcaff 100644 --- a/backend/tests/unit/platform/sources/test_linear.py +++ b/backend/tests/unit/platform/sources/test_linear.py @@ -1,270 +1,6 @@ -"""Unit tests for Linear source connector with continuous sync.""" +"""Offline entity timestamp parsing retained after canonical source replacement.""" -import asyncio -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from airweave.platform.configs.config import LinearConfig -from airweave.platform.entities.linear import ( - LinearCommentEntity, - LinearIssueEntity, - LinearProjectEntity, - LinearTeamEntity, - LinearUserEntity, - _parse_dt, -) -from airweave.platform.sources.linear import LinearSource - - -def _mock_auth(token="test-token"): - auth = AsyncMock() - auth.get_token = AsyncMock(return_value=token) - auth.force_refresh = AsyncMock(return_value="refreshed-token") - auth.supports_refresh = True - auth.provider_kind = "oauth" - return auth - - -def _mock_http_client(): - client = AsyncMock() - client.get = AsyncMock() - client.post = AsyncMock() - return client - - -def _mock_logger(): - return MagicMock() - - -def _graphql_response(collection_key: str, nodes: list, has_next_page: bool = False): - """Build a mock Linear GraphQL response.""" - return { - "data": { - collection_key: { - "nodes": nodes, - "pageInfo": {"hasNextPage": has_next_page, "endCursor": "cursor-1"}, - } - } - } - - -TEAM_NODE = { - "id": "team-1", - "name": "Engineering", - "key": "ENG", - "description": "The engineering team", - "color": "#00f", - "icon": None, - "private": False, - "timezone": "UTC", - "createdAt": "2024-01-01T00:00:00Z", - "updatedAt": "2024-06-01T00:00:00Z", - "parent": None, - "issueCount": 42, -} - -PROJECT_NODE = { - "id": "proj-1", - "name": "Q1 Launch", - "slugId": "q1-launch", - "description": "Q1 launch project", - "priority": 1, - "startDate": "2024-01-01", - "targetDate": "2024-03-31", - "state": "started", - "createdAt": "2024-01-01T00:00:00Z", - "updatedAt": "2024-06-01T00:00:00Z", - "completedAt": None, - "startedAt": "2024-01-01T00:00:00Z", - "progress": 0.5, - "teams": {"nodes": [{"id": "team-1", "name": "Engineering"}]}, - "lead": {"name": "Alice"}, -} - -USER_NODE = { - "id": "user-1", - "name": "Alice", - "displayName": "alice", - "email": "alice@example.com", - "avatarUrl": None, - "description": None, - "timezone": "UTC", - "active": True, - "admin": False, - "guest": False, - "lastSeen": "2024-06-01T00:00:00Z", - "statusEmoji": None, - "statusLabel": None, - "statusUntilAt": None, - "createdIssueCount": 10, - "createdAt": "2024-01-01T00:00:00Z", - "updatedAt": "2024-06-01T00:00:00Z", - "teams": {"nodes": [{"id": "team-1", "name": "Engineering", "key": "ENG"}]}, -} - -ISSUE_NODE = { - "id": "issue-1", - "identifier": "ENG-123", - "title": "Fix the bug", - "description": "There is a bug that needs fixing.", - "priority": 2, - "completedAt": None, - "createdAt": "2024-03-01T00:00:00Z", - "updatedAt": "2024-06-15T00:00:00Z", - "dueDate": "2024-07-01", - "archivedAt": None, - "state": {"name": "In Progress"}, - "team": {"id": "team-1", "name": "Engineering"}, - "project": {"id": "proj-1", "name": "Q1 Launch"}, - "assignee": {"name": "Alice"}, - "comments": { - "nodes": [ - { - "id": "comment-1", - "body": "Working on this now.", - "createdAt": "2024-03-02T00:00:00Z", - "updatedAt": "2024-03-02T00:00:00Z", - "user": {"id": "user-1", "name": "Alice"}, - } - ] - }, -} - - -@pytest.fixture -def linear_source(): - """Create a Linear source instance via the v2 ``create()`` contract.""" - - async def _create(): - return await LinearSource.create( - auth=_mock_auth(), - logger=_mock_logger(), - http_client=_mock_http_client(), - config=LinearConfig(), - ) - - return asyncio.run(_create()) - - -# --------------------------------------------------------------------------- -# create -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_create_returns_instance(): - """create() should return a LinearSource with injected deps.""" - auth = _mock_auth() - http = _mock_http_client() - source = await LinearSource.create( - auth=auth, - logger=_mock_logger(), - http_client=http, - config=LinearConfig(), - ) - assert isinstance(source, LinearSource) - assert source.auth is auth - assert source.http_client is http - - -# --------------------------------------------------------------------------- -# full sync (no cursor) -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_full_sync_generates_all_entity_types(linear_source): - """On first sync (no cursor), generate_entities yields teams, projects, users, issues, comments.""" - responses = [ - _graphql_response("teams", [TEAM_NODE]), - _graphql_response("projects", [PROJECT_NODE]), - _graphql_response("users", [USER_NODE]), - _graphql_response("issues", [ISSUE_NODE]), - ] - - with patch.object(linear_source, "_post", new_callable=AsyncMock, side_effect=responses): - entities = [] - async for e in linear_source.generate_entities(): - entities.append(e) - - teams = [e for e in entities if isinstance(e, LinearTeamEntity)] - projects = [e for e in entities if isinstance(e, LinearProjectEntity)] - users = [e for e in entities if isinstance(e, LinearUserEntity)] - issues = [e for e in entities if isinstance(e, LinearIssueEntity)] - comments = [e for e in entities if isinstance(e, LinearCommentEntity)] - - assert len(teams) == 1 - assert teams[0].team_id == "team-1" - assert len(projects) == 1 - assert projects[0].project_id == "proj-1" - assert len(users) == 1 - assert users[0].user_id == "user-1" - assert len(issues) == 1 - assert issues[0].issue_id == "issue-1" - assert issues[0].identifier == "ENG-123" - assert len(comments) == 1 - assert comments[0].comment_id == "comment-1" - - -# --------------------------------------------------------------------------- -# incremental sync (with cursor) -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_incremental_sync_passes_updated_filter(linear_source): - """Incremental sync should include updatedAt filter in GraphQL queries.""" - mock_cursor = MagicMock() - mock_cursor.data = {"last_synced_at": "2024-06-01T00:00:00Z"} - - captured_queries = [] - responses = [ - _graphql_response("teams", []), - _graphql_response("projects", []), - _graphql_response("users", []), - _graphql_response("issues", []), - ] - - async def capture_post(query): - captured_queries.append(query) - return responses[len(captured_queries) - 1] - - with patch.object(linear_source, "_post", side_effect=capture_post): - async for _ in linear_source.generate_entities(cursor=mock_cursor): - pass - - for query in captured_queries[:4]: - assert "updatedAt" in query - assert "2024-06-01T00:00:00Z" in query - - -@pytest.mark.asyncio -async def test_incremental_sync_updates_cursor(linear_source): - """Incremental sync should update the cursor with the sync start timestamp.""" - mock_cursor = MagicMock() - mock_cursor.data = {"last_synced_at": "2024-06-01T00:00:00Z"} - - responses = [ - _graphql_response("teams", []), - _graphql_response("projects", []), - _graphql_response("users", []), - _graphql_response("issues", []), - ] - - with patch.object(linear_source, "_post", new_callable=AsyncMock, side_effect=responses): - async for _ in linear_source.generate_entities(cursor=mock_cursor): - pass - - mock_cursor.update.assert_called_once() - call_kwargs = mock_cursor.update.call_args[1] - assert "last_synced_at" in call_kwargs - assert call_kwargs["last_synced_at"].endswith("+00:00") - - -# --------------------------------------------------------------------------- -# parse datetime (entity helpers) -# --------------------------------------------------------------------------- +from airweave.platform.entities.linear import _parse_dt def test_parse_dt_valid(): diff --git a/backend/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py new file mode 100644 index 0000000000..ccb657774e --- /dev/null +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -0,0 +1,307 @@ +"""Credential-free tests of exact Linear originals and whole-scope page authority.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import httpx +import pytest +from pydantic import ValidationError + +from airweave.domains.entities.canonical.page_source import ScopeAccessLost +from airweave.domains.entities.canonical.requests import CompletedScope +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 ManagedAuthProvider +from airweave.platform.configs.config import LinearConfig +from airweave.platform.sources.linear import LinearSource + +WORKSPACE = UUID(int=1) +TEAM = UUID(int=2) +ISSUE = UUID(int=3) +COMMENT = UUID(int=4) +STAMP = "2026-09-30T00:00:00Z" + + +def connection(nodes, cursor=None): + return {"nodes": nodes, "pageInfo": {"hasNextPage": cursor is not None, "endCursor": cursor}} + + +def node(identity=ISSUE): + return { + "id": str(identity), + "team": {"id": str(TEAM)}, + "createdAt": STAMP, + "updatedAt": STAMP, + "description": "Original **Markdown**", + "unknown": {"kept": True}, + } + + +def response(data, *, errors=None, status=200, headers=None): + return httpx.Response( + status, + json={"data": data, **({"errors": errors} if errors else {})}, + headers=headers, + request=httpx.Request("POST", "https://api.linear.app/graphql"), + ) + + +async def source(*pages, workspace=WORKSPACE, config=None): + client = AsyncMock() + client.post.side_effect = [ + response({"organization": {"id": str(workspace)}, "viewer": {"id": str(UUID(int=9))}}), + response({"teams": connection([{"id": str(TEAM)}])}), + *pages, + ] + auth = ManagedAuthProvider( + api_key="test-only", + connected_account_id="ca_test", + allowed_hosts=frozenset({"api.linear.app"}), + ) + result = await LinearSource.create( + auth=auth, + logger=MagicMock(), + http_client=client, + config=config or LinearConfig(workspace_id=WORKSPACE, team_ids=(TEAM,)), + ) + return result, client + + +@pytest.mark.asyncio +async def test_native_payload_pagination_and_managed_auth(): + original = node() + capture, client = await source( + response({"issues": connection([original], "next")}), response({"issues": connection([])}) + ) + first = await capture.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + assert not first.final + assert first.records[0].payload == original + assert first.records[0].identity.native_id == str(ISSUE) + second = await capture.capture_page(CompletedScope(record_type="issue"), first.continuation) + assert second.final + assert client.post.call_args.kwargs["json"]["variables"]["after"] == "next" + assert client.post.call_args.kwargs["headers"] == {} + + +@pytest.mark.asyncio +async def test_children_preserve_reply_and_independent_edit_without_parent_timestamp(): + original = { + **node(COMMENT), + "issue": {"id": str(ISSUE)}, + "parent": {"id": str(UUID(int=8))}, + "body": "edited", + "updatedAt": "2026-09-30T01:00:00Z", + } + capture, _ = await source( + response( + { + "issue": { + "id": str(ISSUE), + "team": {"id": str(TEAM)}, + "comments": connection([original]), + } + } + ) + ) + page = await capture.capture_page( + CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + ) + assert page.final and page.records[0].payload == original + assert page.records[0].identity.container_id == str(ISSUE) + assert page.records[0].source_updated_at.hour == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field,value", [("issue", {"id": str(UUID(int=99))}), ("issue", None)]) +async def test_wrong_or_null_comment_parent_fails(field, value): + original = {**node(COMMENT), field: value} + capture, _ = await source( + response( + { + "issue": { + "id": str(ISSUE), + "team": {"id": str(TEAM)}, + "comments": connection([original]), + } + } + ) + ) + with pytest.raises(ValueError): + await capture.capture_page( + CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + ) + + +@pytest.mark.asyncio +async def test_partial_graphql_response_never_returns_a_page(): + capture, _ = await source( + response( + {"issues": connection([node()])}, + errors=[{"extensions": {"code": "INTERNAL_SERVER_ERROR"}}], + ) + ) + with pytest.raises(SourceError): + await capture.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + + +@pytest.mark.asyncio +async def test_moved_child_scope_reports_scope_removed(): + capture, _ = await source( + response( + { + "issue": { + "id": str(ISSUE), + "team": {"id": str(UUID(int=99))}, + "comments": connection([]), + } + } + ) + ) + with pytest.raises(ScopeAccessLost) as error: + await capture.capture_page( + CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + ) + assert error.value.removal_reason == "scope_removed" + + +@pytest.mark.asyncio +async def test_missing_issue_is_not_access_revoked_or_empty_success(): + capture, _ = await source(response({"issue": None})) + with pytest.raises(ValidationError): + await capture.capture_page( + CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + ) + + +@pytest.mark.asyncio +async def test_omitted_readable_selected_root_blocks_reconciliation(): + capture, _ = await source( + response({"issue": node()}), + response({"issue": {**node(), "team": {"id": str(UUID(int=99))}}}), + ) + with pytest.raises(ValueError, match="still readable"): + await capture.confirm_root_absent(str(ISSUE)) + await capture.confirm_root_absent(str(ISSUE)) + + +@pytest.mark.asyncio +async def test_repeated_cursor_fails_without_final_page(): + capture, _ = await source(response({"issues": connection([node()], "same")})) + with pytest.raises(ValueError, match="non-advancing"): + await capture.capture_page( + CompletedScope(record_type="issue"), ScanContinuation(value={"after": "same"}) + ) + + +@pytest.mark.asyncio +async def test_workspace_mismatch_and_unavailable_team_fail_before_capture(): + with pytest.raises(ValueError, match="workspace"): + await source(workspace=UUID(int=99)) + with pytest.raises(ValueError, match="teams are unavailable"): + await source(config=LinearConfig(workspace_id=WORKSPACE, team_ids=(UUID(int=99),))) + + +def test_old_configuration_and_duplicates_fail_explicitly(): + with pytest.raises(ValidationError): + LinearConfig() + with pytest.raises(ValidationError): + LinearConfig(workspace_id=WORKSPACE, team_ids=(TEAM, TEAM)) + + +@pytest.mark.asyncio +async def test_config_fingerprint_and_binary_coverage(): + capture, _ = await source() + original = {**node(), "description": "![image](https://uploads.linear.app/file)"} + record = capture._capture(CompletedScope(record_type="issue"), original) + assert record.completeness == "partial" and not record.blobs + assert record.payload == original + assert capture.canonical_container_parents == {"comment": "issue", "attachment": "issue"} + assert len(capture.capture_cycle_configuration.fingerprint) == 64 + + +@pytest.mark.asyncio +async def test_rate_limit_uses_provider_reset_without_accepting_partial_data(): + capture, _ = await source( + response( + None, + errors=[{"extensions": {"code": "RATELIMITED"}}], + status=400, + headers={"X-RateLimit-Requests-Reset": "9999999999999"}, + ) + ) + # Exercise one attempt without sleeping or retrying an invented rate-limit response. + with pytest.raises(SourceRateLimitError) as error: + await LinearSource._query.__wrapped__(capture, "issues", {}) + assert error.value.retry_after > 60 + + +def test_registry_detects_canonical_class_before_instance_creation(): + from airweave.domains.entities.canonical.page_source import CanonicalPageSource + + assert isinstance(LinearSource, CanonicalPageSource) + + +@pytest.mark.asyncio +async def test_managed_linear_binding_only_allows_graphql_origin(monkeypatch): + from airweave.domains.auth_provider.exceptions import AuthProviderConfigError + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "test-only"}, + config={"account_id": "ca_test", "auth_config_id": "ac_test"}, + ) + fetch = AsyncMock( + return_value={ + "toolkit": {"slug": "linear"}, + "status": "ACTIVE", + "auth_config": {"id": "ac_test"}, + } + ) + monkeypatch.setattr(provider, "_get_with_auth", fetch) + result = await provider.get_auth_result("linear", ["access_token"]) + assert result.credentials is None + assert result.managed_auth.allowed_hosts == {"api.linear.app"} + fetch.return_value["toolkit"]["slug"] = "gmail" + with pytest.raises(AuthProviderConfigError, match="toolkit"): + await provider.get_auth_result("linear", ["access_token"]) + + +@pytest.mark.asyncio +async def test_fresh_instance_resumes_exact_saved_cursor_with_same_scope_fingerprint(): + first_source, _ = await source(response({"issues": connection([node()], "saved-page")})) + first = await first_source.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + fresh_source, client = await source(response({"issues": connection([])})) + resumed = await fresh_source.capture_page( + CompletedScope(record_type="issue"), + ScanContinuation.model_validate_json(first.continuation.model_dump_json()), + ) + assert resumed.final + assert client.post.call_args.kwargs["json"]["variables"]["after"] == "saved-page" + assert fresh_source.capture_cycle_configuration == first_source.capture_cycle_configuration + + +@pytest.mark.asyncio +async def test_attachment_page_is_independent_and_does_not_claim_file_bytes(): + attachment = { + **node(UUID(int=5)), + "issue": {"id": str(ISSUE)}, + "url": "https://external.example/document", + "metadata": {"native": True}, + } + capture, _ = await source( + response( + { + "issue": { + "id": str(ISSUE), + "team": {"id": str(TEAM)}, + "attachments": connection([attachment]), + } + } + ) + ) + page = await capture.capture_page( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), ScanContinuation() + ) + assert page.records[0].payload == attachment + assert page.records[0].completeness == "partial" + assert page.records[0].blobs == () From 1b46aac2dcb7478f56cc2c0c40848ccbb195784f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:38:23 -0700 Subject: [PATCH 044/266] test: verify bounded cross-process Slack capture resume --- .../canonical/tests/test_resume_probe.py | 221 ++++++++++++++++ backend/tests/live/README.md | 65 ++++- backend/tests/live/capture_resume.py | 146 ++++++++++ backend/tests/live/provider_lifecycle.py | 250 +++++++++++++++--- 4 files changed, 626 insertions(+), 56 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_resume_probe.py create mode 100644 backend/tests/live/capture_resume.py 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..8b6e6f98c4 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py @@ -0,0 +1,221 @@ +"""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 = 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" diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 2928bcd78d..63696313fa 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -76,17 +76,19 @@ 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. -## Slack and Wispr durable samples - -`LIVE_PROVIDERS=slack,wispr` selects the existing source adapters. Set -`LIVE_SLACK_ACCOUNT_ID`, `LIVE_WISPR_ACCOUNT_ID` and `LIVE_WISPR_USER_ID` for -previously connected, authorized accounts. Slack verifies `auth.test` identity -and the matching `users.info` email against `LIVE_EXPECTED_EMAIL`. Wispr verifies +## Historical Slack sample and current Wispr sampler + +The earlier Slack sample below used the old observation adapter. Slack now uses +durable pages and the sampler explicitly directs it to `provider_lifecycle.py`; +there is no second capture implementation. `LIVE_PROVIDERS=wispr` remains supported, +with `LIVE_WISPR_ACCOUNT_ID` and `LIVE_WISPR_USER_ID` for a previously connected +account. Slack's lifecycle verifies `auth.test` and the matching `users.info` +email against `LIVE_EXPECTED_EMAIL`. Wispr 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. A Wispr-only run does not require the email variable. -Slack/Wispr stop after three distinct records containing a message/meeting, or +The historical Slack and current Wispr sampler stop after three distinct records, or 25 observations. No file/blob presence is required. Repeated native identities are reduced to their latest raw observation in this bounded sample before capture and replay; Slack history/replies can return the same parent twice. @@ -217,9 +219,10 @@ The earlier smaller-budget failure evidence remains unchanged. Sanitized result: ### Slack and Wispr full lifecycle attempts -Both modes use the real canonical orchestrator and production `cursor=None`. -Slack reconciles only completed channel/message scopes; Wispr has no deletion -scope or durable provider checkpoint. Neither is described as delta resume. +These historical attempts used the real canonical orchestrator and production +`cursor=None`. Slack's current durable-page trial is described below. Wispr still +has no deletion scope or durable provider checkpoint; its enumeration is not a +delta-resume guarantee. Actual network requests are metered, including Wispr session/account operations. The first Slack full-accessible-history attempt stopped on a rate limit after @@ -236,8 +239,9 @@ became cancelled; two scopes started, none completed, and the checkpoint stayed unchanged. Cancellation telemetry then accessed missing fixture connection metadata, masking the timeout with AttributeError. The fixture is corrected; this was not rerun. No second process started, and private schema/files were -removed. Full Slack coverage remains unverified. A durable per-scope sweep and -continuation design is needed before repeatedly rescanning this history. +removed. Full Slack coverage remains unverified. The durable per-scope recovery +implementation below addresses this retry limitation; it has only synthetic +HTTP/PostgreSQL and local cross-process verification so far. See `evidence/slack-lifecycle-backoff-incomplete-20260930.json`. Wispr enumeration and one diagnostic retry each failed after 18 requests and nine @@ -250,3 +254,40 @@ attempts cleaned private storage. This is additional failure evidence alongside the earlier successful three-record sample, not a replacement for that sample. See `evidence/slack-lifecycle-incomplete-20260930.json` and `evidence/wispr-lifecycle-incomplete-20260930.json`. + + +### Current Slack cross-process recovery trial (not yet run live) + +`LIVE_LIFECYCLE_PROVIDER=slack` uses the production page adapter, capture pipeline, +SQL state, and job state machine. It has three stages within **one aggregate cap**: +600 provider requests, 10,000 observed records, and 1,800 seconds including all +Retry-After waits and subprocess startup. + +1. Start a new job. After a history page commits pending thread IDs, the test-only + wrapper verifies the exact SQL continuation, active cycle, current attempt and + running job. It writes private UUID/version evidence, flushes sanitized counters, + and exits 75 before another provider request. This simulates process loss without + marking the job failed or cancelled. +2. Start a fresh process using the same job and attempt 2. The production driver + refreshes membership, then resumes the saved child state. A rejected cursor is + reported separately from successful cursor reuse; the normal one-restart bound + still applies. No complete-scope claim is made until actual finalization. +3. Only after completion, start a genuinely new job and verify its new cycle. + +The parent subtracts each child's counters from the remaining limits and enforces +one wall-clock deadline even while a child sleeps. No pending thread found, changed +scope, or exhausted budget leaves the requested resume proof incomplete. A forcibly +killed child's counters are unknown: aggregate exact counts are null and known +counts are reported only as lower bounds. No further stage follows that failure. +After the resumed child exits (including a forced timeout), the parent reads the +saved scan directly from SQL. Same tenant, sync, job, attempt 2, cycle and sweep, +with an advanced revision, proves a resumed page committed. This is reported as +`resumed_page_committed`; `resumed_partial` means that proof exists but the cycle +has not completed. Scope withdrawal and expired-cursor restarts cannot satisfy +this proof. Full-cycle and subsequent-new-cycle completion remain separate checks. +Private schema, downloaded files, and resume evidence are removed by the parent. + +`test_resume_probe.py` verifies the actual special process exit and fresh-process +resume using synthetic Slack responses and isolated PostgreSQL. Its budget tests +use fake child processes to check decreasing limits and unknown killed-child usage. +These are harness tests, not evidence of live Slack coverage or throughput. diff --git a/backend/tests/live/capture_resume.py b/backend/tests/live/capture_resume.py new file mode 100644 index 0000000000..3ef259c186 --- /dev/null +++ b/backend/tests/live/capture_resume.py @@ -0,0 +1,146 @@ +"""Test-only process-loss injection after a verified production page commit.""" + +import json +import os +from pathlib import Path +from typing import Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import CaptureCycle +from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.entities.canonical.scan_models import ScanContinuation +from airweave.models.capture_scan import CaptureScan +from airweave.models.sync import Sync +from airweave.models.sync_cursor import SyncCursor +from airweave.models.sync_job import SyncJob + + +class ResumeTarget(BaseModel): + """Private test evidence references SQL state; no duplicated cursor or record payload.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + scan_id: UUID + cycle_id: UUID + sweep_id: UUID + revision: int + + +class PageResumeProbe: + """The child exits before its next provider call, leaving the actual job RUNNING.""" + + def __init__( + self, + sessions, + organization_id, + sync_id, + job_id, + attempt_id, + target_path: Path, + mode: Literal["interrupt", "resume"], + counters, + ): + self.sessions = sessions + self.organization_id, self.sync_id = organization_id, sync_id + self.job_id, self.attempt_id = job_id, attempt_id + self.target_path, self.mode, self.counters = target_path, mode, counters + self.resumed_saved_page = False + self.saved_cursor_expired = False + self._checking_saved_page = False + self.target = ( + ResumeTarget.model_validate_json(target_path.read_text()) if mode == "resume" else None + ) + + async def before_page(self, scope: CompletedScope, continuation: ScanContinuation): + """Verify committed pending replies and writer/cycle ownership before fault injection.""" + self._checking_saved_page = False + if scope.record_type != "message" or not continuation.value.get("pending_threads"): + return + async with self.sessions() as db: + row = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == self.organization_id, + CaptureScan.sync_id == self.sync_id, + CaptureScan.record_type == scope.record_type, + CaptureScan.container_id == scope.container_id, + ) + ) + sync = await db.get(Sync, self.sync_id) + job = await db.get(SyncJob, self.job_id) + cursor = await db.scalar( + select(SyncCursor).where( + SyncCursor.organization_id == self.organization_id, + SyncCursor.sync_id == self.sync_id, + ) + ) + if row is None or sync is None or job is None or cursor is None: + raise AssertionError("Durable interruption state is missing") + cycle = CaptureCycle.model_validate(cursor.cursor_data["canonical_cycle"]) + assert sync.organization_id == self.organization_id + assert sync.writer_job_id == self.job_id and sync.writer_attempt_id == self.attempt_id + assert job.organization_id == self.organization_id and job.sync_id == self.sync_id + assert job.status == "running" and cycle.phase == "active" + assert cycle.root_writer_attempt_id == self.attempt_id + assert row.cycle_id == cycle.version.cycle_id and row.phase == "collecting" + assert row.continuation == continuation.value + target = ResumeTarget( + scan_id=row.id, cycle_id=row.cycle_id, sweep_id=row.sweep_id, revision=row.revision + ) + if self.mode == "interrupt": + self.target_path.write_text(target.model_dump_json()) + os.chmod(self.target_path, 0o600) + print( + json.dumps( + { + "failed": False, + "counters_complete": True, + "intentional_interruption": True, + "durable_pending_threads_verified": True, + "job_status": "running", + "full_scope_completed": False, + "checkpoint_saved": False, + **self.counters, + } + ), + flush=True, + ) + os._exit(75) # Test child only: simulate process loss without cancellation transitions. + if target == self.target: + self.resumed_saved_page = True + self._checking_saved_page = True + + def cursor_expired(self): + """A resumed request that is rejected is distinct from successful cursor reuse.""" + if self._checking_saved_page: + self.saved_cursor_expired = True + + +async def inspect_resume_progress(sessions, organization_id, sync_id, job_id, target_path: Path): + """Parent-side readback survives child death without copying provider content or cursors.""" + target = ResumeTarget.model_validate_json(target_path.read_text()) + async with sessions() as db: + row = await db.scalar( + select(CaptureScan).where( + CaptureScan.id == target.scan_id, + CaptureScan.organization_id == organization_id, + CaptureScan.sync_id == sync_id, + ) + ) + sync = await db.get(Sync, sync_id) + job = await db.get(SyncJob, job_id) + advanced = bool( + row is not None + and sync is not None + and job is not None + and sync.organization_id == organization_id + and job.organization_id == organization_id + and job.sync_id == sync_id + and sync.writer_job_id == job_id + and sync.writer_attempt_number == 2 + and row.cycle_id == target.cycle_id + and row.sweep_id == target.sweep_id + and row.revision > target.revision + ) + return {"resumed_page_committed": advanced} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 1a652e8cca..e10c8faced 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -27,7 +27,10 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from airweave.adapters.event_bus.fake import FakeEventBus -from airweave.domains.entities.canonical.page_source import CanonicalPageSource +from airweave.domains.entities.canonical.page_source import ( + CanonicalPageSource, + InvalidScanContinuation, +) from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope from airweave.domains.entities.canonical.service import CanonicalCaptureService from airweave.domains.entities.canonical.store import CanonicalRecordStore @@ -76,14 +79,22 @@ async def write_file(self, path, content): class BoundedPageSource: """Test budget around the real page adapter; production driver still owns all progress.""" - def __init__(self, source, counters, record_limit): + def __init__(self, source, counters, record_limit, resume_probe=None): self.source, self.counters, self.record_limit = source, counters, record_limit + self.resume_probe = resume_probe self.canonical_record_types = source.canonical_record_types self.canonical_container_parents = source.canonical_container_parents self.capture_cycle_configuration = source.capture_cycle_configuration async def capture_page(self, scope, continuation): - page = await self.source.capture_page(scope, continuation) + if self.resume_probe is not None: + await self.resume_probe.before_page(scope, continuation) + try: + page = await self.source.capture_page(scope, continuation) + except InvalidScanContinuation: + if self.resume_probe is not None: + self.resume_probe.cursor_expired() + raise self.counters["records_observed"] += len(page.records) if self.counters["records_observed"] > self.record_limit: raise BudgetExceeded("records") @@ -185,6 +196,42 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert saved != previous +async def prepare_job(db, organization_id, sync_id, job_id, attempt_number): + """Only an explicit retry may reuse a still-running job in the private schema.""" + if attempt_number == 1: + job = SyncJob(id=job_id, organization_id=organization_id, sync_id=sync_id, status="pending") + db.add(job) + await db.commit() + return job + job = await db.get(SyncJob, job_id) + if ( + job is None + or job.sync_id != sync_id + or job.organization_id != organization_id + or job.status != "running" + ): + raise ValueError("Resume requires the same interrupted running job") + return job + + +def prepare_resume_probe(manifest, sessions, organization_id, sync_id, job_id, attempt, counters): + """Fault injection is opt-in for the two explicit cross-process trial stages.""" + if manifest.get("resume_stage") not in {"interrupt", "resume"}: + return None + from capture_resume import PageResumeProbe + + return PageResumeProbe( + sessions, + organization_id, + sync_id, + job_id, + attempt.id, + Path(manifest["root"]) / "resume-target.json", + manifest["resume_stage"], + counters, + ) + + async def child(manifest): engine = create_async_engine( harness.test_database_url(), @@ -192,7 +239,8 @@ async def child(manifest): ) sessions = async_sessionmaker(engine, expire_on_commit=False) organization_id, sync_id = UUID(manifest["organization_id"]), UUID(manifest["sync_id"]) - job_id = uuid4() + job_id = UUID(manifest["job_id"]) if "job_id" in manifest else uuid4() + attempt_number = manifest.get("attempt_number", 1) counters = {"provider_requests": 0, "records_observed": 0, "started": 0, "completed": 0} result = {"failed": True} previous = {} @@ -202,6 +250,7 @@ async def child(manifest): counters["resumed_changes_requests"] = 0 last_operation = "identity" rate_limits = [] + resume_probe = None async def observe_slack_rate_limit(response): await response.aread() @@ -248,11 +297,8 @@ async def db_context(): async with sessions() as db: organization = await db.get(Organization, organization_id) sync = await db.get(Sync, sync_id) - job = SyncJob( - id=job_id, organization_id=organization_id, sync_id=sync_id, status="pending" - ) - db.add(job) - await db.commit() + job = await prepare_job(db, organization_id, sync_id, job_id, attempt_number) + config = SyncConfig() config.behavior.skip_guardrails = True logger = logging.getLogger("private-lifecycle") @@ -321,9 +367,12 @@ async def db_context(): ): assert source.cursor_class == cursor_schema bus = FakeEventBus() - attempt = CaptureAttempt(id=uuid4(), number=1) + attempt = CaptureAttempt(id=uuid4(), number=attempt_number) + resume_probe = prepare_resume_probe( + manifest, sessions, organization_id, sync_id, job_id, attempt, counters + ) page_source = ( - BoundedPageSource(source, counters, manifest["record_limit"]) + BoundedPageSource(source, counters, manifest["record_limit"], resume_probe) if isinstance(source, CanonicalPageSource) else None ) @@ -417,6 +466,10 @@ async def db_context(): ) result = { "failed": False, + "full_scope_completed": name != "wispr", + "counters_complete": True, + "resumed_saved_page": resume_probe.resumed_saved_page if resume_probe else False, + "saved_cursor_expired": resume_probe.saved_cursor_expired if resume_probe else False, "provider": name, "identity_verification": identity_verification, "almanac_consumer": consumer_result, @@ -451,6 +504,10 @@ async def db_context(): status = await db.scalar(select(SyncJob.status).where(SyncJob.id == job_id)) result = { "failed": True, + "full_scope_completed": False, + "counters_complete": True, + "resumed_saved_page": resume_probe.resumed_saved_page if resume_probe else False, + "saved_cursor_expired": resume_probe.saved_cursor_expired if resume_probe else False, "error_type": type(error).__name__, "safe_reason": safe_failure_reason(error), "last_operation": last_operation, @@ -473,6 +530,108 @@ def check_free_disk(provider): return available +async def execute_trial(path, timeout): + """Return only sanitized child evidence; forced termination has unknown usage.""" + process = await asyncio.create_subprocess_exec( + sys.executable, + str(Path(__file__).absolute()), + str(path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, _ = await asyncio.wait_for(process.communicate(), timeout=timeout) + except asyncio.TimeoutError: + return -9, { + "failed": True, + "safe_reason": "aggregate_deadline", + "full_scope_completed": False, + "counters_complete": False, + } + finally: + if process.returncode is None: + process.kill() + await process.wait() + lines = [json.loads(line) for line in stdout.decode().splitlines() if line.startswith("{")] + if not lines: + return process.returncode, { + "failed": True, + "safe_reason": "child_no_result", + "full_scope_completed": False, + "counters_complete": False, + } + return process.returncode, lines[-1] + + +async def run_trials(manifest, evidence_reader=None): + """One parent-owned aggregate budget across interruption, retry, and new-job proof.""" + path = Path(manifest["root"]) / "manifest.json" + results = [] + name = manifest["provider"] + resume_trial = name == "slack" + stages = ("interrupt", "resume", "new_cycle") if resume_trial else ("first", "second") + total_requests, total_records = manifest["request_limit"], manifest["record_limit"] + deadline = time.monotonic() + manifest["timeout"] if resume_trial else None + interrupted_job_id = str(uuid4()) + for stage in stages: + check_free_disk(name) + if resume_trial: + remaining = deadline - time.monotonic() + requests_left = total_requests - sum( + item.get("provider_requests", 0) for item in results + ) + records_left = total_records - sum(item.get("records_observed", 0) for item in results) + if remaining <= 0 or requests_left <= 0 or records_left <= 0: + results.append( + { + "failed": True, + "safe_reason": "aggregate_budget", + "full_scope_completed": False, + } + ) + break + manifest.update( + resume_stage=stage, + job_id=interrupted_job_id if stage != "new_cycle" else str(uuid4()), + attempt_number=2 if stage == "resume" else 1, + request_limit=requests_left, + record_limit=records_left, + timeout=remaining, + ) + path.write_text(json.dumps(manifest)) + code, result = await execute_trial( + path, + max(0.001, deadline - time.monotonic()) if resume_trial else manifest["timeout"] + 60, + ) + if stage == "resume" and evidence_reader is not None: + result.update(await evidence_reader()) + result["resumed_partial"] = bool( + result.get("resumed_page_committed") and not result.get("full_scope_completed") + ) + results.append(result) + print(json.dumps({"run": len(results), "stage": stage, **result}), flush=True) + if stage == "interrupt": + if code != 75 or not result.get("intentional_interruption"): + break + elif code: + break + return results + + +def trial_counter_summary(results): + """A killed child has unknown usage; known values are only a lower bound.""" + complete = all(item.get("counters_complete", True) for item in results) + requests = sum(item.get("provider_requests", 0) for item in results) + records = sum(item.get("records_observed", 0) for item in results) + return { + "aggregate_counters_complete": complete, + "aggregate_provider_requests": requests if complete else None, + "aggregate_records_observed": records if complete else None, + "known_provider_requests": requests, + "known_records_observed": records, + } + + async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") if name not in {"gmail", "google_calendar", "google_drive", "slack", "wispr"}: @@ -557,36 +716,20 @@ async def main(): "end": (now + timedelta(days=7)).isoformat(), }, } - path = root / "manifest.json" - path.write_text(json.dumps(manifest)) - for _ in range(2): - check_free_disk(name) - process = await asyncio.create_subprocess_exec( - sys.executable, - str(Path(__file__).absolute()), - str(path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, + + async def resume_evidence(): + from capture_resume import inspect_resume_progress + + return await inspect_resume_progress( + sessions, + organization_id, + sync_id, + UUID(manifest["job_id"]), + root / "resume-target.json", ) - try: - stdout, _ = await asyncio.wait_for( - process.communicate(), timeout=manifest["timeout"] + 60 - ) - finally: - if process.returncode is None: - process.kill() - await process.wait() - lines = [ - json.loads(line) - for line in stdout.decode().splitlines() - if line.startswith("{") - ] - if not lines: - raise RuntimeError("Child failed without sanitized result") - results.append(lines[-1]) - print(json.dumps({"run": len(results), **lines[-1]}), flush=True) - if process.returncode: - break + + results = await run_trials(manifest, resume_evidence if name == "slack" else None) + finally: os.umask(old_umask) await engine.dispose() @@ -594,6 +737,15 @@ async def main(): async with admin.begin() as connection: await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) await admin.dispose() + completed = [item for item in results if item.get("full_scope_completed")] + resume_verified = ( + name == "slack" + and len(results) == 3 + and results[0].get("intentional_interruption") + and results[1].get("resumed_page_committed") + and results[1].get("full_scope_completed") + and results[2].get("full_scope_completed") + ) print( json.dumps( { @@ -601,13 +753,23 @@ async def main(): "free_disk_bytes_before": free_disk_bytes, "schema_removed": True, "blob_directory_removed": True, - "identical_second_read": len(results) == 2 - and results[0].get("payload_revision_digest") - == results[1].get("payload_revision_digest"), + "resume_verified": bool(resume_verified), + "resumed_page_committed": any( + item.get("resumed_page_committed", False) for item in results + ), + "resumed_partial": any(item.get("resumed_partial", False) for item in results), + "saved_cursor_reused": bool( + resume_verified and not results[1].get("saved_cursor_expired") + ), + **trial_counter_summary(results), + "identical_second_read": len(completed) == 2 + and completed[0].get("payload_revision_digest") + == completed[1].get("payload_revision_digest"), } ) ) - return 0 if len(results) == 2 and not any(r["failed"] for r in results) else 1 + success = resume_verified if name == "slack" else len(results) == 2 + return 0 if success and not any(item["failed"] for item in results) else 1 if __name__ == "__main__": From 82464bd060d0002858d411af686776ddd6e06fb6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:41:53 -0700 Subject: [PATCH 045/266] Wire owned file storage into durable page capture --- .../domains/entities/canonical/page_source.py | 3 ++- .../canonical/tests/test_slack_membership.py | 1 + .../canonical/tests/test_slack_recovery.py | 5 +++-- .../domains/sync_pipeline/canonical_capture.py | 6 ++++++ .../domains/sync_pipeline/canonical_scan.py | 5 ++++- backend/airweave/domains/sync_pipeline/factory.py | 1 + backend/airweave/platform/sources/slack.py | 2 +- backend/tests/live/README.md | 2 +- backend/tests/live/provider_lifecycle.py | 7 ++++--- .../unit/platform/sources/test_slack_capture.py | 14 +++++++++----- 10 files changed, 32 insertions(+), 14 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index c241add9ed..052529c4a9 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -11,6 +11,7 @@ ScopeRemovalReason, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation +from airweave.domains.storage.file_service import FileService class CapturePage(BaseModel): @@ -49,7 +50,7 @@ class CanonicalPageSource(Protocol): capture_cycle_configuration: CycleConfiguration async def capture_page( - self, scope: CompletedScope, continuation: ScanContinuation + self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService ) -> CapturePage: """Fetch one page without advancing durable state or performing writes.""" ... diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py index 30c6d287d0..98e34c7af1 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py @@ -54,6 +54,7 @@ async def test_first_completed_membership_hides_prior_partial_channel( CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), SlackSource.canonical_container_parents, page_source=connector, + files=MagicMock(), ) await pipeline.start(ctx) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index 0b2b9876dc..bc747d9aaa 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -39,6 +39,7 @@ def runner(database, source, responses, attempt=1, service=None): 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 @@ -439,12 +440,12 @@ async def test_scope_withdrawal_keeps_uncertain_reason_without_claiming_access_r instance, _, pipeline = runner(database, source, [ROOT, HISTORY]) original_page = pipeline.page_source.capture_page - async def unavailable(scope, continuation): + async def unavailable(scope, continuation, *, files): 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) + return await original_page(scope, continuation, files=files) pipeline.page_source.capture_page = unavailable await run(instance) diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index cc26a74720..a433ecb0fb 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -23,6 +23,7 @@ ) 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 @@ -43,6 +44,7 @@ def __init__( attempt: CaptureAttempt, container_parents: dict[str, str] | None = None, page_source: CanonicalPageSource | None = None, + files: FileService | None = None, ): """Inject transactional persistence and existing progress-event infrastructure.""" if not record_types: @@ -60,6 +62,7 @@ def __init__( if parent in self._container_parents: raise ValueError("Only root-container visibility relationships are supported") 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 @@ -93,6 +96,8 @@ async def run_scans( 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: @@ -113,6 +118,7 @@ async def progress(result: CaptureResult, records: tuple[CaptureRecord, ...]) -> progress, self._with_parent, check_limits, + self.files, ).run() def _writer(self) -> WriterFence: diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index b249340c68..acf9f6acb4 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -31,6 +31,7 @@ ScanState, ) from airweave.domains.entities.canonical.service import CanonicalCaptureService +from airweave.domains.storage.file_service import FileService class CanonicalScanDriver: @@ -45,9 +46,11 @@ def __init__( progress: Callable[[CaptureResult, tuple[CaptureRecord, ...]], Awaitable[None]], with_parent: Callable[[CaptureRecord], CaptureRecord], check_limits: Callable[[], Awaitable[None]], + files: FileService, ): """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.progress, self.with_parent, self.check_limits = progress, with_parent, check_limits async def run(self) -> CaptureCycle: @@ -127,7 +130,7 @@ async def scan( while state.phase == "collecting": await self.check_limits() try: - page = await self.source.capture_page(scope, state.continuation) + page = await self.source.capture_page(scope, state.continuation, files=self.files) except InvalidScanContinuation: if restarts >= 1: raise diff --git a/backend/airweave/domains/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index b9757b832c..60b62391b7 100644 --- a/backend/airweave/domains/sync_pipeline/factory.py +++ b/backend/airweave/domains/sync_pipeline/factory.py @@ -248,6 +248,7 @@ async def create_orchestrator( 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 ), diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 35283e102b..c266351234 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -217,7 +217,7 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: ) async def capture_page( - self, scope: CompletedScope, continuation: ScanContinuation + self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService ) -> CapturePage: """Fetch one page; record and nested reply progress are committed by the pipeline.""" try: diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 63696313fa..949e7dd3f2 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -260,7 +260,7 @@ See `evidence/slack-lifecycle-incomplete-20260930.json` and `LIVE_LIFECYCLE_PROVIDER=slack` uses the production page adapter, capture pipeline, SQL state, and job state machine. It has three stages within **one aggregate cap**: -600 provider requests, 10,000 observed records, and 1,800 seconds including all +600 provider requests, 10,000 observed records, and 7,200 seconds including all Retry-After waits and subprocess startup. 1. Start a new job. After a history page commits pending thread IDs, the test-only diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index e10c8faced..2c2714b02a 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -86,11 +86,11 @@ def __init__(self, source, counters, record_limit, resume_probe=None): self.canonical_container_parents = source.canonical_container_parents self.capture_cycle_configuration = source.capture_cycle_configuration - async def capture_page(self, scope, continuation): + async def capture_page(self, scope, continuation, *, files: FileService): if self.resume_probe is not None: await self.resume_probe.before_page(scope, continuation) try: - page = await self.source.capture_page(scope, continuation) + page = await self.source.capture_page(scope, continuation, files=files) except InvalidScanContinuation: if self.resume_probe is not None: self.resume_probe.cursor_expired() @@ -384,6 +384,7 @@ async def db_context(): attempt, container_parents=getattr(source, "canonical_container_parents", {}), page_source=page_source, + files=files, ) runtime = SyncRuntime( source=source, @@ -694,7 +695,7 @@ async def main(): "gmail": 600, "google_calendar": 180, "google_drive": 300, - "slack": 1800, + "slack": 7200, "wispr": 300, }[name], "blob_byte_limit": (512 if name == "google_drive" else 256) * 1024 * 1024, diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 84270d8692..3114578411 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -31,15 +31,17 @@ async def test_page_commits_history_before_threads_and_preserves_native_fields() {"messages": [{"ts": "2", "files": [{"id": "F1"}]}]}, ] ) - root = await connector.capture_page(CompletedScope(record_type="channel"), ScanContinuation()) + root = await connector.capture_page( + CompletedScope(record_type="channel"), ScanContinuation(), files=MagicMock() + ) assert root.final and root.records[0].payload["unknown_native_field"] == "kept" scope = CompletedScope(record_type="message", container_id="C1") - history = await connector.capture_page(scope, ScanContinuation()) + history = await connector.capture_page(scope, ScanContinuation(), files=MagicMock()) assert history.records[0].payload["blocks"] == [{"type": "rich_text"}] assert history.continuation.value["pending_threads"] == ["1"] and not history.final - replies = await connector.capture_page(scope, history.continuation) + replies = await connector.capture_page(scope, history.continuation, files=MagicMock()) assert replies.records[0].payload["thread_ts"] == "1" and not replies.final - last = await connector.capture_page(scope, replies.continuation) + last = await connector.capture_page(scope, replies.continuation, files=MagicMock()) assert last.final and last.records[0].completeness == "partial" calls = connector._get.call_args_list assert calls[0].args[1]["types"] == "public_channel,private_channel,im,mpim" @@ -62,7 +64,9 @@ async def test_incomplete_pagination_fails(pages): with pytest.raises(ValueError): for _ in pages: result = await connector.capture_page( - CompletedScope(record_type="message", container_id="C1"), progress + CompletedScope(record_type="message", container_id="C1"), + progress, + files=MagicMock(), ) progress = result.continuation From 02deb9214d0ee39ccbd2c03abda0126aa7d70a36 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:45:15 -0700 Subject: [PATCH 046/266] Retain bounded Linear files and project captured originals offline --- .../auth_provider/providers/composio.py | 2 +- .../entities/canonical/linear_projection.py | 152 +++++++++++++++ .../entities/canonical/projection_mappers.py | 4 + .../canonical/tests/test_linear_projection.py | 147 +++++++++++++++ .../airweave/platform/entities/__init__.py | 2 + backend/airweave/platform/entities/linear.py | 23 +++ backend/airweave/platform/sources/linear.py | 72 ++++++- .../linear_queries/attachments.graphql | 2 +- .../sources/linear_queries/comments.graphql | 2 +- .../platform/sources/test_linear_capture.py | 177 ++++++++++++++++-- 10 files changed, 566 insertions(+), 17 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/linear_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_linear_projection.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index cda796285a..7d5995202a 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -116,7 +116,7 @@ async def get_auth_result( "google_calendar": frozenset({"www.googleapis.com"}), "google_drive": frozenset({"www.googleapis.com"}), "slack": frozenset({"slack.com"}), - "linear": frozenset({"api.linear.app"}), + "linear": frozenset({"api.linear.app", "uploads.linear.app"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( 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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 686c009cf4..0fefada5fb 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -225,6 +225,10 @@ async def map_record( from airweave.domains.entities.canonical.gmail_projection import map_gmail entities = await map_gmail(record, storage, directory) + elif source_name == "linear": + from airweave.domains.entities.canonical.linear_projection import map_linear + + entities = await map_linear(record, storage, directory) elif source_name == "google_drive": entities = await _drive(record, storage, directory) elif source_name == "google_calendar": 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..8ddc0417c8 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py @@ -0,0 +1,147 @@ +"""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: + 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: + 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: + 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/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index bad98fcf8b..d0269675f5 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -181,6 +181,7 @@ LinearAttachmentEntity, LinearCommentEntity, LinearIssueEntity, + LinearLinkedAttachmentEntity, LinearProjectEntity, LinearTeamEntity, LinearUserEntity, @@ -551,6 +552,7 @@ LinearAttachmentEntity, LinearCommentEntity, LinearIssueEntity, + LinearLinkedAttachmentEntity, LinearProjectEntity, LinearTeamEntity, LinearUserEntity, 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/sources/linear.py b/backend/airweave/platform/sources/linear.py index 6771e82a02..9726229e8d 100644 --- a/backend/airweave/platform/sources/linear.py +++ b/backend/airweave/platform/sources/linear.py @@ -6,12 +6,15 @@ import json import time from collections.abc import AsyncGenerator +from contextlib import aclosing from datetime import datetime, timezone from functools import partial from pathlib import Path from typing import Literal +from urllib.parse import urlsplit from uuid import UUID +import httpx from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool from tenacity import retry, retry_if_exception_type, stop_after_attempt @@ -33,6 +36,7 @@ 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 LinearConfig from airweave.platform.decorators import source @@ -50,7 +54,7 @@ _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 = 1 +FIELD_SET_VERSION = 2 class _Model(BaseModel): @@ -240,7 +244,7 @@ def _check_issue(self, payload: JsonValue, expected: str | None = None) -> _Issu return issue async def capture_page( - self, scope: CompletedScope, continuation: ScanContinuation + self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService ) -> CapturePage: """Fetch exactly one whole-scope page; engine commits it and its continuation.""" progress = _Progress.model_validate(continuation.value) @@ -272,11 +276,73 @@ async def capture_page( raise ValueError("Unsupported Linear scope") following = self._following(page, progress.after) return CapturePage( - records=tuple(self._capture(scope, node) for node in page.nodes), + 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 _retain( + self, + scope: CompletedScope, + payload: dict[str, JsonValue], + files: FileService, + ) -> 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 + ) + 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() + ) + 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", + } + ) + + @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 + ) + except ValueError: + return False + @staticmethod def _capture(scope: CompletedScope, payload: dict[str, JsonValue]) -> CaptureRecord: identity = _Identity.model_validate(payload) diff --git a/backend/airweave/platform/sources/linear_queries/attachments.graphql b/backend/airweave/platform/sources/linear_queries/attachments.graphql index 8702f9a525..4481afe0d9 100644 --- a/backend/airweave/platform/sources/linear_queries/attachments.graphql +++ b/backend/airweave/platform/sources/linear_queries/attachments.graphql @@ -2,7 +2,7 @@ query OwnedLinearIssueAttachments($issueId: String!, $first: Int!, $after: Strin issue(id: $issueId) { id team { id } attachments(first: $first, after: $after, includeArchived: true, orderBy: createdAt) { - nodes { id title url metadata createdAt updatedAt archivedAt issue { id } } + nodes { id title url metadata createdAt updatedAt archivedAt issue { id identifier title url team { id name } project { id name } } } pageInfo { hasNextPage endCursor } } } diff --git a/backend/airweave/platform/sources/linear_queries/comments.graphql b/backend/airweave/platform/sources/linear_queries/comments.graphql index 0656009ffb..fbb8ec803e 100644 --- a/backend/airweave/platform/sources/linear_queries/comments.graphql +++ b/backend/airweave/platform/sources/linear_queries/comments.graphql @@ -6,7 +6,7 @@ query OwnedLinearIssueComments($issueId: String!, $first: Int!, $after: String) id body url createdAt updatedAt archivedAt resolvedAt user { id name updatedAt } parent { id } - issue { id } + issue { id identifier title url team { id name } project { id name } } } pageInfo { hasNextPage endCursor } } diff --git a/backend/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py index ccb657774e..3941139ebb 100644 --- a/backend/tests/unit/platform/sources/test_linear_capture.py +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -73,11 +73,15 @@ async def test_native_payload_pagination_and_managed_auth(): capture, client = await source( response({"issues": connection([original], "next")}), response({"issues": connection([])}) ) - first = await capture.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + first = await capture.capture_page( + CompletedScope(record_type="issue"), ScanContinuation(), files=MagicMock() + ) assert not first.final assert first.records[0].payload == original assert first.records[0].identity.native_id == str(ISSUE) - second = await capture.capture_page(CompletedScope(record_type="issue"), first.continuation) + second = await capture.capture_page( + CompletedScope(record_type="issue"), first.continuation, files=MagicMock() + ) assert second.final assert client.post.call_args.kwargs["json"]["variables"]["after"] == "next" assert client.post.call_args.kwargs["headers"] == {} @@ -104,7 +108,9 @@ async def test_children_preserve_reply_and_independent_edit_without_parent_times ) ) page = await capture.capture_page( - CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + CompletedScope(record_type="comment", container_id=str(ISSUE)), + ScanContinuation(), + files=MagicMock(), ) assert page.final and page.records[0].payload == original assert page.records[0].identity.container_id == str(ISSUE) @@ -128,7 +134,9 @@ async def test_wrong_or_null_comment_parent_fails(field, value): ) with pytest.raises(ValueError): await capture.capture_page( - CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + CompletedScope(record_type="comment", container_id=str(ISSUE)), + ScanContinuation(), + files=MagicMock(), ) @@ -141,7 +149,9 @@ async def test_partial_graphql_response_never_returns_a_page(): ) ) with pytest.raises(SourceError): - await capture.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + await capture.capture_page( + CompletedScope(record_type="issue"), ScanContinuation(), files=MagicMock() + ) @pytest.mark.asyncio @@ -159,7 +169,9 @@ async def test_moved_child_scope_reports_scope_removed(): ) with pytest.raises(ScopeAccessLost) as error: await capture.capture_page( - CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + CompletedScope(record_type="comment", container_id=str(ISSUE)), + ScanContinuation(), + files=MagicMock(), ) assert error.value.removal_reason == "scope_removed" @@ -169,7 +181,9 @@ async def test_missing_issue_is_not_access_revoked_or_empty_success(): capture, _ = await source(response({"issue": None})) with pytest.raises(ValidationError): await capture.capture_page( - CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation() + CompletedScope(record_type="comment", container_id=str(ISSUE)), + ScanContinuation(), + files=MagicMock(), ) @@ -189,7 +203,9 @@ async def test_repeated_cursor_fails_without_final_page(): capture, _ = await source(response({"issues": connection([node()], "same")})) with pytest.raises(ValueError, match="non-advancing"): await capture.capture_page( - CompletedScope(record_type="issue"), ScanContinuation(value={"after": "same"}) + CompletedScope(record_type="issue"), + ScanContinuation(value={"after": "same"}), + files=MagicMock(), ) @@ -260,7 +276,7 @@ async def test_managed_linear_binding_only_allows_graphql_origin(monkeypatch): monkeypatch.setattr(provider, "_get_with_auth", fetch) result = await provider.get_auth_result("linear", ["access_token"]) assert result.credentials is None - assert result.managed_auth.allowed_hosts == {"api.linear.app"} + assert result.managed_auth.allowed_hosts == {"api.linear.app", "uploads.linear.app"} fetch.return_value["toolkit"]["slug"] = "gmail" with pytest.raises(AuthProviderConfigError, match="toolkit"): await provider.get_auth_result("linear", ["access_token"]) @@ -269,11 +285,14 @@ async def test_managed_linear_binding_only_allows_graphql_origin(monkeypatch): @pytest.mark.asyncio async def test_fresh_instance_resumes_exact_saved_cursor_with_same_scope_fingerprint(): first_source, _ = await source(response({"issues": connection([node()], "saved-page")})) - first = await first_source.capture_page(CompletedScope(record_type="issue"), ScanContinuation()) + first = await first_source.capture_page( + CompletedScope(record_type="issue"), ScanContinuation(), files=MagicMock() + ) fresh_source, client = await source(response({"issues": connection([])})) resumed = await fresh_source.capture_page( CompletedScope(record_type="issue"), ScanContinuation.model_validate_json(first.continuation.model_dump_json()), + files=MagicMock(), ) assert resumed.final assert client.post.call_args.kwargs["json"]["variables"]["after"] == "saved-page" @@ -300,8 +319,144 @@ async def test_attachment_page_is_independent_and_does_not_claim_file_bytes(): ) ) page = await capture.capture_page( - CompletedScope(record_type="attachment", container_id=str(ISSUE)), ScanContinuation() + CompletedScope(record_type="attachment", container_id=str(ISSUE)), + ScanContinuation(), + files=MagicMock(), ) assert page.records[0].payload == attachment assert page.records[0].completeness == "partial" assert page.records[0].blobs == () + + +@pytest.mark.asyncio +async def test_retained_upload_is_stored_before_reference_without_mutating_original( + monkeypatch, tmp_path +): + import hashlib + from contextlib import asynccontextmanager + + from airweave.domains.storage.file_service import FileService + + @asynccontextmanager + async def stream(*args, **kwargs): + assert args == ("GET", "https://uploads.linear.app/workspace/object") + assert kwargs["follow_redirects"] is False + assert kwargs["headers"] == {} + yield httpx.Response( + 200, + content=b"retained", + headers={"content-type": "application/pdf"}, + request=httpx.Request("GET", args[1]), + ) + + capture, client = await source() + client.stream = MagicMock(side_effect=stream) + storage = AsyncMock() + monkeypatch.setattr(FileService, "_ensure_base_dir", lambda self: None) + files = FileService(sync_job_id=UUID(int=10), sync_id=UUID(int=11), storage_backend=storage) + original = {**node(UUID(int=5)), "url": "https://uploads.linear.app/workspace/object"} + retained = await capture._retain( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), original, files + ) + assert retained.payload == original and retained.completeness == "complete" + ref = retained.blobs[0] + assert ref.sha256 == hashlib.sha256(b"retained").hexdigest() + assert ref.source_path == "/url" + assert ref.key.startswith(f"canonical/{UUID(int=11)}/blobs/sha256/") + storage.write_file.assert_awaited_once_with(ref.key, b"retained") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "https://external.example/a", + "http://uploads.linear.app/a", + "https://uploads.linear.app.evil.example/a", + "https://user@uploads.linear.app/a", + "https://uploads.linear.app:444/a", + "https://uploads.linear.app/a#fragment", + ], +) +async def test_non_audited_file_url_never_receives_credentials(url): + capture, client = await source() + files = MagicMock() + original = {**node(UUID(int=5)), "url": url} + retained = await capture._retain( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), original, files + ) + assert retained.completeness == "partial" and not retained.blobs + client.stream.assert_not_called() + files.store_canonical_blob.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status,content,length", [(302, b"", "0"), (200, b"12345", "1"), (200, b"", "5")] +) +async def test_redirect_or_oversize_never_publishes_a_blob(monkeypatch, status, content, length): + from contextlib import asynccontextmanager + + import airweave.platform.sources.linear as linear_module + + @asynccontextmanager + async def stream(*args, **kwargs): + assert kwargs["follow_redirects"] is False + yield httpx.Response( + status, + content=content, + headers={"content-length": length, "location": "https://external.example/"}, + request=httpx.Request("GET", args[1]), + ) + + capture, client = await source() + client.stream = MagicMock(side_effect=stream) + monkeypatch.setattr(linear_module, "MAX_FILE_SIZE_BYTES", 4) + files = AsyncMock() + original = {**node(UUID(int=5)), "url": "https://uploads.linear.app/a"} + if status == 302: + with pytest.raises(SourceError): + await capture._retain( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), original, files + ) + else: + retained = await capture._retain( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), original, files + ) + assert retained.completeness == "partial" and not retained.blobs + files.store_canonical_blob.assert_not_called() + + +@pytest.mark.asyncio +async def test_file_storage_failure_prevents_page_success(): + from contextlib import asynccontextmanager + + @asynccontextmanager + async def stream(*args, **kwargs): + yield httpx.Response(200, content=b"bytes", request=httpx.Request("GET", args[1])) + + attachment = { + **node(UUID(int=5)), + "issue": {"id": str(ISSUE)}, + "url": "https://uploads.linear.app/a", + } + capture, client = await source( + response( + { + "issue": { + "id": str(ISSUE), + "team": {"id": str(TEAM)}, + "attachments": connection([attachment]), + } + } + ) + ) + client.stream = MagicMock(side_effect=stream) + files = AsyncMock() + files.store_canonical_blob.side_effect = OSError("fixture storage failure") + with pytest.raises(OSError): + await capture.capture_page( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), + ScanContinuation(), + files=files, + ) From d796f92d92292d435a8327b3425e5619b870d413 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:49:20 -0700 Subject: [PATCH 047/266] Preserve status classification for unread source error streams --- .../airweave/platform/sources/http_helpers.py | 4 +- .../platform/sources/test_linear_capture.py | 54 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/backend/airweave/platform/sources/http_helpers.py b/backend/airweave/platform/sources/http_helpers.py index 7743a5fb3b..02edda11c6 100644 --- a/backend/airweave/platform/sources/http_helpers.py +++ b/backend/airweave/platform/sources/http_helpers.py @@ -160,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"): @@ -168,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/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py index 3941139ebb..5c27a0dac3 100644 --- a/backend/tests/unit/platform/sources/test_linear_capture.py +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -460,3 +460,57 @@ async def stream(*args, **kwargs): ScanContinuation(), files=files, ) + + +class _UnreadErrorBody(httpx.AsyncByteStream): + def __init__(self): + self.read_count = 0 + self.closed = False + + async def __aiter__(self): + self.read_count += 1 + raise AssertionError("Error response body must not be drained") + yield b"private error detail" # pragma: no cover + + async def aclose(self): + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [302, 401, 429]) +async def test_unread_stream_error_keeps_status_closes_and_never_stores(status): + from airweave.domains.sources.exceptions import SourceAuthError + from airweave.platform.http_client.airweave_client import AirweaveHttpClient + + body = _UnreadErrorBody() + requests = [] + + def handle(request): + requests.append(request) + return httpx.Response( + status, + stream=body, + headers={"Retry-After": "61", "Location": "https://external.example/"}, + ) + + capture, _ = await source() + logger = MagicMock() + files = AsyncMock() + original = {**node(UUID(int=5)), "url": "https://uploads.linear.app/a?signed=private-value"} + expected = {302: SourceError, 401: SourceAuthError, 429: SourceRateLimitError}[status] + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client: + capture._http_client = AirweaveHttpClient( + client, org_id=UUID(int=1), source_short_name="linear", logger=logger + ) + with pytest.raises(expected) as error: + await capture._retain( + CompletedScope(record_type="attachment", container_id=str(ISSUE)), original, files + ) + if status == 429: + assert error.value.retry_after == 61 + if status == 401: + assert error.value.status_code == 401 + assert len(requests) == 1 and body.closed and body.read_count == 0 + files.store_canonical_blob.assert_not_called() + assert "private-value" not in str(logger.debug.call_args_list) + assert "private error detail" not in str(logger.debug.call_args_list) From 01719fc5396287f42d7f39d4fea571bee79b6b9e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 05:59:21 -0700 Subject: [PATCH 048/266] Confirm Linear scope loss through exact-ID connections --- .../canonical/tests/test_linear_recovery.py | 202 ++++++++++++++++++ backend/airweave/platform/sources/linear.py | 30 ++- .../linear_queries/attachments.graphql | 7 +- .../sources/linear_queries/comments.graphql | 7 +- .../linear_queries/issue-membership.graphql | 7 +- .../platform/sources/test_linear_capture.py | 9 +- 6 files changed, 248 insertions(+), 14 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_linear_recovery.py 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/platform/sources/linear.py b/backend/airweave/platform/sources/linear.py index 9726229e8d..dd16c93cbf 100644 --- a/backend/airweave/platform/sources/linear.py +++ b/backend/airweave/platform/sources/linear.py @@ -54,7 +54,7 @@ _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 = 2 +FIELD_SET_VERSION = 3 class _Model(BaseModel): @@ -262,8 +262,12 @@ async def capture_page( variables["issueId"] = scope.container_id name = "comments" if scope.record_type == "comment" else "attachments" data = await self._query(name, variables) - issue_payload = data.get("issue") - # Null without documented error evidence is not a deletion or access-loss proof. + 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 @@ -366,13 +370,27 @@ def _capture(scope: CompletedScope, payload: dict[str, JsonValue]) -> CaptureRec observed_at=datetime.now(timezone.utc), ) + @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_root_absent(self, native_id: str) -> None: """Fail closed on ambiguous absence; readable out-of-scope roots may be removed.""" UUID(native_id) data = await self._query("issue-membership", {"issueId": native_id}) - issue = _Issue.model_validate(data.get("issue")) - if str(issue.id) != native_id: - raise ValueError("Linear returned the wrong issue identity") + 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") diff --git a/backend/airweave/platform/sources/linear_queries/attachments.graphql b/backend/airweave/platform/sources/linear_queries/attachments.graphql index 4481afe0d9..034702fa0e 100644 --- a/backend/airweave/platform/sources/linear_queries/attachments.graphql +++ b/backend/airweave/platform/sources/linear_queries/attachments.graphql @@ -1,9 +1,12 @@ -query OwnedLinearIssueAttachments($issueId: String!, $first: Int!, $after: String) { - issue(id: $issueId) { +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 index fbb8ec803e..ef6ac1b97c 100644 --- a/backend/airweave/platform/sources/linear_queries/comments.graphql +++ b/backend/airweave/platform/sources/linear_queries/comments.graphql @@ -1,5 +1,6 @@ -query OwnedLinearIssueComments($issueId: String!, $first: Int!, $after: String) { - issue(id: $issueId) { +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 { @@ -10,5 +11,7 @@ query OwnedLinearIssueComments($issueId: String!, $first: Int!, $after: String) } pageInfo { hasNextPage endCursor } } + } + pageInfo { hasNextPage endCursor } } } diff --git a/backend/airweave/platform/sources/linear_queries/issue-membership.graphql b/backend/airweave/platform/sources/linear_queries/issue-membership.graphql index af2825ab67..69418b8c2f 100644 --- a/backend/airweave/platform/sources/linear_queries/issue-membership.graphql +++ b/backend/airweave/platform/sources/linear_queries/issue-membership.graphql @@ -1,3 +1,6 @@ -query OwnedLinearIssueMembership($issueId: String!) { - issue(id: $issueId) { id team { id } archivedAt trashed } +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/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py index 5c27a0dac3..cb67fb270c 100644 --- a/backend/tests/unit/platform/sources/test_linear_capture.py +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -38,6 +38,9 @@ def node(identity=ISSUE): def response(data, *, errors=None, status=200, headers=None): + # Child and membership fixtures use the actual exact-ID outer connection. + if isinstance(data, dict) and "issue" in data: + data = {"issues": connection([] if data["issue"] is None else [data["issue"]])} return httpx.Response( status, json={"data": data, **({"errors": errors} if errors else {})}, @@ -177,15 +180,17 @@ async def test_moved_child_scope_reports_scope_removed(): @pytest.mark.asyncio -async def test_missing_issue_is_not_access_revoked_or_empty_success(): +async def test_exhausted_zero_issue_connection_withdraws_scope_without_deletion_claim(): capture, _ = await source(response({"issue": None})) - with pytest.raises(ValidationError): + with pytest.raises(ScopeAccessLost) as error: await capture.capture_page( CompletedScope(record_type="comment", container_id=str(ISSUE)), ScanContinuation(), files=MagicMock(), ) + assert error.value.removal_reason == "scope_removed" + @pytest.mark.asyncio async def test_omitted_readable_selected_root_blocks_reconciliation(): From 634422524322d58a1c138516f3d307b396e97629 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 06:09:17 -0700 Subject: [PATCH 049/266] Read exact captured Calendar identities with bounded version resolution --- .../api/v1/endpoints/calendar_records.py | 20 +- backend/airweave/api/v1/endpoints/records.py | 1 + .../entities/canonical/calendar_exact.py | 157 +++++++++++++ .../canonical/tests/test_calendar_exact.py | 221 ++++++++++++++++++ 4 files changed, 398 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/domains/entities/canonical/calendar_exact.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_calendar_exact.py diff --git a/backend/airweave/api/v1/endpoints/calendar_records.py b/backend/airweave/api/v1/endpoints/calendar_records.py index 5f11715b8d..26f3846c28 100644 --- a/backend/airweave/api/v1/endpoints/calendar_records.py +++ b/backend/airweave/api/v1/endpoints/calendar_records.py @@ -3,7 +3,7 @@ from uuid import UUID from zoneinfo import ZoneInfoNotFoundError -from fastapi import Depends, HTTPException, Query +from fastapi import Depends, HTTPException, Path, Query from pydantic import AwareDatetime, ValidationError from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -13,11 +13,13 @@ from airweave.api.router import TrailingSlashRouter from airweave.core.config import settings from airweave.db.session import get_db +from airweave.domains.entities.canonical.calendar_exact import read_event from airweave.domains.entities.canonical.calendar_query import ( CalendarRange, CalendarRangePage, CalendarRangeService, ) +from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.store import SourceNotFound from airweave.models.source_connection import SourceConnection @@ -63,3 +65,19 @@ async def calendar_events( if await _authorized_connection(db, ctx.organization.id, sync_id) != connection_id: raise SourceNotFound("Calendar source changed while reading") return result + + +@router.get("/{sync_id}/calendar/events/{event_id}", response_model=SourceRecord) +async def calendar_event( + sync_id: UUID, + event_id: str = Path(min_length=1, max_length=1024), + calendar_id: str = Query(min_length=1, max_length=1024), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), +) -> SourceRecord: + """Read one captured provider identity; no live lookup or window expansion.""" + connection = await _authorized_connection(db, ctx.organization.id, sync_id) + result = await read_event(db, ctx.organization.id, sync_id, calendar_id, event_id) + if await _authorized_connection(db, ctx.organization.id, sync_id) != connection: + raise SourceNotFound("Calendar source changed while reading") + return result diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index 42d1811830..b9f463b5f1 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -48,6 +48,7 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> "stale_record_revision": 409, "blob_unavailable": 503, "calendar_changed_restart": 409, + "calendar_read_incomplete": 409, "calendar_range_not_captured": 409, }.get(error.code, 400) detail = {"code": error.code, "message": str(error), "retryable": status == 503} diff --git a/backend/airweave/domains/entities/canonical/calendar_exact.py b/backend/airweave/domains/entities/canonical/calendar_exact.py new file mode 100644 index 0000000000..2575d17c05 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/calendar_exact.py @@ -0,0 +1,157 @@ +"""Bounded Google event identity resolution over already captured observations.""" + +from datetime import date as Date +from typing import Literal, Self +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError, model_validator +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event +from airweave.domains.entities.canonical.calendar_query import CalendarChanged +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.query import RecordNotFound +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.store import ( + CanonicalStoreError, + content_is_available, + parent_is_visible, + source_record, +) +from airweave.models.entity import Entity +from airweave.models.sync import Sync + + +class CalendarReadIncomplete(CanonicalStoreError): + """Existing observations cannot safely establish one scheduling state.""" + + code = "calendar_read_incomplete" + + +class Version(BaseModel): + """Only provider version evidence, never capture wall time.""" + + model_config = ConfigDict(extra="ignore") + etag: str | None = None + updated: AwareDatetime | None = None + status: str | None = None + + +class CancellationStart(BaseModel): + """Only the original occurrence identity may survive a tombstone read.""" + + model_config = ConfigDict(extra="ignore") + date: Date | None = None + dateTime: AwareDatetime | None = None + timeZone: str | None = None + + @model_validator(mode="after") + def exact_start(self) -> Self: + """Retain exactly one date or instant from the provider identity.""" + if (self.date is None) == (self.dateTime is None): + raise ValueError("Exactly one original start is required") + return self + + +class CancellationIdentity(BaseModel): + """Safe recurrence exclusion fields retained when a body is deleted.""" + + model_config = ConfigDict(extra="ignore") + id: str = Field(min_length=1) + status: Literal["cancelled"] + recurringEventId: str = Field(min_length=1) + originalStartTime: CancellationStart + + +def choose(records: tuple[SourceRecord, ...]) -> SourceRecord: + """Resolve at most one raw and one expanded record; denied content never falls through.""" + if any( + r.content_access != "available" or r.removal_reason == "access_revoked" for r in records + ): + raise RecordNotFound("Calendar event is unavailable") + if any( + r.identity.record_type == "event" and r.removal_reason == "scope_removed" for r in records + ): + raise RecordNotFound("Calendar event is unavailable") + candidates = tuple(r for r in records if r.removal_reason != "scope_removed") + if not candidates: + raise RecordNotFound("Calendar event has not been captured or is unavailable") + if len(candidates) == 1: + return candidates[0] + raw, expanded = sorted(candidates, key=lambda r: r.identity.record_type != "event") + try: + a, b = Version.model_validate(raw.payload), Version.model_validate(expanded.payload) + except ValidationError: + raise CalendarReadIncomplete("Calendar observations need a refresh") from None + same_state = (raw.deleted_at is not None or a.status == "cancelled") == ( + expanded.deleted_at is not None or b.status == "cancelled" + ) and a.status == b.status + fields = ("start", "end", "recurringEventId", "originalStartTime") + same_schedule = all(raw.payload.get(key) == expanded.payload.get(key) for key in fields) + if a.etag and a.etag == b.etag and not (same_state and same_schedule): + raise CalendarReadIncomplete("Calendar observations disagree; refresh this calendar") + if a.updated and b.updated and a.updated != b.updated: + return raw if a.updated > b.updated else expanded + if ( + same_state + and same_schedule + and ((a.etag and a.etag == b.etag) or raw.payload == expanded.payload) + ): + return raw + raise CalendarReadIncomplete("Calendar observations disagree; refresh this calendar") + + +async def read_event( + db: AsyncSession, organization_id: UUID, sync_id: UUID, calendar_id: str, event_id: str +) -> SourceRecord: + """Use the existing compound identity index and current calendar visibility.""" + scope = select(Sync.observed_change_sequence).where( + Sync.organization_id == organization_id, Sync.id == sync_id + ) + sequence = await db.scalar(scope) + parent = await db.scalar( + select(Entity.id).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.entity_definition_short_name == "calendar", + Entity.native_id == calendar_id, + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + content_is_available(), + ) + ) + if sequence is None or parent is None: + raise RecordNotFound("Calendar is unavailable") + identity = RecordIdentity(record_type="event", container_id=calendar_id, native_id=event_id) + rows = ( + await db.scalars( + select(Entity) + .where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.entity_id == identity.entity_key, + Entity.entity_definition_short_name.in_(("event", "event_occurrence")), + Entity.record_revision > 0, + parent_is_visible(), + ) + .limit(2) + ) + ).all() + result = choose(tuple(source_record(row) for row in rows)) + db.expire_all() + if await db.scalar(scope) != sequence: + raise CalendarChanged("Calendar changed while reading; try again") + if result.deleted_at is not None: + try: + identity = ( + CancellationIdentity.model_validate(result.payload).model_dump( + mode="json", exclude_none=True + ) + if is_cancelled_recurring_event(result.payload) + else {} + ) + except ValidationError: + raise CalendarReadIncomplete("Calendar cancellation identity needs a refresh") from None + return result.model_copy(update={"payload": identity, "blobs": ()}) + return result diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exact.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exact.py new file mode 100644 index 0000000000..d527d6ef3e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exact.py @@ -0,0 +1,221 @@ +"""Real SQL checks for native identity, version conflicts and withdrawn scope.""" + +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.calendar_exact import CalendarReadIncomplete, read_event +from airweave.domains.entities.canonical.query import RecordNotFound +from airweave.domains.entities.canonical.requests import CaptureBatch +from airweave.platform.sources.records.google_calendar import record + + +async def seed( + database, + source, + raw_status="confirmed", + expanded_status="cancelled", + raw_updated="2026-09-01T00:00:00Z", + expanded_updated="2026-09-02T00:00:00Z", +): + capture, fence = source + parent = record("calendar", {"id": "cal"}) + + def item(kind, status, updated): + payload = { + "id": "exact", + "status": status, + "updated": updated, + "summary": "private body", + "start": {"date": "2026-09-01"}, + "end": {"date": "2026-09-02"}, + } + return record(kind, payload, "cal").model_copy(update={"parent": parent.identity}) + + observations = ( + parent, + item("event", raw_status, raw_updated), + item("event_occurrence", expanded_status, expanded_updated), + ) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=observations)) + return parent + + +async def test_newer_provider_deletion_redacts_body_and_wrong_scope_is_absent(database, source): + await seed(database, source) + _, fence = source + async with database() as db: + result = await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + assert result.deleted_at and result.payload == {} and result.blobs == () + assert result.identity.record_type == "event_occurrence" + with pytest.raises(RecordNotFound): + await read_event(db, fence.organization_id, fence.sync_id, "other", "exact") + with pytest.raises(RecordNotFound): + await read_event(db, uuid4(), fence.sync_id, "cal", "exact") + + +async def test_reinstatement_uses_provider_version_not_capture_order(database, source): + await seed(database, source, "cancelled", "confirmed") + _, fence = source + async with database() as db: + result = await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + assert result.deleted_at is None and result.payload["status"] == "confirmed" + + +async def test_equal_versions_conflicting_status_fail_closed(database, source): + await seed(database, source, expanded_updated="2026-09-01T00:00:00Z") + _, fence = source + async with database() as db: + with pytest.raises(CalendarReadIncomplete): + await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + + +async def test_parent_withdrawal_blocks_both_candidates(database, source): + parent = await seed(database, source) + capture, fence = source + withdrawn = parent.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(withdrawn,))) + async with database() as db: + with pytest.raises(RecordNotFound): + await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + + +async def test_horizon_eviction_does_not_delete_raw_event(database, source): + await seed(database, source, expanded_status="confirmed") + capture, fence = source + eviction = record("event_occurrence", {"id": "exact"}, "cal").model_copy( + update={"kind": "delete", "removal_reason": "scope_removed"} + ) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(eviction,))) + async with database() as db: + result = await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + assert result.deleted_at is None and result.identity.record_type == "event" + + +async def test_missing_native_versions_fail_closed(database, source): + await seed(database, source, raw_updated=None, expanded_updated=None) + _, fence = source + async with database() as db: + with pytest.raises(CalendarReadIncomplete): + await read_event(db, fence.organization_id, fence.sync_id, "cal", "exact") + + +async def test_cancelled_recurring_tombstone_keeps_only_exclusion_identity(database, source): + capture, fence = source + parent = record("calendar", {"id": "cal"}) + exclusion = record( + "event_occurrence", + { + "id": "exception", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-09-01"}, + "summary": "must not leak", + "attendees": [{"email": "private@example.com"}], + }, + "cal", + ).model_copy(update={"parent": parent.identity}) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(parent, exclusion))) + async with database() as db: + result = await read_event(db, fence.organization_id, fence.sync_id, "cal", "exception") + assert result.deleted_at is not None + assert result.payload == { + "id": "exception", + "status": "cancelled", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-09-01"}, + } + + +async def test_exact_http_is_source_authenticated_and_tenant_scoped(database, source): + from types import SimpleNamespace + + from fastapi import FastAPI + from httpx import ASGITransport, AsyncClient + from sqlalchemy import update + + from airweave.api import deps + from airweave.api.v1.endpoints.calendar_records import router + from airweave.api.v1.endpoints.records import record_error_response + from airweave.db.session import get_db + from airweave.domains.entities.canonical.store import CanonicalStoreError + from airweave.models.source_connection import SourceConnection + + await seed(database, source) + _, fence = source + async with database() as db: + db.add( + SourceConnection( + organization_id=fence.organization_id, + sync_id=fence.sync_id, + name="Synthetic", + short_name="google_calendar", + is_authenticated=True, + ) + ) + await db.commit() + app = FastAPI() + app.include_router(router, prefix="/sync") + app.add_exception_handler(CanonicalStoreError, record_error_response) + owner = fence.organization_id + + 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 + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + path = f"/sync/{fence.sync_id}/calendar/events/exact" + result = await client.get(path, params={"calendar_id": "cal"}) + assert result.status_code == 200 and result.json()["payload"] == {} + owner = uuid4() + assert (await client.get(path, params={"calendar_id": "cal"})).status_code == 404 + owner = fence.organization_id + async with database() as db: + await db.execute( + update(SourceConnection) + .where(SourceConnection.sync_id == fence.sync_id) + .values(is_authenticated=False) + ) + await db.commit() + assert (await client.get(path, params={"calendar_id": "cal"})).status_code == 404 + + +async def test_same_cancelled_exception_raw_state_beats_duplicate_tombstone(database, source): + from airweave.domains.entities.canonical.calendar_exact import choose + + capture, fence = source + parent = record("calendar", {"id": "cal"}) + payload = { + "id": "exception", + "status": "cancelled", + "etag": "same", + "recurringEventId": "series", + "originalStartTime": {"date": "2026-09-01"}, + } + observations = tuple( + record(kind, payload, "cal").model_copy(update={"parent": parent.identity}) + for kind in ("event", "event_occurrence") + ) + async with database() as db: + await capture.capture(db, CaptureBatch(fence=fence, records=(parent, *observations))) + async with database() as db: + result = await read_event(db, fence.organization_id, fence.sync_id, "cal", "exception") + assert result.deleted_at is None and result.identity.record_type == "event" + assert result.payload["originalStartTime"] == {"date": "2026-09-01"} + with pytest.raises(RecordNotFound): + choose( + ( + result.model_copy( + update={"content_access": "unavailable", "removal_reason": None} + ), + ) + ) From ebb0bbdc63e3d65c182a342d708a3a82161c884c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 06:42:29 -0700 Subject: [PATCH 050/266] Pin private staging dependencies and separate schema bootstrap --- deploy/DEPENDENCIES.md | 16 +- deploy/staging/README.md | 132 +++++++++ deploy/staging/cnpg-operator.values.yaml | 12 + deploy/staging/network-policy.yaml | 30 +++ deploy/staging/postgres.yaml | 27 ++ deploy/staging/services.yaml | 297 +++++++++++++++++++++ deploy/staging/temporal-schema.values.yaml | 15 ++ deploy/staging/temporal.values.yaml | 57 ++++ deploy/versions.lock.json | 117 ++++++++ 9 files changed, 695 insertions(+), 8 deletions(-) create mode 100644 deploy/staging/README.md create mode 100644 deploy/staging/cnpg-operator.values.yaml create mode 100644 deploy/staging/network-policy.yaml create mode 100644 deploy/staging/postgres.yaml create mode 100644 deploy/staging/services.yaml create mode 100644 deploy/staging/temporal-schema.values.yaml create mode 100644 deploy/staging/temporal.values.yaml create mode 100644 deploy/versions.lock.json diff --git a/deploy/DEPENDENCIES.md b/deploy/DEPENDENCIES.md index b0d1881806..f55e5517c5 100644 --- a/deploy/DEPENDENCIES.md +++ b/deploy/DEPENDENCIES.md @@ -65,9 +65,9 @@ API/worker traffic. The current Compose file's `vespa:8`, `redis:7-alpine`, Svix Docling and inference image tags are floating and are not a release lockfile. Resolve and record their architecture-specific digests before applying manifests. -The official Temporal chart's current main source identifies chart **1.7.0** and -server **1.32.0**. This is a candidate requiring published-chart retrieval and -compatibility verification, not an already tested upgrade. The repository's +Published official Temporal chart **1.7.0** and server **1.32.0** were retrieved +and pinned; the chart renders with the proposed values. This is not yet a live +interoperability test or an already tested upgrade. The repository's Compose `auto-setup:1.24.2` is old development configuration; do not reuse it as production bootstrap. For a fresh isolated namespace, prefer current official server/chart with PostgreSQL default and visibility stores, no Cassandra, @@ -94,11 +94,11 @@ Bootstrap in this order: Vespa query/feed and worker polling. Private service reachability from Almanac must be explicitly arranged; cluster-private is not cross-cloud connectivity. -No executable Helm values are included yet: cluster storage class, database -operator availability, published chart pin and IAM secret references are not -verified. Inventing them would produce an apparently complete but unusable -manifest. The next deliverable is a small dependency version lock and rendered -values against those facts, not a new infrastructure framework. +Published artifacts are now pinned in `versions.lock.json`; review-only values and +private service templates are in `staging/`. Official Temporal and CNPG charts +render locally. Explicit unresolved cluster/storage/IAM/secret inputs prevent a +deployment-ready claim. See `staging/README.md` for bootstrap ownership and +remaining gates; no resources were provisioned. ## Cost envelope and tradeoff diff --git a/deploy/staging/README.md b/deploy/staging/README.md new file mode 100644 index 0000000000..651b0e51e3 --- /dev/null +++ b/deploy/staging/README.md @@ -0,0 +1,132 @@ +# Private staging review templates + +**Prepared, not deployed.** These files deliberately contain `REQUIRED_*` values. +Do not apply this directory. Registry retrieval and Helm rendering were checked on +2026-09-30; no Kubernetes admission, service startup, restore, capacity or cost +claim follows from that check. The $500 monthly cap and quota approval remain open. + +## What is pinned + +`../versions.lock.json` records retrieved chart archives and SHA-256, image indexes +and architecture-specific digests. These manifests choose **linux/amd64**; confirm +actual nodes before use. Multi-platform indexes/arm64 alternatives are evidence, +not a reason to silently change the deployed platform. + +| Component | Published artifact | +| --- | --- | +| Temporal | Official chart 1.7.0, server/admin-tools 1.32.0 | +| PostgreSQL | CNPG chart 0.29.1 / operator 1.30.1; PG 16.15-standard-trixie | +| Redis | 7.4.11-alpine3.21 | +| Svix | Server v1.101.0, not the unrelated SDK version | +| Vespa | Existing CI-tested image index, pinned amd64 child | +| MiniLM | all-MiniLM-L6-v2 inference image, pinned amd64 child | + +Primary evidence: [Temporal chart index](https://temporalio.github.io/helm-charts/index.yaml), +[CNPG chart index](https://cloudnative-pg.github.io/charts/index.yaml), +[CNPG image requirements](https://cloudnative-pg.io/documentation/current/image_requirements/), +[Svix pinned CLI](https://github.com/svix/svix-webhooks/blob/server-v1.101.0/server/svix-server/src/main.rs). +Registry manifests were retrieved from Docker Hub/GHCR by the exact references in +the lockfile (`reference` may be a tag or an index digest). Image bytes were not pulled or executed during preparation. + +## Required facts before deployment + +- Dedicated dependency namespace, actual Kubernetes version/CNPG compatibility, + existing operator/add-on ownership, node architecture and quota. Reuse a maintained + PostgreSQL add-on if suitable; do not install a second CNPG operator by accident. +- Verified persistent storage class, encryption, reclaim policy and volume/AZ + placement; working backup destination/IAM and a restore drill. `postgres.yaml` + has **no backup policy yet** and provisions only the initial source-store DB. +- Provision separate `source_store`, `temporal`, `temporal_visibility`, `svix` + databases and schema/runtime roles. Existing-secret references are names, not + credentials. Grant runtime DML/sequence access and default privileges from the + schema owner; no runtime role owns schemas. Review grants against actual migrations. +- Secrets: `source-store-schema` (CNPG basic-auth username/password matching owner), + Temporal `*-schema` / `*-runtime` (password), `source-redis` (password), Svix + `source-svix-schema` / `source-svix-runtime` (environment described below). + App credentials stay in its separately scoped staging environment group. +- Actual application namespace/pod identity, CNI NetworkPolicy enforcement, DNS and + private Almanac-to-service connectivity. Policies deny incoming traffic then allow + dedicated-namespace peers and explicitly selected application callers. They do + **not** restrict egress or supply cross-cloud connectivity. Operator management, + monitoring and bootstrap callers require reviewed ingress rules before use. +- S3 canonical bucket/workload identity, separate backup bucket, restricted egress + and model cache. No secret values, IAM policies or public ingress are invented here. +- Full cost estimate including nodes, disks, load balancers, NAT/egress and backups; + $500 is the cap, not an estimate already proved by these resource requests. + +`ClusterIP` is not authentication. Redis has a password; Temporal and Vespa rely +on the reviewed private network boundary in this initial topology. Do not expose +those ports publicly. Stronger client isolation/TLS must be chosen if the cluster +is shared with untrusted workloads. The default-deny policy is not a substitute +for verifying the CNI enforces it. + +## Bootstrap ownership + +1. Resolve infrastructure facts above. Create secrets out of band; keep them out + of rendered artifacts/version control. Provision DBs/roles before schema jobs. +2. Render Temporal schema **only** with `temporal-schema.values.yaml` overlay and + `--show-only templates/server-job.yaml`. Never install the entire chart with + that overlay: it would give server pods schema credentials. Use a distinct job + name per release; the chart's runtime no-op schema job otherwise has the same + generated name. Runtime values have `manageSchema: false` and DML credentials. + History shards are deliberately 128 for this fresh staging proposal and become + immutable at initialization. Review before the first bootstrap. +3. Run Svix `svix-server migrate` using schema secret, wait for successful exit. + Runtime explicitly runs `svix-server` without `--run-migrations`; its upstream + default launch script would run migrations on every startup. Both secrets need + appropriate `SVIX_DB_DSN`, `SVIX_REDIS_DSN`, `SVIX_JWT_SECRET`, + `SVIX_QUEUE_TYPE=redis`, `SVIX_CACHE_TYPE=none`. Same JWT secret; distinct DB roles. + Do not copy Compose's allow-all webhook subnet whitelist into staging. +4. Start Temporal; create the dedicated Temporal namespace chosen for the app and + the `SyncId` Keyword search attribute, then verify both. Do not mask errors with + `|| true`. Application settings must match this namespace and private frontend. +5. Deploy the committed `vespa/app` package after replacing `{{EMBEDDING_DIM}}` + with 384, `{{VERSION}}` with the release identifier, and the host in `hosts.xml` + with the stable StatefulSet hostname. Wait for prepare/activate **and convergence**. + `vespa/init-vespa.sh` is a Compose helper whose convergence timeout only warns; + it is not a reviewed staging bootstrap job. The application package artifact + and bootstrap job remain a required next step after hostname/release resolution. +6. Explicit application jobs own `alembic upgrade head`, then + `python -m airweave.db.bootstrap` from the same immutable app image. Create the + scoped service-key organization separately per `backend/SERVICE_AUTH.md`. + API/worker entrypoints do not migrate or seed. Do not put schema credentials in + `porter.yaml` runtime configuration. Existing maintenance schedule ownership stays. +7. Verify Redis persistence, inference `/health`, Vespa feed/query, Temporal worker + polling, source-store API authorization, scoped capture and restore. Probe and + resource tuning remains necessary for Svix/Vespa; these templates don't claim it. + +## Local rendering performed + +Helm 3.22.0 was downloaded into a temporary directory from the official release, +verified against its published SHA-256 and used without changing the system PATH. +Chart archives were verified against the chart-index checksums in the lockfile. +With extracted chart paths and an isolated `HELM` executable: + +```sh +"$HELM" template source-temporal "$TEMPORAL_CHART" --namespace validation-only \ + -f deploy/staging/temporal.values.yaml > /tmp/temporal-review.yaml +"$HELM" template source-temporal "$TEMPORAL_CHART" --namespace validation-only \ + -f deploy/staging/temporal.values.yaml \ + -f deploy/staging/temporal-schema.values.yaml \ + --show-only templates/server-job.yaml > /tmp/temporal-schema-review.yaml +"$HELM" template source-cnpg "$CNPG_CHART" --namespace validation-only \ + -f deploy/staging/cnpg-operator.values.yaml > /tmp/cnpg-review.yaml +``` + +Results: 13 Temporal objects, 1 schema-only job, 22 CNPG objects; YAML parsed and +all rendered workload images carried digests. Temporal emits four single-replica +server deployments, no UI/admin deployment or public service. Schema overlay +renders two schema init containers using separate schema secrets. Runtime render +contains no schema setup command. Initial rendering rejected obsolete top-level +`cassandra`/`elasticsearch`/`prometheus`/`grafana` keys; they were removed because +chart 1.7 no longer bundles these subcharts. This is why copied older values were +not accepted on faith. + +Static service files were parsed and checked for image digests, private Services +and explicit unresolved namespace/storage inputs. These checks do not validate +CRD admission, deployment ordering, credentials or live vendor interoperability. + +Parent review corrected the Vespa reference from an invalid digest-as-tag form +to repository@digest. Six static workload image references were independently +checked for syntax and exact agreement with the lockfile amd64 digests. This is +still not image pull, startup or Kubernetes admission proof. diff --git a/deploy/staging/cnpg-operator.values.yaml b/deploy/staging/cnpg-operator.values.yaml new file mode 100644 index 0000000000..764aaf4a22 --- /dev/null +++ b/deploy/staging/cnpg-operator.values.yaml @@ -0,0 +1,12 @@ +# REVIEW TEMPLATE: resolve every REQUIRED value before any deployment. +image: + repository: ghcr.io/cloudnative-pg/cloudnative-pg + tag: 1.30.1@sha256:b0f9805bd5f98592e7320b5ea624dd10a7a6ef8e477df5650655c94f0b242497 +nodeSelector: + kubernetes.io/arch: amd64 +resources: + requests: + cpu: 100m + memory: 256Mi + limits: + memory: 512Mi diff --git a/deploy/staging/network-policy.yaml b/deploy/staging/network-policy.yaml new file mode 100644 index 0000000000..ab392ad571 --- /dev/null +++ b/deploy/staging/network-policy.yaml @@ -0,0 +1,30 @@ +# Requires a dedicated namespace and verified CNI enforcement. +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: deny-ingress + namespace: REQUIRED_NAMESPACE +spec: + podSelector: {} + policyTypes: + - Ingress + ingress: [] +--- +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: allow-reviewed-dependency-callers + namespace: REQUIRED_NAMESPACE +spec: + podSelector: {} + policyTypes: + - Ingress + ingress: + - from: + - podSelector: {} + - namespaceSelector: + matchLabels: + kubernetes.io/metadata.name: REQUIRED_APPLICATION_NAMESPACE + podSelector: + matchLabels: + REQUIRED_CALLER_LABEL: REQUIRED_CALLER_VALUE diff --git a/deploy/staging/postgres.yaml b/deploy/staging/postgres.yaml new file mode 100644 index 0000000000..dad03ea299 --- /dev/null +++ b/deploy/staging/postgres.yaml @@ -0,0 +1,27 @@ +# REVIEW TEMPLATE: resolve every REQUIRED value before any deployment. +apiVersion: postgresql.cnpg.io/v1 +kind: Cluster +metadata: + name: source-postgres + namespace: REQUIRED_NAMESPACE +spec: + instances: 1 + imageName: ghcr.io/cloudnative-pg/postgresql:16.15-standard-trixie@sha256:1ac54199af47712e3c17753a7b6feb87ce2c7627362bd45a80f91d9fb9075249 + storage: + storageClass: REQUIRED_STORAGE_CLASS + size: 30Gi + resources: + requests: + cpu: 500m + memory: 2Gi + limits: + memory: 3Gi + affinity: + nodeSelector: + kubernetes.io/arch: amd64 + bootstrap: + initdb: + database: source_store + owner: source_store_schema + secret: + name: source-store-schema diff --git a/deploy/staging/services.yaml b/deploy/staging/services.yaml new file mode 100644 index 0000000000..408b1bd37b --- /dev/null +++ b/deploy/staging/services.yaml @@ -0,0 +1,297 @@ +# REVIEW TEMPLATE: not apply-ready; see README gates. +apiVersion: v1 +kind: Service +metadata: + name: redis-headless + namespace: REQUIRED_NAMESPACE +spec: + clusterIP: None + selector: + app.kubernetes.io/name: redis + ports: + - port: 6379 +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: redis + namespace: REQUIRED_NAMESPACE +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: redis + template: + metadata: + labels: + app.kubernetes.io/name: redis + spec: + automountServiceAccountToken: false + nodeSelector: + kubernetes.io/arch: amd64 + containers: + - name: redis + image: docker.io/library/redis:7.4.11-alpine3.21@sha256:ca0acbb137c1dc3339c8b147a58fd6f42775d4599327b50e7b116c23de501af2 + ports: + - containerPort: 6379 + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + memory: 256Mi + command: + - sh + - -ec + args: + - exec redis-server --appendonly yes --appendfsync everysec --requirepass + "$REDIS_PASSWORD" + env: + - name: REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: source-redis + key: password + volumeMounts: + - name: data + mountPath: /data + serviceName: redis-headless + volumeClaimTemplates: + - metadata: + name: data + spec: + accessModes: + - ReadWriteOnce + storageClassName: REQUIRED_STORAGE_CLASS + resources: + requests: + storage: 5Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: redis + namespace: REQUIRED_NAMESPACE +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: redis + ports: + - name: service + port: 6379 + targetPort: 6379 +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: minilm + namespace: REQUIRED_NAMESPACE +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: minilm + template: + metadata: + labels: + app.kubernetes.io/name: minilm + spec: + automountServiceAccountToken: false + nodeSelector: + kubernetes.io/arch: amd64 + containers: + - name: minilm + image: docker.io/semitechnologies/transformers-inference:sentence-transformers-all-MiniLM-L6-v2@sha256:c9c13f047ee31488e686c4a7b2ea7ccfcdec16ee825447450dc047cf97791590 + ports: + - containerPort: 8080 + resources: + requests: + cpu: 500m + memory: 1536Mi + limits: + memory: 1536Mi + env: + - name: ENABLE_CUDA + value: '0' + readinessProbe: + httpGet: + path: /health + port: 8080 + initialDelaySeconds: 15 + periodSeconds: 10 +--- +apiVersion: v1 +kind: Service +metadata: + name: minilm + namespace: REQUIRED_NAMESPACE +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: minilm + ports: + - name: service + port: 8080 + targetPort: 8080 +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: svix + namespace: REQUIRED_NAMESPACE +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: svix + template: + metadata: + labels: + app.kubernetes.io/name: svix + spec: + automountServiceAccountToken: false + nodeSelector: + kubernetes.io/arch: amd64 + containers: + - name: svix + image: docker.io/svix/svix-server:v1.101.0@sha256:378a15491d525efeaf93257b710aca9e65b80fb995345b13d86a77df763e588a + ports: + - containerPort: 8071 + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + memory: 256Mi + command: + - svix-server + envFrom: + - secretRef: + name: source-svix-runtime +--- +apiVersion: v1 +kind: Service +metadata: + name: svix + namespace: REQUIRED_NAMESPACE +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: svix + ports: + - name: service + port: 8071 + targetPort: 8071 +--- +apiVersion: v1 +kind: Service +metadata: + name: vespa-headless + namespace: REQUIRED_NAMESPACE +spec: + clusterIP: None + selector: + app.kubernetes.io/name: vespa + ports: + - port: 8081 +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: vespa + namespace: REQUIRED_NAMESPACE +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: vespa + template: + metadata: + labels: + app.kubernetes.io/name: vespa + spec: + automountServiceAccountToken: false + nodeSelector: + kubernetes.io/arch: amd64 + containers: + - name: vespa + image: docker.io/vespaengine/vespa@sha256:ded3e50c1b74ffddbef7944c13cfdfb699db238a2ae377539caf515708fd4ac5 + ports: + - containerPort: 8081 + - containerPort: 19071 + resources: + requests: + cpu: 2000m + memory: 8Gi + limits: + memory: 8Gi + env: + - name: VESPA_CONFIGSERVERS + value: vespa-0.vespa-headless.REQUIRED_NAMESPACE.svc.cluster.local + volumeMounts: + - name: data + mountPath: /opt/vespa/var + serviceName: vespa-headless + volumeClaimTemplates: + - metadata: + name: data + spec: + accessModes: + - ReadWriteOnce + storageClassName: REQUIRED_STORAGE_CLASS + resources: + requests: + storage: 60Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: vespa + namespace: REQUIRED_NAMESPACE +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: vespa + ports: + - name: service + port: 8081 + targetPort: 8081 +--- +apiVersion: v1 +kind: Service +metadata: + name: vespa-config + namespace: REQUIRED_NAMESPACE +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: vespa + ports: + - port: 19071 +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: svix-schema-REQUIRED_RELEASE + namespace: REQUIRED_NAMESPACE +spec: + backoffLimit: 1 + activeDeadlineSeconds: 600 + template: + metadata: + labels: + app.kubernetes.io/name: svix-schema + spec: + restartPolicy: Never + automountServiceAccountToken: false + nodeSelector: + kubernetes.io/arch: amd64 + containers: + - name: migrate + image: docker.io/svix/svix-server:v1.101.0@sha256:378a15491d525efeaf93257b710aca9e65b80fb995345b13d86a77df763e588a + command: + - svix-server + - migrate + envFrom: + - secretRef: + name: source-svix-schema diff --git a/deploy/staging/temporal-schema.values.yaml b/deploy/staging/temporal-schema.values.yaml new file mode 100644 index 0000000000..539f00124c --- /dev/null +++ b/deploy/staging/temporal-schema.values.yaml @@ -0,0 +1,15 @@ +# REVIEW TEMPLATE: resolve every REQUIRED value before any deployment. +server: + config: + persistence: + datastores: + default: + sql: + user: temporal_schema + existingSecret: temporal-schema + manageSchema: true + visibility: + sql: + user: temporal_visibility_schema + existingSecret: temporal-visibility-schema + manageSchema: true diff --git a/deploy/staging/temporal.values.yaml b/deploy/staging/temporal.values.yaml new file mode 100644 index 0000000000..3c56944506 --- /dev/null +++ b/deploy/staging/temporal.values.yaml @@ -0,0 +1,57 @@ +# REVIEW TEMPLATE: resolve REQUIRED values before deployment. +server: + replicaCount: 1 + image: + repository: docker.io/temporalio/server + tag: 1.32.0@sha256:ca47d4de249b9cc28137628dba77ae5e75e8b313ebb5c801c64615c1c99cbb09 + nodeSelector: + kubernetes.io/arch: amd64 + resources: + requests: + cpu: 250m + memory: 512Mi + limits: + memory: 1Gi + config: + numHistoryShards: 128 + namespaces: + create: false + persistence: + defaultStore: default + visibilityStore: visibility + datastores: + default: + sql: + pluginName: postgres12 + databaseName: temporal + connectAddr: REQUIRED_POSTGRES_HOST:5432 + user: temporal_runtime + existingSecret: temporal-runtime + secretKey: password + createDatabase: false + manageSchema: false + maxConns: 10 + maxIdleConns: 5 + visibility: + sql: + pluginName: postgres12 + databaseName: temporal_visibility + connectAddr: REQUIRED_POSTGRES_HOST:5432 + user: temporal_visibility_runtime + existingSecret: temporal-visibility-runtime + secretKey: password + createDatabase: false + manageSchema: false + maxConns: 10 + maxIdleConns: 5 +admintools: + enabled: false + image: + repository: docker.io/temporalio/admin-tools + tag: 1.32.0@sha256:8bb84e24db4336de48031a66dbda25a7477d300736fe825893e467042135da1f +web: + enabled: false +schema: + useHelmHooks: false + backoffLimit: 2 + activeDeadlineSeconds: 600 diff --git a/deploy/versions.lock.json b/deploy/versions.lock.json new file mode 100644 index 0000000000..b62d5448ec --- /dev/null +++ b/deploy/versions.lock.json @@ -0,0 +1,117 @@ +{ + "verified_at": "2026-09-30", + "status": "registry-and-render-verification-only", + "candidate_platform": "linux/amd64", + "images": [ + { + "registry": "docker.io", + "repository": "temporalio/server", + "manifest_digest": "sha256:c3e752127759616bb1615e0f9ba0e21635aeb5fdeb922de4f371c350955f46ae", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "amd64": "sha256:ca47d4de249b9cc28137628dba77ae5e75e8b313ebb5c801c64615c1c99cbb09", + "arm64": "sha256:c75012952e4c29acb36ce24243dcdb47822c0134719ea48a1870248620a336d4" + }, + "reference": "1.32.0" + }, + { + "registry": "docker.io", + "repository": "temporalio/admin-tools", + "manifest_digest": "sha256:a9f84fb9a374b2374fe2e67c8efc0468ff3f1c66c8a0b14597ec86e349e62bca", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "amd64": "sha256:8bb84e24db4336de48031a66dbda25a7477d300736fe825893e467042135da1f", + "arm64": "sha256:2fe54408e988e3b53860f77f06742592f1b50a4256d64f28f02fc73213ed6687" + }, + "reference": "1.32.0" + }, + { + "registry": "docker.io", + "repository": "library/redis", + "manifest_digest": "sha256:858f009f9709ce576febc734aa78b8f6d624b82571f9ddb6bda4377c833b3499", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "amd64": "sha256:ca0acbb137c1dc3339c8b147a58fd6f42775d4599327b50e7b116c23de501af2", + "arm64": "sha256:1f09a89a207d794a8c61d9edfc26e7c58427de10ccef7c5d18d638df79a63b85" + }, + "reference": "7.4.11-alpine3.21" + }, + { + "registry": "docker.io", + "repository": "svix/svix-server", + "manifest_digest": "sha256:7ba594c766c882dea9c868cd38957755b1d5c8bc5020f81cf5521ba5ab5adc11", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "arm64": "sha256:fc0fb52e82e739bab030b782835e0e539f64e6686076cb6b03f74d0224460b12", + "amd64": "sha256:378a15491d525efeaf93257b710aca9e65b80fb995345b13d86a77df763e588a" + }, + "reference": "v1.101.0" + }, + { + "registry": "docker.io", + "repository": "semitechnologies/transformers-inference", + "manifest_digest": "sha256:75359f8b886cf51b1e4d971172da1c22776acca7321043b63a24b8d23b710879", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "arm64": "sha256:695dbbeaf67d30a585db42a1c2ecebd8351a08ea11197022db680002095227a2", + "amd64": "sha256:c9c13f047ee31488e686c4a7b2ea7ccfcdec16ee825447450dc047cf97791590" + }, + "reference": "sentence-transformers-all-MiniLM-L6-v2" + }, + { + "registry": "docker.io", + "repository": "vespaengine/vespa", + "manifest_digest": "sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd", + "media_type": "application/vnd.docker.distribution.manifest.list.v2+json", + "platform_digests": { + "amd64": "sha256:ded3e50c1b74ffddbef7944c13cfdfb699db238a2ae377539caf515708fd4ac5", + "arm64": "sha256:125ea9b64d0b4c31bd84502a6dfc2f3c1e44487b4fbee606959e3c418b1d1964" + }, + "reference": "sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd" + }, + { + "registry": "ghcr.io", + "repository": "cloudnative-pg/postgresql", + "manifest_digest": "sha256:1dd42563007500b94b943f2a1620f724efedfc9eea2d07a50c6db4582c2431fd", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "amd64": "sha256:1ac54199af47712e3c17753a7b6feb87ce2c7627362bd45a80f91d9fb9075249", + "arm64": "sha256:d825f9b9b11ffd1772504d9cc5a358c7ea524f33aefb38a3f095615acbc04845" + }, + "reference": "16.15-standard-trixie" + }, + { + "registry": "ghcr.io", + "repository": "cloudnative-pg/cloudnative-pg", + "manifest_digest": "sha256:923c267ec29636db3bee20f993d0ec4973fa22998e1adad37da79e4d32b5bc07", + "media_type": "application/vnd.oci.image.index.v1+json", + "platform_digests": { + "amd64": "sha256:b0f9805bd5f98592e7320b5ea624dd10a7a6ef8e477df5650655c94f0b242497", + "arm64": "sha256:9a10dd7d7466310d5d543f618478cb905093cfc7a5463b538f449da626660083" + }, + "reference": "1.30.1" + } + ], + "charts": [ + { + "name": "temporal", + "version": "1.7.0", + "app_version": "1.32.0", + "url": "https://github.com/temporalio/helm-charts/releases/download/temporal-1.7.0/temporal-1.7.0.tgz", + "sha256": "f7b443f637b71fcbddd12900e5be9fe4c18cf741495d7b2cb247670681bf57d4" + }, + { + "name": "cloudnative-pg", + "version": "0.29.1", + "app_version": "1.30.1", + "url": "https://github.com/cloudnative-pg/charts/releases/download/cloudnative-pg-v0.29.1/cloudnative-pg-0.29.1.tgz", + "sha256": "b53d3991fe84bcf38767e7702cae78666265427a127a26fff168ab4207d2b1df" + } + ], + "render_tool": { + "version": "3.22.0", + "url": "https://get.helm.sh/helm-v3.22.0-darwin-arm64.tar.gz", + "sha256": "4c9982a6cdeb458b60258df66b55398ca5b19293f6877faffe2909ad6f23dfe0", + "platform": "darwin/arm64" + } +} From 587e7412ad973b59118b687fd8709c01215efce0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 06:59:16 -0700 Subject: [PATCH 051/266] Exercise real MiniLM retrieval against synthetic Vespa corpus --- .github/workflows/owned-source-store.yml | 8 ++ backend/tests/integration/test_owned_vespa.py | 121 ++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 110eeff482..6be1cac93b 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -37,6 +37,13 @@ jobs: --hostname vespa --health-cmd "curl -f http://localhost:19071/state/v1/health" --health-interval 10s --health-timeout 5s --health-retries 30 + minilm: + image: semitechnologies/transformers-inference@sha256:c9c13f047ee31488e686c4a7b2ea7ccfcdec16ee825447450dc047cf97791590 + env: + ENABLE_CUDA: '0' + ports: + - 8080:8080 + options: --memory=2g --cpus=2 postgres: image: postgres:16.15 env: @@ -78,6 +85,7 @@ jobs: - name: Verify real Vespa schema and retrieval env: OWNED_VESPA_TEST: '1' + OWNED_MINILM_TEST: '1' EMBEDDING_DIMENSIONS: '384' run: make test-index diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index cd5fe2388a..dc1a75891d 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -89,3 +89,124 @@ async def test_real_schema_retrieval_collection_isolation_and_delete(): for url in urls: response = await http.delete(url) assert response.status_code in (200, 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 + + 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) + urls = [] + async with httpx.AsyncClient(timeout=120) as http: + try: + for _attempt in range(60): + try: + health = await http.get("http://localhost:8080/health", timeout=5) + if health.status_code == 200: + 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]) + identities = {} + for (key, title, text), vector in zip(corpus, 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": "{}", + "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}, + } + }, + ) + 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, + ) + 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) From 7073a6e289519440727e7c29fcf1203896a51b39 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:02:46 -0700 Subject: [PATCH 052/266] Preserve Vespa date units and fractional filter boundaries --- .../adapters/vector_db/filter_translator.py | 11 ++++-- .../destinations/vespa/filter_translator.py | 23 ++++++----- backend/tests/integration/test_owned_vespa.py | 31 +++++++++++++++ .../vespa/test_filter_translator.py | 36 +++++++++++++++-- .../search/test_date_filter_translation.py | 39 +++++++++++++++++++ 5 files changed, 125 insertions(+), 15 deletions(-) create mode 100644 backend/tests/unit/search/test_date_filter_translation.py 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/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/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index dc1a75891d..71328c5a74 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -6,6 +6,7 @@ import json import os +from datetime import datetime, timezone from uuid import uuid4 import httpx @@ -18,6 +19,7 @@ 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.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( @@ -41,6 +43,7 @@ async def test_real_schema_retrieval_collection_isolation_and_delete(): 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"}), @@ -81,6 +84,34 @@ async def test_real_schema_retrieval_collection_isolation_and_delete(): 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) diff --git a/backend/tests/unit/platform/destinations/vespa/test_filter_translator.py b/backend/tests/unit/platform/destinations/vespa/test_filter_translator.py index 1084e4a19b..2f20132331 100644 --- a/backend/tests/unit/platform/destinations/vespa/test_filter_translator.py +++ b/backend/tests/unit/platform/destinations/vespa/test_filter_translator.py @@ -165,16 +165,17 @@ def test_range_multiple_operators(self, filter_translator): assert " AND " in result def test_range_datetime_conversion(self, filter_translator): - """Test range condition with datetime string converts to epoch ms.""" + """Test range condition with datetime string converts to epoch seconds.""" filter_dict = { "must": [ {"key": "created_at", "range": {"gte": "2024-01-01T00:00:00Z"}} ] } result = filter_translator.translate(filter_dict) - # Should convert to epoch milliseconds + # Match the transformer and schema: stored timestamps are epoch seconds. assert "created_at >=" in result - assert "1704067200000" in result # 2024-01-01T00:00:00Z in epoch ms + assert "1704067200" in result + assert "1704067200000" not in result # ========================================================================= # HasId Conditions @@ -565,3 +566,32 @@ def test_field_name_map_values_correct(self): assert FIELD_NAME_MAP["source_name"] == "airweave_system_metadata_source_name" assert FIELD_NAME_MAP["access.is_public"] == "access_is_public" + + +@pytest.mark.parametrize("field", ["created_at", "updated_at"]) +@pytest.mark.parametrize("boundary", ["2024-01-01T00:00:00Z", "2023-12-31T16:00:00-08:00"]) +def test_date_filter_matches_actual_feed_timestamp(filter_translator, field, boundary): + """A document on the inclusive boundary must not disappear due to unit mismatch.""" + from airweave.platform.entities.slack import SlackChannelEntity + from airweave.platform.destinations.vespa.transformer import EntityTransformer + + instant = datetime.fromisoformat("2024-01-01T00:00:00+00:00") + entity = SlackChannelEntity(channel_id="boundary", title="Boundary", purpose="", topic="", breadcrumbs=[]) + entity.created_at = entity.updated_at = instant + stored = EntityTransformer()._build_base_fields(entity)[field] + assert stored == 1704067200 + clause = filter_translator._translate_range_condition( + {"key": field, "range": {"gte": boundary, "lt": "2024-01-01T00:00:01Z"}} + ) + assert clause == f"{field} >= {stored} AND {field} < {stored + 1}" + + +def test_subsecond_boundary_is_not_rounded_into_another_second(filter_translator): + clause = filter_translator._translate_range_condition( + {"key": "created_at", "range": {"gte": "2024-01-01T00:00:00.5Z"}} + ) + assert clause == "created_at >= 1704067200.5" + + +def test_naive_legacy_timestamp_is_explicitly_utc(filter_translator): + assert filter_translator._parse_datetime_to_epoch_seconds("2024-01-01T00:00:00") == 1704067200 diff --git a/backend/tests/unit/search/test_date_filter_translation.py b/backend/tests/unit/search/test_date_filter_translation.py new file mode 100644 index 0000000000..bd84db6d02 --- /dev/null +++ b/backend/tests/unit/search/test_date_filter_translation.py @@ -0,0 +1,39 @@ +"""Date bounds use the same seconds unit as Vespa documents, without truncation.""" + +from unittest.mock import Mock + +import pytest + +from airweave.domains.search.adapters.vector_db.filter_translator import FilterTranslator +from airweave.domains.search.types.filters import FilterCondition, FilterGroup + + +@pytest.mark.parametrize("field", ["created_at", "updated_at"]) +@pytest.mark.parametrize( + "instant", ["2024-01-01T00:00:00Z", "2023-12-31T16:00:00-08:00", "2024-01-01T00:00:00"] +) +def test_integral_boundary_uses_utc_seconds(field, instant): + translator = FilterTranslator(Mock()) + assert translator._parse_datetime_to_epoch(instant, field) == 1704067200 + + +def test_fractional_half_open_bounds_survive_typed_filter_translation(): + translator = FilterTranslator(Mock()) + result = translator.translate( + [ + FilterGroup( + conditions=[ + FilterCondition( + field="created_at", + operator="greater_than_or_equal", + value="2024-01-01T00:00:00.5Z", + ), + FilterCondition( + field="created_at", operator="less_than", value="2024-01-01T00:00:01.5Z" + ), + ] + ) + ] + ) + assert "created_at >= 1704067200.5" in result + assert "created_at < 1704067201.5" in result From 4061dfb14670034b38439a91dbe76f84621a7d7a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:10:43 -0700 Subject: [PATCH 053/266] Use actual transformer readiness endpoint across runtime and deployment --- backend/airweave/domains/embedders/config.py | 2 +- .../embedders/tests/test_local_reachability.py | 4 ++-- backend/tests/integration/test_owned_vespa.py | 4 ++-- deploy/staging/README.md | 12 +++++++++++- deploy/staging/services.yaml | 2 +- docker/docker-compose.test.yml | 2 +- docker/docker-compose.yml | 2 +- 7 files changed, 19 insertions(+), 9 deletions(-) diff --git a/backend/airweave/domains/embedders/config.py b/backend/airweave/domains/embedders/config.py index 125c20f3a4..94849fcda0 100644 --- a/backend/airweave/domains/embedders/config.py +++ b/backend/airweave/domains/embedders/config.py @@ -165,7 +165,7 @@ def _validate_local_reachability(dense_spec: DenseEmbedderEntry) -> None: return inference_url = settings.TEXT2VEC_INFERENCE_URL - health_url = f"{inference_url}/health" + health_url = f"{inference_url}/.well-known/ready" try: with httpx.Client(timeout=httpx.Timeout(5.0)) as client: diff --git a/backend/airweave/domains/embedders/tests/test_local_reachability.py b/backend/airweave/domains/embedders/tests/test_local_reachability.py index 4be2eb8c44..9110cfc153 100644 --- a/backend/airweave/domains/embedders/tests/test_local_reachability.py +++ b/backend/airweave/domains/embedders/tests/test_local_reachability.py @@ -62,7 +62,7 @@ def test_skips_non_local_embedder(self): @patch("airweave.domains.embedders.config.settings") @patch("airweave.domains.embedders.config.httpx") def test_passes_when_service_reachable(self, mock_httpx, mock_settings): - """Passes when the health endpoint returns 200.""" + """Passes when the readiness endpoint responds successfully.""" mock_settings.TEXT2VEC_INFERENCE_URL = "http://localhost:9878" mock_httpx.Timeout = httpx.Timeout client = _mock_client() @@ -71,7 +71,7 @@ def test_passes_when_service_reachable(self, mock_httpx, mock_settings): entry = _make_entry(LocalDenseEmbedder) _validate_local_reachability(entry) - client.get.assert_called_once_with("http://localhost:9878/health") + client.get.assert_called_once_with("http://localhost:9878/.well-known/ready") @patch("airweave.domains.embedders.config.settings") @patch("airweave.domains.embedders.config.httpx") diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index 71328c5a74..45f1f96d7c 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -163,8 +163,8 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): try: for _attempt in range(60): try: - health = await http.get("http://localhost:8080/health", timeout=5) - if health.status_code == 200: + health = await http.get("http://localhost:8080/.well-known/ready", timeout=5) + if health.status_code == 204: break except httpx.HTTPError: pass diff --git a/deploy/staging/README.md b/deploy/staging/README.md index 651b0e51e3..32a6d8ff55 100644 --- a/deploy/staging/README.md +++ b/deploy/staging/README.md @@ -91,7 +91,7 @@ for verifying the CNI enforces it. scoped service-key organization separately per `backend/SERVICE_AUTH.md`. API/worker entrypoints do not migrate or seed. Do not put schema credentials in `porter.yaml` runtime configuration. Existing maintenance schedule ownership stays. -7. Verify Redis persistence, inference `/health`, Vespa feed/query, Temporal worker +7. Verify Redis persistence, inference `/.well-known/ready` (HTTP 204), Vespa feed/query, Temporal worker polling, source-store API authorization, scoped capture and restore. Probe and resource tuning remains necessary for Svix/Vespa; these templates don't claim it. @@ -130,3 +130,13 @@ Parent review corrected the Vespa reference from an invalid digest-as-tag form to repository@digest. Six static workload image references were independently checked for syntax and exact agreement with the lockfile amd64 digests. This is still not image pull, startup or Kubernetes admission proof. + +### Readiness correction from actual-model CI + +Run36725698424 reached the real inference container but its old `/health` probe +returned404, so model retrieval assertions did not run. Startup validation, +Compose, staging readiness and the integration probe now use +`/.well-known/ready`, matching the [upstream inference API](https://docs.weaviate.io/weaviate/modules/custom-modules). +The integration test requires204 and still verifies actual384-dimensional vectors +and retrieval afterward. This fixes the contract; a successful rerun is required +before claiming the pinned deployment works. diff --git a/deploy/staging/services.yaml b/deploy/staging/services.yaml index 408b1bd37b..7a5de632a2 100644 --- a/deploy/staging/services.yaml +++ b/deploy/staging/services.yaml @@ -115,7 +115,7 @@ spec: value: '0' readinessProbe: httpGet: - path: /health + path: /.well-known/ready port: 8080 initialDelaySeconds: 15 periodSeconds: 10 diff --git a/docker/docker-compose.test.yml b/docker/docker-compose.test.yml index 8fb7324ce6..f84ff257da 100644 --- a/docker/docker-compose.test.yml +++ b/docker/docker-compose.test.yml @@ -40,7 +40,7 @@ services: ENABLE_CUDA: 0 WORKERS_PER_NODE: 1 healthcheck: - test: [ "CMD", "wget", "--spider", "-q", "http://localhost:8080/health" ] + test: [ "CMD", "wget", "-q", "-O", "/dev/null", "http://localhost:8080/.well-known/ready" ] interval: 30s timeout: 10s retries: 3 diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index a50d553ddc..9a0a75f172 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -157,7 +157,7 @@ services: ENABLE_CUDA: 0 WORKERS_PER_NODE: 1 healthcheck: - test: [ "CMD", "wget", "--spider", "-q", "http://localhost:8080/health" ] + test: [ "CMD", "wget", "-q", "-O", "/dev/null", "http://localhost:8080/.well-known/ready" ] interval: 30s timeout: 10s retries: 3 From 638c86a7db067f2bb428d76ca9dc225aa94c7336 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:13:58 -0700 Subject: [PATCH 054/266] Keep readiness test fixture compliant with lint --- .../tests/test_local_reachability.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/backend/airweave/domains/embedders/tests/test_local_reachability.py b/backend/airweave/domains/embedders/tests/test_local_reachability.py index 9110cfc153..36d02e0599 100644 --- a/backend/airweave/domains/embedders/tests/test_local_reachability.py +++ b/backend/airweave/domains/embedders/tests/test_local_reachability.py @@ -16,19 +16,19 @@ def _make_entry(embedder_class: type, **overrides) -> DenseEmbedderEntry: """Build a minimal DenseEmbedderEntry for testing.""" - defaults = dict( - short_name="test", - name="Test", - description="test entry", - class_name=embedder_class.__name__, - provider="local", - api_model_name="test-model", - max_dimensions=384, - max_tokens=512, - supports_matryoshka=False, - embedder_class_ref=embedder_class, - required_setting="TEXT2VEC_INFERENCE_URL", - ) + defaults = { + "short_name": "test", + "name": "Test", + "description": "test entry", + "class_name": embedder_class.__name__, + "provider": "local", + "api_model_name": "test-model", + "max_dimensions": 384, + "max_tokens": 512, + "supports_matryoshka": False, + "embedder_class_ref": embedder_class, + "required_setting": "TEXT2VEC_INFERENCE_URL", + } defaults.update(overrides) return DenseEmbedderEntry(**defaults) From 792d4d7fd53c6e29620787854b9d102562295383 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:20:14 -0700 Subject: [PATCH 055/266] Exercise actual BM25 document and query vectors in search baseline --- backend/tests/integration/test_owned_vespa.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index 45f1f96d7c..ff746f7127 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -131,6 +131,7 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): import asyncio from airweave.domains.embedders.dense.local import LocalDenseEmbedder + from airweave.domains.embedders.sparse.fastembed import FastEmbedSparseEmbedder corpus = [ ( @@ -158,6 +159,7 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): ] 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: @@ -173,8 +175,13 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): 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 in zip(corpus, vectors, strict=True): + 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 @@ -190,13 +197,21 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): "entity_id": document_id, "name": title, "textual_representation": text, - "payload": "{}", + "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 + ) + } + }, } }, ) @@ -220,6 +235,9 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): 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( From 4e180f9cc42459a336fd929d66f2f3546f121558 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:21:21 -0700 Subject: [PATCH 056/266] Prefilter canonical dates and types with explicit scoped re-projection --- Makefile | 2 +- .../domains/entities/canonical/projector.py | 16 ++- .../entities/canonical/reprojection.py | 110 ++++++++++++++++++ .../entities/canonical/search_metadata.py | 30 +++++ .../canonical/tests/test_owned_search.py | 48 +++++++- .../canonical/tests/test_reprojection.py | 96 +++++++++++++++ .../canonical/tests/test_search_metadata.py | 44 +++++++ backend/airweave/domains/search/owned.py | 82 +++++++++++-- .../airweave/domains/search/types/filters.py | 10 ++ .../destinations/vespa/transformer.py | 10 ++ backend/airweave/platform/entities/_base.py | 7 ++ backend/scripts/REPROJECTION.md | 50 ++++++++ backend/scripts/reproject_canonical.py | 62 ++++++++++ .../test_canonical_search_prefilters.py | 97 +++++++++++++++ vespa/app/schemas/base_entity.sd | 25 ++++ 15 files changed, 674 insertions(+), 15 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/reprojection.py create mode 100644 backend/airweave/domains/entities/canonical/search_metadata.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_reprojection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_search_metadata.py create mode 100644 backend/scripts/REPROJECTION.md create mode 100644 backend/scripts/reproject_canonical.py create mode 100644 backend/tests/integration/test_canonical_search_prefilters.py diff --git a/Makefile b/Makefile index 722a384bbc..0a26ec1be2 100644 --- a/Makefile +++ b/Makefile @@ -48,6 +48,6 @@ test-health: test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py tests/integration/test_canonical_search_prefilters.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py check: test-store test-capture test-search test-auth test-health diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 66277fa480..ddbb09a20c 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -1,6 +1,6 @@ """Bounded original-record projection using existing conversion/embed/feed primitives.""" -from collections.abc import Callable +from collections.abc import Callable, Iterable from contextlib import AbstractAsyncContextManager from dataclasses import dataclass from uuid import UUID, uuid4 @@ -8,6 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.core.logging import ContextualLogger +from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.projection_models import ( ProjectionBatchResult, ProjectionDocument, @@ -16,11 +17,12 @@ scope_projection_document_id, ) from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore +from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata from airweave.domains.storage.protocols import StorageBackend from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields from airweave.platform.destinations.vespa.destination import VespaDestination -from airweave.platform.entities._base import AirweaveSystemMetadata +from airweave.platform.entities._base import AirweaveSystemMetadata, BaseEntity @dataclass @@ -47,6 +49,14 @@ class ProjectionRuntime: entity_tracker: StrictProjectionTracker +def _stamp_chunks(chunks: Iterable[BaseEntity], record: SourceRecord) -> None: + """Conversion cannot replace the source-owned searchable metadata.""" + for chunk in chunks: + if chunk.airweave_system_metadata is None: + raise ValueError("Projection chunk lost canonical metadata") + stamp_search_metadata(chunk.airweave_system_metadata, record) + + class CanonicalProjector: """Read snapshot, compute outside transaction, then atomically publish generation.""" @@ -103,12 +113,14 @@ async def project_one( meta.sync_id = work.record.sync_id meta.sync_job_id = None meta.db_entity_id = work.record.id + stamp_search_metadata(meta, work.record) chunks = await self._processor.process( list(mapped), ProjectionContext(logger, source_name), ProjectionRuntime(StrictProjectionTracker()), strict=True, ) + _stamp_chunks(chunks, work.record) prepared = destination.prepare_documents(chunks) for group in prepared.values(): for document in group: diff --git a/backend/airweave/domains/entities/canonical/reprojection.py b/backend/airweave/domains/entities/canonical/reprojection.py new file mode 100644 index 0000000000..c01984d622 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/reprojection.py @@ -0,0 +1,110 @@ +"""Scoped operator version change; existing projection rows remain the durable queue.""" + +from uuid import UUID + +from pydantic import BaseModel, ConfigDict +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.search_metadata import SEARCH_METADATA_PIPELINE_VERSION +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 + + +class ReprojectionPlan(BaseModel): + """Inspectable target and counts, with no provider credentials.""" + + model_config = ConfigDict(frozen=True) + organization_id: UUID + sync_id: UUID + source_connection_id: UUID + provider: str + current_version: int + target_version: int + captured_records: int + target_pending_records: int + workflow_id: str + applied: bool + + +async def plan_reprojection( + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + *, + expected_version: int, + target_version: int, + apply: bool = False, +) -> ReprojectionPlan: + """Caller commits before starting workflow; identical retries never bump twice.""" + if expected_version < 1 or target_version != expected_version + 1: + raise ValueError("Target version must be exactly expected version plus one") + if target_version < SEARCH_METADATA_PIPELINE_VERSION: + raise ValueError("Target version lacks current canonical search metadata") + sync = await db.scalar( + select(Sync) + .where( + Sync.id == sync_id, + Sync.organization_id == organization_id, + ) + .with_for_update() + ) + if sync is None: + raise ValueError("Selected organization/sync is unavailable") + connection = ( + await db.scalars( + select(SourceConnection) + .join( + Collection, + ( + (Collection.readable_id == SourceConnection.readable_collection_id) + & (Collection.organization_id == organization_id) + ), + ) + .where( + SourceConnection.organization_id == organization_id, + SourceConnection.sync_id == sync_id, + SourceConnection.is_authenticated.is_(True), + ) + ) + ).one_or_none() + if connection is None: + raise ValueError("Selected authenticated source binding is unavailable") + if sync.index_pipeline_version not in (expected_version, target_version): + raise ValueError("Projection version changed; inspect before choosing a new operation") + count, pending = ( + await db.execute( + select( + func.count(), + func.count().filter( + Entity.indexed_pipeline_version.is_distinct_from(target_version) + | Entity.indexed_revision.is_distinct_from(Entity.record_revision) + | Entity.indexed_generation.is_(None) + ), + ).where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.record_revision > 0, + ) + ) + ).one() + if not count: + raise ValueError("Selected source has no captured canonical records") + plan = ReprojectionPlan( + organization_id=organization_id, + sync_id=sync_id, + source_connection_id=connection.id, + provider=connection.short_name, + current_version=sync.index_pipeline_version, + target_version=target_version, + captured_records=count, + target_pending_records=pending, + workflow_id=f"reprojection:{organization_id}:{sync_id}:{target_version}", + applied=apply, + ) + if apply: + sync.index_pipeline_version = target_version + await db.flush() + return plan diff --git a/backend/airweave/domains/entities/canonical/search_metadata.py b/backend/airweave/domains/entities/canonical/search_metadata.py new file mode 100644 index 0000000000..d05e1e627c --- /dev/null +++ b/backend/airweave/domains/entities/canonical/search_metadata.py @@ -0,0 +1,30 @@ +"""Canonical-owned search attributes, independent of provider entity date fallbacks.""" + +from datetime import datetime, timezone + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.platform.entities._base import AirweaveSystemMetadata + +SEARCH_METADATA_PIPELINE_VERSION = 2 +_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc) + + +def epoch_microseconds(value: datetime) -> int: + """Preserve source precision without float timestamp rounding.""" + if value.tzinfo is None: + raise ValueError("Canonical search timestamp must be timezone-aware") + delta = value - _EPOCH + return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds + + +def stamp_search_metadata(meta: AirweaveSystemMetadata, record: SourceRecord) -> None: + """Use only the committed source record, never mapper or observation timestamps.""" + meta.canonical_record_type = record.identity.record_type + meta.source_created_known = int(record.source_created_at is not None) + meta.source_updated_known = int(record.source_updated_at is not None) + meta.source_created_us = ( + epoch_microseconds(record.source_created_at) if record.source_created_at else None + ) + meta.source_updated_us = ( + epoch_microseconds(record.source_updated_at) if record.source_updated_at else None + ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 3a11f915cc..388727c3b5 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -23,6 +23,7 @@ from airweave.domains.search.types import SearchResults from airweave.models.collection import Collection from airweave.models.source_connection import SourceConnection +from airweave.models.sync import Sync from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata @@ -54,6 +55,11 @@ async def indexed(database, source): ) db.add(connection) await db.commit() + async with database() as db: + await db.execute( + update(Sync).where(Sync.id == fence.sync_id).values(index_pipeline_version=2) + ) + await db.commit() store = CanonicalProjectionStore() async with database() as db: work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] @@ -63,7 +69,7 @@ async def indexed(database, source): locator = ProjectionLocator( record_id=work.record.id, revision=1, - pipeline_version=1, + pipeline_version=2, generation=generation, part_index=0, ) @@ -398,3 +404,43 @@ async def test_email_route_uses_visible_canonical_provider_payload( ) assert response.status_code == 200, response.text assert response.json()["items"][0]["email_thread_id"] == expected + + +async def test_filtered_search_requires_explicit_reprojection(database, indexed, http_search): + fence, _, _ = indexed + client, vector, _, _, _ = http_search + async with database() as db: + await db.execute( + update(Sync).where(Sync.id == fence.sync_id).values(index_pipeline_version=1) + ) + await db.commit() + result = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "record_types": ["event"]}, + ) + assert result.status_code == 409 + assert result.json()["detail"]["code"] == "reindex_required" + assert vector._calls == [] + + +async def test_dates_and_types_prefilter_before_candidate_limit(indexed, http_search): + fence, _, _ = indexed + client, vector, _, _, _ = http_search + response = await client.post( + "/sync/search", + json={ + "query": "budget", + "sync_ids": [str(fence.sync_id)], + "mode": "keyword", + "record_types": ["event"], + "created_after": "2026-01-01T00:00:00.000001Z", + "updated_before": "2026-02-01T00:00:00Z", + }, + ) + assert response.status_code == 200, response.text + conditions = vector._calls[0][1].filter_groups[0].conditions + by_field = {condition.field.value: condition.value for condition in conditions} + assert by_field["airweave_system_metadata.canonical_record_type"] == ["event"] + assert by_field["airweave_system_metadata.source_created_known"] == 1 + assert by_field["airweave_system_metadata.source_created_us"] == 1767225600000001 + assert by_field["airweave_system_metadata.source_updated_known"] == 1 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..94a8571bfd --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_reprojection.py @@ -0,0 +1,96 @@ +"""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, monkeypatch # noqa: F811 +): + 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"].endswith(":3") + 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_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/search/owned.py b/backend/airweave/domains/search/owned.py index fba0efe002..cebb66b458 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -12,6 +12,10 @@ from airweave.domains.entities.canonical.projection_models import ProjectionLocator from airweave.domains.entities.canonical.projection_store import publication_matches 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.store import content_is_available from airweave.domains.search.owned_models import ( OwnedSearchCoverage, @@ -42,6 +46,22 @@ async def search( ) -> OwnedSearchResponse: """Resolve exact authorized scopes before any embedding/index request.""" scopes, groups = await self._resolve_scopes(db, ctx, request) + if self._filtered(request): + versions = ( + await db.scalars( + select(Sync.index_pipeline_version).where( + Sync.organization_id == ctx.organization.id, Sync.id.in_(request.sync_ids) + ) + ) + ).all() + if any(version < SEARCH_METADATA_PIPELINE_VERSION for version in versions): + raise HTTPException( + 409, + { + "code": "reindex_required", + "message": "Selected sources need canonical search metadata re-projection", + }, + ) scope_snapshot = self._scope_identity(scopes) hits, scores, exclusions, postfiltered = {}, {}, 0, 0 engine_partial, full = False, False @@ -53,17 +73,7 @@ async def search( offset=0, retrieval_strategy=request.mode, ), - user_filter=[ - FilterGroup( - conditions=[ - FilterCondition( - field="airweave_system_metadata.sync_id", - operator="in", - value=[str(item) for item in sync_ids], - ) - ] - ) - ], + user_filter=[FilterGroup(conditions=self._prefilters(request, sync_ids))], collection_id=str(collection_id), db=db, ctx=ctx, @@ -102,6 +112,56 @@ async def search( or any(row.pending_records for row in sources), ) + @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, SourceConnection], dict[tuple[UUID, str], list[UUID]]]: 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/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/_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/scripts/REPROJECTION.md b/backend/scripts/REPROJECTION.md new file mode 100644 index 0000000000..091eb63e3c --- /dev/null +++ b/backend/scripts/REPROJECTION.md @@ -0,0 +1,50 @@ +# 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 Sync creation still defaults to version 1 in this isolated implementation. +Until integrated initialization is changed, newly captured sources also need the +explicit upgrade. The next coherent creation change is in the existing sync +creation transaction (`domains/syncs/service.py::_create_sync_records`): initialize +new canonical sources to the current projection contract after source capability +resolution, without resetting any existing Sync version. That change is deferred +while the independent capture engine/live lifecycle is active; it must be finished +before presenting new-source date/type search as automatically ready. diff --git a/backend/scripts/reproject_canonical.py b/backend/scripts/reproject_canonical.py new file mode 100644 index 0000000000..b6199a1df0 --- /dev/null +++ b/backend/scripts/reproject_canonical.py @@ -0,0 +1,62 @@ +"""Dry-run by default; scoped re-projection without provider capture or writes.""" + +import argparse +import asyncio +from uuid import UUID + +from temporalio.common import WorkflowIDConflictPolicy, WorkflowIDReusePolicy + +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() + await client.start_workflow( + ProjectCanonicalRecordsWorkflow.run, + args=[str(args.organization), str(args.sync)], + id=plan.workflow_id, + task_queue=settings.TEMPORAL_TASK_QUEUE, + id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, + id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, + ) + 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/tests/integration/test_canonical_search_prefilters.py b/backend/tests/integration/test_canonical_search_prefilters.py new file mode 100644 index 0000000000..f211fb5d7d --- /dev/null +++ b/backend/tests/integration/test_canonical_search_prefilters.py @@ -0,0 +1,97 @@ +"""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(203): + identity = f"canonical-prefilter-{uuid4()}" + url = f"http://localhost:8081/document/v1/airweave/base_entity/docid/{identity}" + urls.append(url) + correct = number == 202 + 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_canonical_record_type": "event", + "airweave_system_metadata_source_created_known": int(known), + "airweave_system_metadata_source_created_us": epoch_microseconds(instant) + if correct + 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/vespa/app/schemas/base_entity.sd b/vespa/app/schemas/base_entity.sd index 09dc6c306d..e3c55838ed 100644 --- a/vespa/app/schemas/base_entity.sd +++ b/vespa/app/schemas/base_entity.sd @@ -66,6 +66,31 @@ schema base_entity { attribute: fast-search } + field airweave_system_metadata_canonical_record_type type string { + indexing: attribute | summary + attribute: fast-search + } + + field airweave_system_metadata_source_created_us type long { + indexing: attribute | summary + attribute: fast-search + } + + field airweave_system_metadata_source_updated_us type long { + indexing: attribute | summary + attribute: fast-search + } + + field airweave_system_metadata_source_created_known type int { + indexing: attribute | summary + attribute: fast-search + } + + field airweave_system_metadata_source_updated_known type int { + indexing: attribute | summary + attribute: fast-search + } + # ============================================================================= # System Metadata Fields (flattened) # ============================================================================= From 1bf91a3a67dbf2f56e38f85d257e90af8ad60e4e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:23:21 -0700 Subject: [PATCH 057/266] Exercise canonical type and unknown date prefilter exclusions --- .../integration/test_canonical_search_prefilters.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backend/tests/integration/test_canonical_search_prefilters.py b/backend/tests/integration/test_canonical_search_prefilters.py index f211fb5d7d..9e212f64a7 100644 --- a/backend/tests/integration/test_canonical_search_prefilters.py +++ b/backend/tests/integration/test_canonical_search_prefilters.py @@ -35,11 +35,11 @@ async def test_prefilter_excludes_over_two_hundred_wrong_dates_before_top_matche urls = [] expected = None try: - for number in range(203): + 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 == 202 + correct = number == 203 known = number != 201 fields = { "entity_id": identity, @@ -49,7 +49,11 @@ async def test_prefilter_excludes_over_two_hundred_wrong_dates_before_top_matche "airweave_system_metadata_collection_id": collection, "airweave_system_metadata_sync_id": str(sync), "airweave_system_metadata_original_entity_id": identity, - "airweave_system_metadata_canonical_record_type": "event", + "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 correct From b825159247141e1e39da14e8a493d6099accd455 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:30:02 -0700 Subject: [PATCH 058/266] Initialize new canonical syncs with current projection metadata version --- backend/airweave/crud/crud_sync.py | 6 ++++++ .../airweave/domains/syncs/fakes/repository.py | 3 ++- backend/airweave/domains/syncs/protocols.py | 1 + backend/airweave/domains/syncs/repository.py | 9 ++++++++- backend/airweave/domains/syncs/service.py | 17 ++++++++++++++++- .../domains/syncs/tests/test_service.py | 17 +++++++++++++++-- backend/scripts/REPROJECTION.md | 15 +++++++-------- 7 files changed, 55 insertions(+), 13 deletions(-) diff --git a/backend/airweave/crud/crud_sync.py b/backend/airweave/crud/crud_sync.py index 2ab913372b..a75bcfc951 100644 --- a/backend/airweave/crud/crud_sync.py +++ b/backend/airweave/crud/crud_sync.py @@ -333,6 +333,7 @@ async def create( obj_in: SyncCreate, ctx: BaseContext, uow: Optional[UnitOfWork] = None, + initial_pipeline_version: int = 1, ) -> schemas.Sync: """Create a sync. @@ -347,11 +348,16 @@ async def create( obj_in (SyncCreate): The sync to create ctx (BaseContext): The API context uow (UnitOfWork, optional): The unit of work + initial_pipeline_version: Server-selected version for this new sync only. + Returns: schemas.Sync: The model validated schema of the created sync """ # Dump the obj_in to a dict + if initial_pipeline_version < 1: + raise ValueError("Initial projection version must be positive") obj_in_dict = obj_in.model_dump() + obj_in_dict["index_pipeline_version"] = initial_pipeline_version obj_in_dict.pop("run_immediately") # Pop off the connection ids 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/protocols.py b/backend/airweave/domains/syncs/protocols.py index 7d17b55fe9..456f3109f4 100644 --- a/backend/airweave/domains/syncs/protocols.py +++ b/backend/airweave/domains/syncs/protocols.py @@ -53,6 +53,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.""" ... 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..f27f0aee97 100644 --- a/backend/airweave/domains/syncs/service.py +++ b/backend/airweave/domains/syncs/service.py @@ -29,6 +29,9 @@ ) 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.sources.exceptions.classifier import classify_error from airweave.domains.sources.types import SourceRegistryEntry from airweave.domains.sync_pipeline.config import SyncConfig @@ -131,6 +134,11 @@ async def create( destination_connection_ids=destination_connection_ids, cron_schedule=cron, run_immediately=run_immediately, + initial_pipeline_version=( + SEARCH_METADATA_PIPELINE_VERSION + if isinstance(source_entry.source_class_ref, (CanonicalSource, CanonicalPageSource)) + else 1 + ), ctx=ctx, uow=uow, ) @@ -470,6 +478,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 +495,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..c4bb68b24e 100644 --- a/backend/airweave/domains/syncs/tests/test_service.py +++ b/backend/airweave/domains/syncs/tests/test_service.py @@ -953,6 +953,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 +1066,10 @@ 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, expected_version", [("legacy", 1), ("gmail", 2), ("slack", 2)] +) +async def test_create_with_cron_calls_temporal_schedule(source_kind, expected_version): from airweave.schemas.source_connection import ScheduleConfig sync_repo = AsyncMock() @@ -1086,6 +1090,13 @@ 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.slack import SlackSource + + source_entry = _mock_source_entry() + source_entry.source_class_ref = {"legacy": object, "gmail": GmailSource, "slack": SlackSource}[ + source_kind + ] result = await svc.create( AsyncMock(), name="test", @@ -1093,7 +1104,7 @@ 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, schedule_config=ScheduleConfig(cron="0 6 * * *"), run_immediately=False, ctx=_mock_ctx(), @@ -1102,6 +1113,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() # --------------------------------------------------------------------------- diff --git a/backend/scripts/REPROJECTION.md b/backend/scripts/REPROJECTION.md index 091eb63e3c..e16513b180 100644 --- a/backend/scripts/REPROJECTION.md +++ b/backend/scripts/REPROJECTION.md @@ -40,11 +40,10 @@ 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 Sync creation still defaults to version 1 in this isolated implementation. -Until integrated initialization is changed, newly captured sources also need the -explicit upgrade. The next coherent creation change is in the existing sync -creation transaction (`domains/syncs/service.py::_create_sync_records`): initialize -new canonical sources to the current projection contract after source capability -resolution, without resetting any existing Sync version. That change is deferred -while the independent capture engine/live lifecycle is active; it must be finished -before presenting new-source date/type search as automatically ready. +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. From 45f0ef5c88d40faa95ce2d2ce38a5dbc189b9697 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:30:02 -0700 Subject: [PATCH 059/266] Include document titles in default keyword retrieval --- vespa/app/schemas/base_entity.sd | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vespa/app/schemas/base_entity.sd b/vespa/app/schemas/base_entity.sd index e3c55838ed..5d3b3083d6 100644 --- a/vespa/app/schemas/base_entity.sd +++ b/vespa/app/schemas/base_entity.sd @@ -34,7 +34,7 @@ schema base_entity { } field name type string { - indexing: attribute | summary + indexing: index | attribute | summary attribute: fast-search } @@ -216,7 +216,7 @@ schema base_entity { # Default fieldset used by userInput() / unspecified text search fieldset default { - fields: textual_representation, payload + fields: name, textual_representation, payload } # Note: We use Vespa's default summary (no explicit document-summary) From 2e30c5b3a1b1792e6b2b870d183ccd226dfd5e5f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:35:34 -0700 Subject: [PATCH 060/266] Make unknown-date and wrong-type fixtures independently match date bounds --- backend/tests/integration/test_canonical_search_prefilters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/tests/integration/test_canonical_search_prefilters.py b/backend/tests/integration/test_canonical_search_prefilters.py index 9e212f64a7..40c47faa95 100644 --- a/backend/tests/integration/test_canonical_search_prefilters.py +++ b/backend/tests/integration/test_canonical_search_prefilters.py @@ -56,7 +56,7 @@ async def test_prefilter_excludes_over_two_hundred_wrong_dates_before_top_matche else "event", "airweave_system_metadata_source_created_known": int(known), "airweave_system_metadata_source_created_us": epoch_microseconds(instant) - if correct + if number >= 201 else 0, "dense_embedding": {"values": vector}, "sparse_embedding": {"cells": {"1": 1.0}}, From fc8e91c99618ab673abe208a87b7597556e4d5bf Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:39:13 -0700 Subject: [PATCH 061/266] Exercise canonical source captures and sync creation in correctness checks --- Makefile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 0a26ec1be2..e0d7f319a8 100644 --- a/Makefile +++ b/Makefile @@ -32,10 +32,10 @@ setup: test-store: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/syncs/tests test-capture: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_slack_capture.py tests/unit/platform/sources/test_wispr_capture.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search From dc036d902abd40d5f674856fc7ce4b3552d2977e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 06:27:30 -0700 Subject: [PATCH 062/266] Attest canonical record forest visibility across ancestor changes --- .../domains/entities/canonical/store.py | 176 ++++++++-- .../entities/canonical/tests/conftest.py | 49 ++- .../canonical/tests/test_projection_gc.py | 7 + .../canonical/tests/test_record_forest.py | 319 ++++++++++++++++++ backend/airweave/models/entity.py | 32 +- .../versions/0006_record_visibility.py | 58 ++++ 6 files changed, 617 insertions(+), 24 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_record_forest.py create mode 100644 backend/alembic/versions/0006_record_visibility.py diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index ca41549852..d4d6da1585 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -5,7 +5,8 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import and_, exists, or_, select, update +from sqlalchemy import and_, any_, exists, func, literal, not_, or_, select, update +from sqlalchemy.dialects.postgresql import array from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import aliased @@ -103,19 +104,79 @@ def source_record(entity: Entity) -> SourceRecord: ) +def ancestor_chain(subject=Entity): + """Finite scoped ancestor walk; repeated UUIDs stop corrupt record cycles.""" + parent = aliased(Entity) + available = and_( + parent.record_revision > 0, + parent.deleted_at.is_(None), + or_( + parent.removal_reason.is_(None), + parent.removal_reason.not_in(("scope_removed", "access_revoked")), + ), + ) + chain = ( + select( + parent.id, + parent.organization_id, + parent.sync_id, + parent.parent_record_type, + parent.parent_native_id, + parent.parent_container_id, + parent.parent_visibility_epoch, + parent.removal_reason, + array([parent.id]).label("path"), + literal(False).label("cycle"), + and_(available, parent.visibility_epoch == subject.parent_visibility_epoch).label( + "valid" + ), + ) + .where( + parent.organization_id == subject.organization_id, + parent.sync_id == subject.sync_id, + parent.entity_definition_short_name == subject.parent_record_type, + parent.native_id == subject.parent_native_id, + parent.container_id.is_not_distinct_from(subject.parent_container_id), + ) + .correlate(subject) + .cte(recursive=True, nesting=True) + ) + return chain.union_all( + select( + parent.id, + parent.organization_id, + parent.sync_id, + parent.parent_record_type, + parent.parent_native_id, + parent.parent_container_id, + parent.parent_visibility_epoch, + parent.removal_reason, + chain.c.path + array([parent.id]), + (parent.id == any_(chain.c.path)).label("cycle"), + and_( + chain.c.valid, available, parent.visibility_epoch == chain.c.parent_visibility_epoch + ), + ) + .join( + chain, + and_( + parent.organization_id == chain.c.organization_id, + parent.sync_id == chain.c.sync_id, + parent.entity_definition_short_name == chain.c.parent_record_type, + parent.native_id == chain.c.parent_native_id, + parent.container_id.is_not_distinct_from(chain.c.parent_container_id), + ), + ) + .where(not_(chain.c.cycle)) + ) + + def active_parent_exists(): - """Root container visibility; caller still supplies tenant/sync authorization.""" - parent = aliased(Entity, name="canonical_parent") + """All parent attestations must reach an available root without a record cycle.""" + chain = ancestor_chain() return exists( - select(parent.id).where( - parent.organization_id == Entity.organization_id, - parent.sync_id == Entity.sync_id, - parent.entity_definition_short_name == Entity.parent_record_type, - parent.native_id == Entity.parent_native_id, - parent.container_id.is_not_distinct_from(Entity.parent_container_id), - parent.record_revision > 0, - parent.deleted_at.is_(None), - parent.parent_record_type.is_(None), + select(chain.c.id).where( + chain.c.parent_record_type.is_(None), chain.c.valid, not_(chain.c.cycle) ) ).correlate(Entity) @@ -236,6 +297,45 @@ async def capture(self, db: AsyncSession, batch: CaptureBatch) -> CaptureResult: sync = await self._fenced_sync(db, batch.fence) return await self._capture_locked(db, sync, batch) + async def _parent_attestation( + self, db: AsyncSession, sync: Sync, observation: CaptureRecord, entity: Entity | None + ) -> int | None: + """Resolve provider parent identity only inside the source writer transaction.""" + if observation.parent is None: + return None + parent_identity = observation.parent + if observation.identity == parent_identity: + raise CanonicalStoreError("A record cannot parent itself") + parent = await db.scalar( + select(Entity).where( + Entity.organization_id == sync.organization_id, + Entity.sync_id == sync.id, + Entity.entity_definition_short_name == parent_identity.record_type, + Entity.entity_id == parent_identity.entity_key, + Entity.record_revision > 0, + ) + ) + if observation.kind == "delete" and parent is None: + return None + if parent is None: + raise CanonicalStoreError("Capture parent must exist before its child") + if entity is not None: + chain = ancestor_chain() + contains_child = exists(select(chain.c.id).where(chain.c.id == entity.id)).correlate( + Entity + ) + cycles = await db.scalar(select(contains_child).where(Entity.id == parent.id)) + if cycles: + raise CanonicalStoreError("Record parent would create a cycle") + if observation.kind == "delete" and entity is not None: + return entity.parent_visibility_epoch + available = await db.scalar(select(content_is_available()).where(Entity.id == parent.id)) + if parent.deleted_at is not None or not available: + if observation.kind == "delete": + return None + raise CanonicalStoreError("Capture parent is unavailable") + return parent.visibility_epoch + async def _capture_locked( self, db: AsyncSession, sync: Sync, batch: CaptureBatch, *, seen_id: UUID | None = None ) -> CaptureResult: @@ -255,8 +355,35 @@ async def _capture_locked( .execution_options(populate_existing=True) ) ).scalar_one_or_none() + parent_epoch = await self._parent_attestation(db, sync, observation, entity) fingerprint = capture_fingerprint(observation) - if entity is not None and entity.capture_hash == fingerprint: + was_available = ( + (await db.scalar(select(content_is_available()).where(Entity.id == entity.id))) + if entity is not None + else True + ) + was_deleted = entity is not None and entity.deleted_at is not None + parent_changed = entity is not None and ( + (entity.parent_record_type, entity.parent_native_id, entity.parent_container_id) + != ( + ( + observation.parent.record_type, + observation.parent.native_id, + observation.parent.container_id, + ) + if observation.parent + else (None, None, None) + ) + ) + attestation_changed = ( + entity is not None and entity.parent_visibility_epoch != parent_epoch + ) + if ( + entity is not None + and entity.capture_hash == fingerprint + and not attestation_changed + and not (observation.kind == "upsert" and not was_available) + ): entity.last_seen_run_id = seen_id or batch.fence.attempt_id entity.observed_at = observation.observed_at unchanged += 1 @@ -269,8 +396,14 @@ async def _capture_locked( entity_id=identity.entity_key, entity_definition_short_name=identity.record_type, record_revision=0, + visibility_epoch=1, ) db.add(entity) + if observation.kind == "upsert" and ( + was_deleted or not was_available or parent_changed + ): + entity.visibility_epoch += 1 + entity.parent_visibility_epoch = parent_epoch entity.sync_job_id = batch.fence.job_id entity.native_id = identity.native_id entity.container_id = identity.container_id @@ -464,16 +597,15 @@ async def reconcile_parents( observations = [] for entity in rows[:limit]: record = source_record(entity) - reason = await db.scalar( - select(Entity.removal_reason).where( - Entity.organization_id == fence.organization_id, - Entity.sync_id == fence.sync_id, - Entity.entity_definition_short_name == record.parent.record_type, - Entity.native_id == record.parent.native_id, - Entity.container_id.is_not_distinct_from(record.parent.container_id), - Entity.record_revision > 0, - ) + chain = ancestor_chain() + nearest_reason = ( + select(chain.c.removal_reason) + .where(chain.c.removal_reason.in_(("access_revoked", "scope_removed"))) + .order_by(func.cardinality(chain.c.path)) + .limit(1) + .scalar_subquery() ) + reason = await db.scalar(select(nearest_reason).where(Entity.id == entity.id)) observations.append( CaptureRecord( identity=record.identity, diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index a777871007..bf5a54b020 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -87,6 +87,50 @@ def seed_legacy(connection): ) +def seed_flat_visibility(connection): + """Canonical flat states from schema0005, including already withdrawn content.""" + import json + + from sqlalchemy import MetaData, Table + + table = Table("entity", MetaData(), autoload_with=connection) + existing = connection.execute( + text("SELECT organization_id,sync_id,sync_job_id FROM entity LIMIT 1") + ).one() + now = datetime.now(timezone.utc) + for native, parent, reason in ( + ("active-root", None, None), + ("withdrawn-root", None, "access_revoked"), + ("active-child", "active-root", None), + ("withdrawn-child", "active-root", "scope_removed"), + ("hidden-child", "withdrawn-root", None), + ("orphan", "missing", None), + ): + connection.execute( + table.insert().values( + id=uuid4(), + organization_id=existing.organization_id, + sync_id=existing.sync_id, + sync_job_id=existing.sync_job_id, + entity_id=json.dumps([None, native], separators=(",", ":")), + entity_definition_short_name="block", + native_id=native, + parent_record_type="block" if parent else None, + parent_native_id=parent, + source_payload={"native": native}, + record_revision=1, + hash="unchanged", + capture_hash="unchanged", + completeness="complete", + observed_at=now, + deleted_at=now if reason else None, + removal_reason=reason, + created_at=now.replace(tzinfo=None), + modified_at=now.replace(tzinfo=None), + ) + ) + + @pytest.fixture async def database(request): """Every test gets its own schema; no global application DSN is consulted.""" @@ -101,13 +145,16 @@ async def database(request): try: async with engine.begin() as connection: await connection.run_sync(migrate, "0000_baseline.py") - if getattr(request, "param", None) == "legacy": + if getattr(request, "param", None) in ("legacy", "forest_upgrade"): await connection.run_sync(seed_legacy) await connection.run_sync(migrate, "0001_canonical_records.py") await connection.run_sync(migrate, "0002_projection_publication.py") await connection.run_sync(migrate, "0003_mail_thread_index.py") await connection.run_sync(migrate, "0004_projection_generation.py") await connection.run_sync(migrate, "0005_capture_scan.py") + if getattr(request, "param", None) == "forest_upgrade": + await connection.run_sync(seed_flat_visibility) + await connection.run_sync(migrate, "0006_record_visibility.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py index 73a635ce87..4675090f82 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py @@ -218,6 +218,13 @@ async def test_hidden_parent_retirement_restoration_requires_fresh_generation(da 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 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..644ba9f585 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py @@ -0,0 +1,319 @@ +"""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}) + 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) + 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"]} + ) + ) + + +async def test_nested_withdrawal_gates_original_blob_history_search_and_publication( + database, source +): + 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") + 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"}), + ) + 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 + assert not await query.queries.list_records( + db, fence.organization_id, fence.sync_id, RecordFilters(), after_id=None, limit=100 + ) + history = await service.store.changes(db, fence.organization_id, fence.sync_id) + assert all(item.record.content_access == "unavailable" for item in history.changes) + 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) diff --git a/backend/airweave/models/entity.py b/backend/airweave/models/entity.py index b22626bf23..2f3682a01f 100644 --- a/backend/airweave/models/entity.py +++ b/backend/airweave/models/entity.py @@ -4,7 +4,16 @@ from typing import TYPE_CHECKING, Optional from uuid import UUID -from sqlalchemy import BigInteger, DateTime, ForeignKey, Index, String, UniqueConstraint, text +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 @@ -47,6 +56,8 @@ class Entity(OrganizationBase): 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") @@ -80,6 +91,25 @@ 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" ), 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") From d433a8eb742d6b4ed9f7e6cedc7359fb09ca8e3a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:10:07 -0700 Subject: [PATCH 063/266] Support nested capture scopes with exact parent ownership --- .../entities/canonical/cycle_models.py | 104 +++++- .../domains/entities/canonical/cycle_store.py | 267 +++++++++----- .../domains/entities/canonical/page_source.py | 18 +- .../domains/entities/canonical/requests.py | 28 +- .../domains/entities/canonical/scan_models.py | 2 + .../domains/entities/canonical/scan_store.py | 282 ++++++++++----- .../domains/entities/canonical/service.py | 51 ++- .../domains/entities/canonical/store.py | 32 +- .../entities/canonical/tests/conftest.py | 126 ++++++- .../entities/canonical/tests/test_cycles.py | 3 +- .../canonical/tests/test_forest_scans.py | 342 ++++++++++++++++++ .../canonical/tests/test_projection_gc.py | 2 +- .../canonical/tests/test_record_forest.py | 9 +- .../canonical/tests/test_slack_recovery.py | 47 ++- .../sync_pipeline/canonical_capture.py | 38 +- .../domains/sync_pipeline/canonical_scan.py | 148 ++++---- backend/airweave/models/capture_scan.py | 24 +- backend/airweave/platform/sources/linear.py | 23 +- backend/airweave/platform/sources/slack.py | 23 +- .../alembic/versions/0007_scan_scope_owner.py | 171 +++++++++ backend/tests/live/canonical_capture.py | 2 + backend/tests/live/capture_resume.py | 13 +- backend/tests/live/provider_lifecycle.py | 13 +- .../platform/sources/test_linear_capture.py | 15 +- .../platform/sources/test_slack_capture.py | 13 +- 25 files changed, 1442 insertions(+), 354 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_forest_scans.py create mode 100644 backend/alembic/versions/0007_scan_scope_owner.py diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index b2916f7f93..2353ba1627 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -5,28 +5,90 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator -from airweave.domains.entities.canonical.requests import WriterFence +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import RecordKind, WriterFence CYCLE_KEY = "canonical_cycle" class CycleConfiguration(BaseModel): - """Whole-scope topology; child containers are visible root native IDs.""" + """Immutable copy of the source's allowed record-parent relationships.""" model_config = ConfigDict(extra="forbid", frozen=True) fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") - root_record_type: str = Field(min_length=1, max_length=200) - child_record_types: tuple[str, ...] = Field(default=(), max_length=20) + parents: dict[RecordKind, tuple[RecordKind | None, ...]] + + @classmethod + def from_source( + cls, + *, + fingerprint: str, + record_types: tuple[str, ...], + container_parents: dict[str, str | tuple[str | None, ...]], + ) -> "CycleConfiguration": + """Snapshot the existing source declaration; no second topology registry.""" + if len(set(record_types)) != len(record_types): + raise ValueError("Duplicate source record kind") + if set(container_parents) - set(record_types): + raise ValueError("Parent relationship has an undeclared source record kind") + parents = {} + for kind in record_types: + declared = container_parents.get(kind, (None,)) + parents[kind] = (declared,) if isinstance(declared, str) else declared + return cls(fingerprint=fingerprint, parents=parents) + + @model_validator(mode="before") + @classmethod + def upgrade_flat(cls, value): + """Read schema1 flat declarations without retaining a competing topology.""" + if isinstance(value, dict) and "root_record_type" in value: + upgraded = dict(value) + root = upgraded.pop("root_record_type") + children = upgraded.pop("child_record_types", ()) + if not isinstance(children, (list, tuple)): + raise ValueError("Legacy child kinds must be an array") + if "parents" in upgraded or root in children or len(set(children)) != len(children): + raise ValueError("Ambiguous capture topology") + upgraded["parents"] = {root: (None,), **{child: (root,) for child in children}} + return upgraded + return value @model_validator(mode="after") - def distinct_types(self) -> "CycleConfiguration": - """Keep scope discovery bounded and unambiguous.""" - if len(set(self.child_record_types)) != len(self.child_record_types): - raise ValueError("Child record types must be unique") - if any(not item or len(item) > 200 for item in self.child_record_types): - raise ValueError("Invalid child record type") + def reachable_types(self) -> "CycleConfiguration": + """Type cycles are valid; every declared kind must be reachable from a root.""" + if not self.parents or any(not kind or not values for kind, values in self.parents.items()): + raise ValueError("Capture topology must declare kinds and their parents") + if any(len(set(values)) != len(values) for values in self.parents.values()): + raise ValueError("Duplicate allowed parent kind") + if any( + parent is not None and parent not in self.parents + for values in self.parents.values() + for parent in values + ): + raise ValueError("Unknown parent kind") + reachable = {kind for kind, values in self.parents.items() if None in values} + while True: + expanded = reachable | { + kind + for kind, values in self.parents.items() + if any(parent in reachable for parent in values) + } + if expanded == reachable: + break + reachable = expanded + if reachable != set(self.parents): + raise ValueError("Capture kinds must be reachable from a root") return self + @property + def root_record_types(self) -> tuple[str, ...]: + """Kinds permitting records without a parent.""" + return tuple(kind for kind, parents in self.parents.items() if None in parents) + + def children_of(self, kind: str) -> tuple[str, ...]: + """Scopes required for a visible record of this kind.""" + return tuple(child for child, parents in self.parents.items() if kind in parents) + class CycleVersion(BaseModel): """Exact checkpoint version observed by the caller.""" @@ -40,13 +102,22 @@ class CaptureCycle(BaseModel): """Completion stays durable until an explicit next-cycle CAS.""" model_config = ConfigDict(extra="forbid", frozen=True) - schema_version: Literal[1] = 1 + schema_version: Literal[2] = 2 version: CycleVersion configuration: CycleConfiguration phase: Literal["active", "complete"] = "active" - root_writer_attempt_id: UUID | None = None completed_job_id: UUID | None = None + @model_validator(mode="before") + @classmethod + def upgrade_version(cls, value): + """Flat persisted cycles normalize through the same configuration boundary.""" + if isinstance(value, dict) and value.get("schema_version") == 1: + upgraded = {**value, "schema_version": 2} + upgraded.pop("root_writer_attempt_id", None) + return upgraded + return value + class BeginCycle(BaseModel): """Resume an active cycle or CAS-advance its completed predecessor.""" @@ -74,9 +145,10 @@ class RestartCycle(BaseModel): configuration: CycleConfiguration -class CycleRoot(BaseModel): - """Bounded operational scope discovery, without copying provider payloads.""" +class ScopeWork(BaseModel): + """An eligible incomplete inventory and its exact captured owner, never a durable queue.""" model_config = ConfigDict(extra="forbid", frozen=True) - id: UUID - native_id: str + record_type: str + parent: SourceRecord | None = None + parent_visibility_epoch: int | None = None diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 472adb483f..763738504f 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -3,8 +3,9 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import exists, select +from sqlalchemy import and_, exists, select from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import aliased from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint from airweave.domains.entities.canonical.cycle_models import ( @@ -12,12 +13,17 @@ BeginCycle, CaptureCycle, CompleteCycle, - CycleRoot, CycleVersion, RestartCycle, + ScopeWork, +) +from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity, WriterFence +from airweave.domains.entities.canonical.store import ( + CanonicalStoreError, + ancestor_chain, + content_is_available, + source_record, ) -from airweave.domains.entities.canonical.requests import CompletedScope, WriterFence -from airweave.domains.entities.canonical.store import CanonicalStoreError, content_is_available from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity from airweave.models.sync import Sync @@ -101,54 +107,188 @@ async def attest_cycle( return cursor, state -async def root_ready(db: AsyncSession, fence: WriterFence, state: CaptureCycle) -> bool: - """A retry must refresh membership before touching or finishing child scopes.""" - root = await db.scalar( - select(CaptureScan).where( +def inventory_complete(subject, fence: WriterFence, state: CaptureCycle): + """Exact inventory that captured a parent must be fresh for this writer attempt.""" + owner = aliased(Entity) + parent_id = ( + select(owner.id) + .where( + owner.organization_id == subject.organization_id, + owner.sync_id == subject.sync_id, + owner.entity_definition_short_name == subject.parent_record_type, + owner.native_id == subject.parent_native_id, + owner.container_id.is_not_distinct_from(subject.parent_container_id), + ) + .correlate(subject) + .scalar_subquery() + ) + return exists( + select(CaptureScan.id).where( CaptureScan.organization_id == fence.organization_id, CaptureScan.sync_id == fence.sync_id, - CaptureScan.record_type == state.configuration.root_record_type, - CaptureScan.container_id.is_(None), CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.record_type == subject.entity_definition_short_name, + CaptureScan.container_id.is_not_distinct_from(subject.container_id), + CaptureScan.parent_record_id.is_not_distinct_from(parent_id), + CaptureScan.parent_visibility_epoch.is_not_distinct_from( + subject.parent_visibility_epoch + ), + CaptureScan.membership_attempt_id == fence.attempt_id, CaptureScan.phase == "complete", ) ) - return root is not None and ( - not state.configuration.child_record_types - or state.root_writer_attempt_id == fence.attempt_id + + +def membership_ready(fence: WriterFence, state: CaptureCycle): + """Current parent and every ancestor need completed current-attempt inventories.""" + chain = ancestor_chain() + ancestor = aliased(Entity) + stale_ancestor = exists( + select(chain.c.id) + .join(ancestor, ancestor.id == chain.c.id) + .where(~inventory_complete(ancestor, fence, state)) ) + return and_(inventory_complete(Entity, fence, state), ~stale_ancestor) -async def attest_scope( +async def root_ready(db: AsyncSession, fence: WriterFence, state: CaptureCycle) -> bool: + """Every declared root inventory completes before descendant work.""" + for kind in state.configuration.root_record_types: + predicates = [ + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.record_type == kind, + CaptureScan.container_id.is_(None), + CaptureScan.parent_record_id.is_(None), + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.phase == "complete", + ] + if state.configuration.children_of(kind): + predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) + if await db.scalar(select(CaptureScan.id).where(*predicates).limit(1)) is None: + return False + return True + + +async def scope_owner( db: AsyncSession, fence: WriterFence, state: CaptureCycle, scope: CompletedScope -) -> bool: - """Only whole root or visible declared child scopes belong to this cycle.""" - if scope.record_type == state.configuration.root_record_type and scope.container_id is None: - return True - if ( - scope.record_type not in state.configuration.child_record_types - or scope.container_id is None - ): - raise CycleConflict("Scope is not declared by the active cycle") - if not await root_ready(db, fence, state): - raise CycleConflict("Refresh and complete root membership before child work") +) -> Entity | None: + """Resolve full parent identity; only old unambiguous flat requests are normalized.""" + allowed = state.configuration.parents.get(scope.record_type) + if allowed is None: + raise CycleConflict("Scope kind is not declared by the active cycle") + identity = scope.parent + if identity is None: + if scope.container_id is None and None in allowed: + return None + flat = tuple(parent for parent in allowed if parent is not None) + if ( + len(flat) != 1 + or state.configuration.parents[flat[0]] != (None,) + or scope.container_id is None + ): + raise CycleConflict("Nested scope requires its exact parent identity") + identity = RecordIdentity(record_type=flat[0], native_id=scope.container_id) + if identity.record_type not in allowed: + raise CycleConflict("Scope parent kind is not declared by the active cycle") parent = await db.scalar( - select(Entity).where( + select(Entity) + .where( Entity.organization_id == fence.organization_id, Entity.sync_id == fence.sync_id, - Entity.entity_definition_short_name == state.configuration.root_record_type, - Entity.container_id.is_(None), - Entity.native_id == scope.container_id, + Entity.entity_definition_short_name == identity.record_type, + Entity.entity_id == identity.entity_key, Entity.record_revision > 0, - Entity.deleted_at.is_(None), - content_is_available(), ) + .execution_options(populate_existing=True) ) if parent is None: + raise CycleConflict("Scope parent is missing; refresh membership before resuming") + return parent + + +async def attest_scope( + db: AsyncSession, fence: WriterFence, state: CaptureCycle, scope: CompletedScope +) -> bool: + """A scope's full owner chain must be visible and freshly enumerated.""" + parent = await scope_owner(db, fence, state, scope) + if parent is None: + return True + if ( + await db.scalar(select(Entity.id).where(Entity.id == parent.id, content_is_available())) + is None + ): raise CycleConflict("Child scope no longer has a visible parent") + ready = await db.scalar( + select(Entity.id).where( + Entity.id == parent.id, content_is_available(), membership_ready(fence, state) + ) + ) + if ready is None or not await root_ready(db, fence, state): + raise CycleConflict("Refresh and complete ancestor membership before child work") return False +def child_scope_complete(fence: WriterFence, state: CaptureCycle, record_type: str): + """A completed child scope belongs to this exact owner generation.""" + predicates = [ + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.record_type == record_type, + CaptureScan.parent_record_id == Entity.id, + CaptureScan.parent_visibility_epoch == Entity.visibility_epoch, + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.phase == "complete", + ] + if state.configuration.children_of(record_type): + predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) + return exists(select(CaptureScan.id).where(*predicates)) + + +async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) -> ScopeWork | None: + """Reevaluate SQL frontier each time; newly discovered earlier UUIDs cannot be skipped.""" + _, state = await attest_cycle(db, fence, cycle_id) + for kind in state.configuration.root_record_types: + predicates = [ + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.record_type == kind, + CaptureScan.container_id.is_(None), + CaptureScan.parent_record_id.is_(None), + CaptureScan.cycle_id == cycle_id, + CaptureScan.phase == "complete", + ] + if state.configuration.children_of(kind): + predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) + if await db.scalar(select(CaptureScan.id).where(*predicates).limit(1)) is None: + return ScopeWork(record_type=kind) + for child_type, allowed in state.configuration.parents.items(): + parent_types = tuple(kind for kind in allowed if kind is not None) + if not parent_types: + continue + parent = await db.scalar( + select(Entity) + .where( + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name.in_(parent_types), + content_is_available(), + membership_ready(fence, state), + ~child_scope_complete(fence, state, child_type), + ) + .order_by(Entity.id) + .limit(1) + .execution_options(populate_existing=True) + ) + if parent is not None: + return ScopeWork( + record_type=child_type, + parent=source_record(parent), + parent_visibility_epoch=parent.visibility_epoch, + ) + return None + + async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) -> CaptureCycle: """Publish the checkpoint only when SQL proves all currently visible scopes complete.""" cursor, state = await attest_cycle(db, request.fence, request.expected.cycle_id) @@ -156,28 +296,20 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - raise CycleConflict("Cycle changed before completion") if not await root_ready(db, request.fence, state): raise CycleConflict("Root enumeration is incomplete or stale") - for child_type in state.configuration.child_record_types: - complete_child = exists( - select(CaptureScan.id).where( - CaptureScan.organization_id == request.fence.organization_id, - CaptureScan.sync_id == request.fence.sync_id, - CaptureScan.record_type == child_type, - CaptureScan.container_id == Entity.native_id, - CaptureScan.cycle_id == state.version.cycle_id, - CaptureScan.phase == "complete", - ) - ) + for child_type, parents in state.configuration.parents.items(): + parent_types = tuple(kind for kind in parents if kind is not None) + if not parent_types: + continue missing = await db.scalar( select(Entity.id) .where( Entity.organization_id == request.fence.organization_id, Entity.sync_id == request.fence.sync_id, - Entity.entity_definition_short_name == state.configuration.root_record_type, - Entity.container_id.is_(None), + Entity.entity_definition_short_name.in_(parent_types), Entity.record_revision > 0, Entity.deleted_at.is_(None), content_is_available(), - ~complete_child, + ~child_scope_complete(request.fence, state, child_type), ) .limit(1) ) @@ -216,48 +348,3 @@ async def restart_cycle(db: AsyncSession, request: RestartCycle) -> CaptureCycle persist(cursor, state) await db.flush() return state - - -async def list_cycle_roots( - db: AsyncSession, - fence: WriterFence, - cycle_id: UUID, - *, - after: UUID | None = None, - missing: bool = False, -) -> tuple[CycleRoot, ...]: - """Read up to 100 roots under the fence; absence checks require the final root page.""" - _, state = await attest_cycle(db, fence, cycle_id) - predicates = [ - Entity.organization_id == fence.organization_id, - Entity.sync_id == fence.sync_id, - Entity.entity_definition_short_name == state.configuration.root_record_type, - Entity.container_id.is_(None), - Entity.record_revision > 0, - Entity.deleted_at.is_(None), - content_is_available(), - ] - if after is not None: - predicates.append(Entity.id > after) - if missing: - root = await db.scalar( - select(CaptureScan).where( - CaptureScan.organization_id == fence.organization_id, - CaptureScan.sync_id == fence.sync_id, - CaptureScan.record_type == state.configuration.root_record_type, - CaptureScan.container_id.is_(None), - CaptureScan.cycle_id == cycle_id, - CaptureScan.phase == "reconciling", - ) - ) - if root is None or state.root_writer_attempt_id != fence.attempt_id: - raise CycleConflict("Root absence checks require completed current-attempt collection") - predicates.append(Entity.last_seen_run_id.is_distinct_from(root.sweep_id)) - elif not await root_ready(db, fence, state): - raise CycleConflict("Root membership must complete before child discovery") - rows = ( - await db.execute( - select(Entity.id, Entity.native_id).where(*predicates).order_by(Entity.id).limit(100) - ) - ).all() - return tuple(CycleRoot(id=row.id, native_id=row.native_id) for row in rows) diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index 052529c4a9..36252787c6 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -5,6 +5,7 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from airweave.domains.entities.canonical.cycle_models import CycleConfiguration +from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.requests import ( CaptureRecord, CompletedScope, @@ -46,15 +47,24 @@ 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] + canonical_container_parents: dict[str, str | tuple[str | None, ...]] capture_cycle_configuration: CycleConfiguration async def capture_page( - self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService + self, + scope: CompletedScope, + continuation: ScanContinuation, + *, + files: FileService, + parent: SourceRecord | None = None, ) -> CapturePage: """Fetch one page without advancing durable state or performing writes.""" ... - async def confirm_root_absent(self, native_id: str) -> None: - """Raise unless an omitted previously captured root is confirmed inaccessible.""" + 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.""" ... diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 4fe249be39..f4c3b8c79c 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -1,19 +1,21 @@ """Validated source capture commands. Provider payloads remain opaque JSON.""" +import hashlib import json -from typing import Literal +from typing import Annotated, Literal from uuid import UUID from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator ScopeRemovalReason = Literal["scope_removed", "access_revoked"] +RecordKind = Annotated[str, Field(min_length=1, max_length=200)] class RecordIdentity(BaseModel): """Native identity inside one sync, including provider-required container.""" model_config = ConfigDict(extra="forbid", frozen=True) - record_type: str = Field(min_length=1) + record_type: RecordKind native_id: str = Field(min_length=1) container_id: str | None = None @@ -25,6 +27,19 @@ def entity_key(self) -> str: ) +def parent_container_key(parent: RecordIdentity) -> str: + """Optional source identity policy: bounded full-parent locator, stable across rebuilds. + + New connectors choosing this policy use it for every child, including root-owned + children. Native parent identity remains in CaptureRecord.parent and native JSON. + Existing connectors may retain their audited provider-global container IDs. + """ + material = json.dumps( + parent.model_dump(mode="json"), sort_keys=True, separators=(",", ":"), ensure_ascii=False + ) + return "parent:" + hashlib.sha256(material.encode()).hexdigest() + + class BlobReference(BaseModel): """An already-written immutable blob required by this record.""" @@ -44,6 +59,10 @@ class CaptureRecord(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) identity: RecordIdentity parent: RecordIdentity | None = None + allow_reparent: bool = Field( + default=False, + description="Audited provider evidence that this same object moved; never a retry override", + ) payload: dict[str, JsonValue] payload_schema_version: int = Field(default=1, ge=1) kind: Literal["upsert", "delete"] = "upsert" @@ -88,11 +107,12 @@ class CaptureBatch(BaseModel): class CompletedScope(BaseModel): - """Evidence that one exact type/container was exhaustively enumerated.""" + """Evidence for one exact type/container and, for nested pages, its full owner.""" model_config = ConfigDict(extra="forbid", frozen=True) - record_type: str = Field(min_length=1) + record_type: RecordKind container_id: str | None = None + parent: RecordIdentity | None = None class ReconcileScope(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index 3c70ad7559..5da784983e 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -58,6 +58,8 @@ class ScanState(BaseModel): continuation: ScanContinuation started_at: AwareDatetime completed_at: AwareDatetime | None + parent_visibility_epoch: int | None = None + membership_attempt_id: UUID | None = None class CommitScanPage(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index a6d3b3d00d..6eefe3bce4 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -4,21 +4,26 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import or_, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.cycle_models import CaptureCycle, CycleVersion +from airweave.domains.entities.canonical.cycle_models import CycleVersion from airweave.domains.entities.canonical.cycle_store import ( CycleConflict, attest_cycle, attest_scope, + cursor_row, + cycle_state, persist, + scope_owner, ) +from airweave.domains.entities.canonical.models import CaptureResult, SourceRecord from airweave.domains.entities.canonical.requests import ( CaptureBatch, + CaptureRecord, CompletedScope, ReconcileScope, - RecordIdentity, + ScopeRemovalReason, WriterFence, ) from airweave.domains.entities.canonical.scan_models import ( @@ -33,7 +38,7 @@ from airweave.domains.entities.canonical.store import ( CanonicalRecordStore, CanonicalStoreError, - content_is_available, + source_record, ) from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity @@ -45,17 +50,22 @@ class ScanConflict(CanonicalStoreError): code = "scan_conflict" -def scope_key(scope: CompletedScope) -> str: +def scope_key(scope: CompletedScope, parent_id: UUID | None = None) -> str: """Distinguish null containers from empty strings without a nullable unique key.""" - return json.dumps( + key = json.dumps( [scope.record_type, scope.container_id], ensure_ascii=False, separators=(",", ":") ) + return key if parent_id is None else key + "|" + str(parent_id) -def scan_state(row: CaptureScan) -> ScanState: +def scan_state(row: CaptureScan, parent: Entity | None = None) -> ScanState: """Do not leak ORM state outside its transaction.""" return ScanState( - scope=CompletedScope(record_type=row.record_type, container_id=row.container_id), + scope=CompletedScope( + record_type=row.record_type, + container_id=row.container_id, + parent=source_record(parent).identity if parent is not None else None, + ), cycle_id=row.cycle_id, version=ScanVersion(sweep_id=row.sweep_id, revision=row.revision), phase=row.phase, @@ -63,6 +73,8 @@ def scan_state(row: CaptureScan) -> ScanState: continuation=ScanContinuation(value=row.continuation), started_at=row.started_at, completed_at=row.completed_at, + parent_visibility_epoch=row.parent_visibility_epoch, + membership_attempt_id=row.membership_attempt_id, ) @@ -76,16 +88,40 @@ def __init__(self, records: CanonicalRecordStore): async def _row( self, db: AsyncSession, fence: WriterFence, scope: CompletedScope ) -> CaptureScan | None: - return await db.scalar( - select(CaptureScan) - .where( - CaptureScan.organization_id == fence.organization_id, - CaptureScan.sync_id == fence.sync_id, - CaptureScan.scope_key == scope_key(scope), + cycle = cycle_state(await cursor_row(db, fence)) + if cycle is None: + raise CycleConflict("Scan requires a persisted cycle") + parent = await scope_owner(db, fence, cycle, scope) + predicates = [ + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + ] + if parent is None: + predicates.append(CaptureScan.scope_key == scope_key(scope)) + else: + predicates.extend( + ( + CaptureScan.parent_record_id == parent.id, + CaptureScan.record_type == scope.record_type, + ) ) + row = await db.scalar( + select(CaptureScan) + .where(*predicates) .with_for_update() .execution_options(populate_existing=True) ) + if row is not None and row.container_id != scope.container_id: + raise CycleConflict( + "Established scope identity mapping changed; explicit migration required" + ) + return row + + async def _state(self, db: AsyncSession, row: CaptureScan) -> ScanState: + parent = await db.get(Entity, row.parent_record_id) if row.parent_record_id else None + if row.parent_record_id is not None and parent is None: + raise CycleConflict("Scope parent disappeared; explicitly restart the cycle") + return scan_state(row, parent) async def read( self, db: AsyncSession, fence: WriterFence, scope: CompletedScope @@ -93,7 +129,7 @@ async def read( """Read only under the active writer fence, including after an uncertain commit.""" await self.records._fenced_sync(db, fence) row = await self._row(db, fence, scope) - return scan_state(row) if row else None + return await self._state(db, row) if row else None @staticmethod def _expect(row: CaptureScan | None, expected: ScanVersion, cycle_id: UUID) -> CaptureScan: @@ -106,30 +142,37 @@ def _expect(row: CaptureScan | None, expected: ScanVersion, cycle_id: UUID) -> C raise ScanConflict("Scan changed; reload its durable state") return row + @staticmethod + def _check_parent_epoch(row: CaptureScan, parent: Entity | None) -> None: + if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): + raise ScanConflict("Scope owner changed; restart its sweep with current CAS") + async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: """Resume unchanged scans; restart and next-cycle transitions require exact CAS.""" await self.records._fenced_sync(db, request.fence) cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) - is_root = await attest_scope(db, request.fence, cycle, request.scope) + await attest_scope(db, request.fence, cycle, request.scope) + parent = await scope_owner(db, request.fence, cycle, request.scope) + inventory = bool(cycle.configuration.children_of(request.scope.record_type)) if request.fingerprint != cycle.configuration.fingerprint: raise CycleConflict("Scan configuration differs from the active cycle") row = await self._row(db, request.fence, request.scope) if ( row is not None and row.cycle_id == request.cycle_id - and is_root - and cycle.configuration.child_record_types - and cycle.root_writer_attempt_id != request.fence.attempt_id + and inventory + and row.membership_attempt_id != request.fence.attempt_id and not request.restart ): - raise CycleConflict("Restart root membership for this writer attempt") + raise CycleConflict("Restart root or ancestor membership for this writer attempt") if row is None: if request.expected is not None or request.restart: raise ScanConflict("Cannot restart a missing scan") row = CaptureScan( organization_id=request.fence.organization_id, sync_id=request.fence.sync_id, - scope_key=scope_key(request.scope), + scope_key=scope_key(request.scope, parent.id if parent is not None else None), + parent_record_id=parent.id if parent is not None else None, record_type=request.scope.record_type, container_id=request.scope.container_id, revision=1, @@ -143,12 +186,15 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: and row.fingerprint == request.fingerprint and not request.restart ): - return scan_state(row) + self._check_parent_epoch(row, parent) + return await self._state(db, row) if request.expected is None: raise ScanConflict("Replacing a scan requires its current version") if row.cycle_id == request.cycle_id and not request.restart: raise ScanConflict("Changed scope configuration requires an explicit restart") row.revision += 1 + row.parent_visibility_epoch = parent.visibility_epoch if parent else None + row.membership_attempt_id = request.fence.attempt_id if inventory else None row.cycle_id = request.cycle_id row.sweep_id = uuid4() row.fingerprint = request.fingerprint @@ -156,33 +202,36 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: row.continuation = request.continuation.value row.started_at = datetime.now(timezone.utc) row.completed_at = None - if is_root: - cycle = cycle.model_copy( + persist( + cursor, + cycle.model_copy( update={ - "root_writer_attempt_id": request.fence.attempt_id, "version": CycleVersion( - cycle_id=cycle.version.cycle_id, revision=cycle.version.revision + 1 - ), + cycle_id=cycle.version.cycle_id, + revision=cycle.version.revision + 1, + ) } - ) - persist(cursor, cycle) + ), + ) await db.flush() - return scan_state(row) + return await self._state(db, row) async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: """Capture and advance the page as one transaction, never a partial acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) - is_root = await attest_scope(db, request.fence, cycle, request.scope) - if ( - is_root - and cycle.configuration.child_record_types - and cycle.root_writer_attempt_id != request.fence.attempt_id - ): - raise CycleConflict("Root membership belongs to an earlier writer attempt") + await attest_scope(db, request.fence, cycle, request.scope) row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) + parent = await scope_owner(db, request.fence, cycle, request.scope) + if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): + raise ScanConflict("Scope owner changed; restart from its current epoch") + if ( + cycle.configuration.children_of(row.record_type) + and row.membership_attempt_id != request.fence.attempt_id + ): + raise CycleConflict("Membership belongs to an earlier writer attempt") if row.phase != "collecting": raise ScanConflict("Scan is not collecting pages") if any( @@ -191,15 +240,9 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: for record in request.records ): raise ScanConflict("Page contains records outside its exact scope") - if not is_root: - parent = RecordIdentity( - record_type=cycle.configuration.root_record_type, - native_id=request.scope.container_id, - ) - if any(record.parent != parent for record in request.records): - raise ScanConflict("Child records must retain their declared parent identity") - if is_root and cycle.configuration.child_record_types: - await self._invalidate_revived_children(db, request, cycle) + expected_parent = source_record(parent).identity if parent is not None else None + if any(record.parent != expected_parent for record in request.records): + raise ScanConflict("Page records must retain their exact declared parent identity") captured = await self.records._capture_locked( db, sync, @@ -211,71 +254,109 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: if request.final: row.phase = "reconciling" await db.flush() - return ScanResult(state=scan_state(row), capture=captured) + return ScanResult(state=await self._state(db, row), capture=captured) - async def _invalidate_revived_children( + async def withdraw_parent( self, db: AsyncSession, - request: CommitScanPage, - cycle: CaptureCycle, - ) -> None: - """Reviving an unavailable parent requires a fresh child scan in this same transaction.""" - upserts = tuple( - record.identity.native_id for record in request.records if record.kind == "upsert" + fence: WriterFence, + parent: SourceRecord, + *, + expected_epoch: int, + removal_reason: ScopeRemovalReason, + ) -> CaptureResult: + """A failed old provider route cannot withdraw a renamed or revived current owner.""" + sync = await self.records._fenced_sync(db, fence) + current = await db.scalar( + select(Entity) + .where( + Entity.id == parent.id, + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + ) + .execution_options(populate_existing=True) ) - if not upserts: - return - revived = tuple( - ( - await db.scalars( - select(Entity.native_id).where( - Entity.organization_id == request.fence.organization_id, - Entity.sync_id == request.fence.sync_id, - Entity.entity_definition_short_name == cycle.configuration.root_record_type, - Entity.container_id.is_(None), - Entity.native_id.in_(upserts), - or_(Entity.deleted_at.is_not(None), ~content_is_available()), - ) - ) - ).all() + if ( + parent.sync_id != fence.sync_id + or current is None + or current.record_revision != parent.revision + or current.visibility_epoch != expected_epoch + or source_record(current).identity != parent.identity + ): + raise ScanConflict("Scope owner changed during provider access check; reload work") + observation = CaptureRecord( + identity=parent.identity, + parent=parent.parent, + payload=parent.payload, + payload_schema_version=parent.payload_schema_version, + completeness=parent.completeness, + content_hash=parent.content_hash, + source_created_at=parent.source_created_at, + source_updated_at=parent.source_updated_at, + blobs=parent.blobs, + kind="delete", + removal_reason=removal_reason, + observed_at=datetime.now(timezone.utc), ) - if not revived: - return - rows = ( - await db.scalars( - select(CaptureScan) - .where( - CaptureScan.organization_id == request.fence.organization_id, - CaptureScan.sync_id == request.fence.sync_id, - CaptureScan.cycle_id == request.cycle_id, - CaptureScan.record_type.in_(cycle.configuration.child_record_types), - CaptureScan.container_id.in_(revived), - ) - .with_for_update() - ) + return await self.records._capture_locked( + db, sync, CaptureBatch(fence=fence, records=(observation,)) + ) + + async def missing( + self, db: AsyncSession, fence: WriterFence, state: ScanState, *, after: UUID | None = None + ) -> tuple[SourceRecord, ...]: + """Bounded exact inventory omissions, only after current-attempt collection.""" + await self.records._fenced_sync(db, fence) + _, cycle = await attest_cycle(db, fence, state.cycle_id) + await attest_scope(db, fence, cycle, state.scope) + row = self._expect(await self._row(db, fence, state.scope), state.version, state.cycle_id) + parent = await scope_owner(db, fence, cycle, state.scope) + if ( + row.phase != "reconciling" + or row.membership_attempt_id != fence.attempt_id + or row.parent_visibility_epoch != (parent.visibility_epoch if parent else None) + ): + raise CycleConflict("Omission checks require completed current-attempt inventory") + identity = source_record(parent).identity if parent else None + predicates = [ + Entity.organization_id == fence.organization_id, + Entity.sync_id == fence.sync_id, + Entity.entity_definition_short_name == row.record_type, + Entity.container_id.is_not_distinct_from(row.container_id), + Entity.parent_record_type.is_not_distinct_from( + identity.record_type if identity else None + ), + Entity.parent_native_id.is_not_distinct_from(identity.native_id if identity else None), + Entity.parent_container_id.is_not_distinct_from( + identity.container_id if identity else None + ), + Entity.record_revision > 0, + Entity.deleted_at.is_(None), + Entity.last_seen_run_id.is_distinct_from(row.sweep_id), + ] + if after is not None: + predicates.append(Entity.id > after) + records = ( + await db.scalars(select(Entity).where(*predicates).order_by(Entity.id).limit(100)) ).all() - for child in rows: - child.sweep_id = uuid4() - child.revision += 1 - child.phase = "collecting" - child.continuation = {} - child.started_at = datetime.now(timezone.utc) - child.completed_at = None + return tuple(source_record(record) for record in records) async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: """A final page is necessary; absence completion is durable and bounded.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) - is_root = await attest_scope(db, request.fence, cycle, request.scope) - if ( - is_root - and cycle.configuration.child_record_types - and cycle.root_writer_attempt_id != request.fence.attempt_id - ): - raise CycleConflict("Root membership belongs to an earlier writer attempt") + await attest_scope(db, request.fence, cycle, request.scope) row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) + parent = await scope_owner(db, request.fence, cycle, request.scope) + if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): + raise ScanConflict("Scope owner changed; restart from its current epoch") + if ( + cycle.configuration.children_of(row.record_type) + and row.membership_attempt_id != request.fence.attempt_id + ): + raise CycleConflict("Membership belongs to an earlier writer attempt") if row.phase != "reconciling": raise ScanConflict("Only fully collected scans can reconcile absence") result = await self.records._reconcile_scope_locked( @@ -283,20 +364,21 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul sync, ReconcileScope( fence=request.fence, - scope=request.scope, + scope=(await self._state(db, row)).scope, removal_reason=( "scope_removed" - if is_root and cycle.configuration.child_record_types + if cycle.configuration.children_of(row.record_type) else request.removal_reason ), observed_at=request.observed_at, limit=request.limit, ), seen_id=row.sweep_id, + parent_scoped=True, ) row.revision += 1 if not result.has_more: row.phase = "complete" row.completed_at = datetime.now(timezone.utc) await db.flush() - return ScanResult(state=scan_state(row), capture=result.capture) + return ScanResult(state=await self._state(db, row), capture=result.capture) diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index b83ef65946..e9d434521c 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -10,23 +10,24 @@ BeginCycle, CaptureCycle, CompleteCycle, - CycleRoot, RestartCycle, + ScopeWork, ) from airweave.domains.entities.canonical.cycle_store import ( begin_cycle, complete_cycle, cursor_row, cycle_state, - list_cycle_roots, + next_scope_work, restart_cycle, ) -from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult +from airweave.domains.entities.canonical.models import CaptureResult, ReconcileResult, SourceRecord from airweave.domains.entities.canonical.requests import ( CaptureBatch, CompletedScope, ReconcileScope, RemovedScope, + ScopeRemovalReason, StartedScope, WriterFence, ) @@ -115,6 +116,36 @@ async def save_checkpoint( async with UnitOfWork(db): await self.store.save_checkpoint(db, fence, cursor_data) + async def next_scope_work( + self, db: AsyncSession, fence: WriterFence, cycle_id: UUID + ) -> ScopeWork | None: + """Read the next eligible incomplete scope under the existing writer fence.""" + async with UnitOfWork(db): + await self.store._fenced_sync(db, fence) + return await next_scope_work(db, fence, cycle_id) + + async def withdraw_scan_parent( + self, + db: AsyncSession, + fence: WriterFence, + parent: SourceRecord, + *, + expected_epoch: int, + removal_reason: ScopeRemovalReason, + ) -> CaptureResult: + """Attest exact fetched routing revision and owner epoch in the withdrawal transaction.""" + async with UnitOfWork(db): + return await self.scans.withdraw_parent( + db, fence, parent, expected_epoch=expected_epoch, removal_reason=removal_reason + ) + + async def scan_missing( + self, db: AsyncSession, fence: WriterFence, state: ScanState, *, after: UUID | None = None + ) -> tuple[SourceRecord, ...]: + """Read exact inventory omissions under the same attempt and scan revision.""" + async with UnitOfWork(db): + return await self.scans.missing(db, fence, state, after=after) + async def read_scan( self, db: AsyncSession, fence: WriterFence, scope: CompletedScope ) -> ScanState | None: @@ -160,17 +191,3 @@ async def restart_cycle(self, db: AsyncSession, request: RestartCycle) -> Captur async with UnitOfWork(db): await self.store._fenced_sync(db, request.fence) return await restart_cycle(db, request) - - async def list_cycle_roots( - self, - db: AsyncSession, - fence: WriterFence, - cycle_id: UUID, - *, - after: UUID | None = None, - missing: bool = False, - ) -> tuple[CycleRoot, ...]: - """Read a bounded scope-discovery page after verifying current writer ownership.""" - async with UnitOfWork(db): - await self.store._fenced_sync(db, fence) - return await list_cycle_roots(db, fence, cycle_id, after=after, missing=missing) diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index d4d6da1585..0034c3adbf 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -61,7 +61,7 @@ class WriterBusy(CanonicalStoreError): def capture_fingerprint(record: CaptureRecord) -> str: """Observation times never create changes; sparse tombstone payloads do.""" - material = record.model_dump(mode="json", exclude={"observed_at"}) + material = record.model_dump(mode="json", exclude={"observed_at", "allow_reparent"}) serialized = json.dumps( material, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False ) @@ -355,6 +355,12 @@ async def _capture_locked( .execution_options(populate_existing=True) ) ).scalar_one_or_none() + if entity is not None and observation.kind == "delete" and observation.parent is None: + # Native deletion feeds may contain only an object ID. Missing parent data + # is not evidence that the object moved out of its known authorization chain. + observation = observation.model_copy( + update={"parent": source_record(entity).parent} + ) parent_epoch = await self._parent_attestation(db, sync, observation, entity) fingerprint = capture_fingerprint(observation) was_available = ( @@ -375,6 +381,11 @@ async def _capture_locked( else (None, None, None) ) ) + if parent_changed and not observation.allow_reparent: + raise CanonicalStoreError( + "Existing identity belongs to a different parent; " + "explicit provider move required" + ) attestation_changed = ( entity is not None and entity.parent_visibility_epoch != parent_epoch ) @@ -459,7 +470,13 @@ async def reconcile_scope(self, db: AsyncSession, request: ReconcileScope) -> Re return await self._reconcile_scope_locked(db, sync, request) async def _reconcile_scope_locked( - self, db: AsyncSession, sync: Sync, request: ReconcileScope, *, seen_id: UUID | None = None + self, + db: AsyncSession, + sync: Sync, + request: ReconcileScope, + *, + seen_id: UUID | None = None, + parent_scoped: bool = False, ) -> ReconcileResult: """Share exact-scope reconciliation with durable whole-scope scans.""" statement = ( @@ -476,6 +493,17 @@ async def _reconcile_scope_locked( .order_by(Entity.id) .limit(request.limit + 1) ) + if parent_scoped: + parent = request.scope.parent + statement = statement.where( + Entity.parent_record_type.is_not_distinct_from( + parent.record_type if parent else None + ), + Entity.parent_native_id.is_not_distinct_from(parent.native_id if parent else None), + Entity.parent_container_id.is_not_distinct_from( + parent.container_id if parent else None + ), + ) entities = list((await db.scalars(statement)).all()) observations = tuple( CaptureRecord( diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index bf5a54b020..76369741d3 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -131,6 +131,101 @@ def seed_flat_visibility(connection): ) +def seed_flat_scan(connection, mode="scan_upgrade"): + """A schema1 active cycle with an acknowledged provider continuation before0007.""" + from sqlalchemy import MetaData, Table + + metadata = MetaData() + cursor = Table("sync_cursor", metadata, autoload_with=connection) + scans = Table("capture_scan", metadata, autoload_with=connection) + parent = connection.execute( + text("SELECT organization_id,sync_id FROM entity WHERE native_id='active-root'") + ).one() + cycle_id, sweep_id = uuid4(), uuid4() + now = datetime.now(timezone.utc) + connection.execute( + cursor.insert().values( + id=uuid4(), + organization_id=parent.organization_id, + sync_id=parent.sync_id, + created_at=now.replace(tzinfo=None), + modified_at=now.replace(tzinfo=None), + last_updated=now, + cursor_data={ + "canonical_cycle": { + "schema_version": 1, + "version": {"cycle_id": str(cycle_id), "revision": 7}, + "configuration": { + "fingerprint": "a" * 64, + "root_record_type": "block", + "child_record_types": ["leaf"], + }, + "phase": "active", + "root_writer_attempt_id": str(uuid4()), + } + }, + ) + ) + connection.execute( + scans.insert().values( + id=uuid4(), + organization_id=parent.organization_id, + sync_id=parent.sync_id, + created_at=now.replace(tzinfo=None), + modified_at=now.replace(tzinfo=None), + scope_key='["leaf","active-root"]', + record_type="leaf", + container_id="active-root", + cycle_id=cycle_id, + sweep_id=sweep_id, + revision=11, + phase="collecting", + fingerprint="a" * 64, + continuation={"cursor": "synthetic-acknowledged-page"}, + started_at=now, + ) + ) + + if mode != "scan_upgrade": + stored = connection.execute(cursor.select()).mappings().one() + changed = dict(stored["cursor_data"]) + legacy = changed["canonical_cycle"] + if mode == "scan_upgrade_children_object": + legacy["configuration"]["child_record_types"] = {"leaf": True} + elif mode == "scan_upgrade_foreign_version": + legacy["schema_version"] = 77 + elif mode == "scan_upgrade_completed": + legacy["phase"] = "complete" + elif mode == "scan_upgrade_too_many_children": + legacy["configuration"]["child_record_types"] = [ + "leaf", + *[f"kind{i}" for i in range(20)], + ] + elif mode == "scan_upgrade_default_children": + legacy["configuration"].pop("child_record_types") + connection.execute( + scans.insert().values( + id=uuid4(), + organization_id=parent.organization_id, + sync_id=parent.sync_id, + created_at=now.replace(tzinfo=None), + modified_at=now.replace(tzinfo=None), + scope_key='["block",null]', + record_type="block", + container_id=None, + cycle_id=cycle_id, + sweep_id=uuid4(), + revision=9, + phase="complete", + fingerprint="a" * 64, + continuation={}, + started_at=now, + completed_at=now, + ) + ) + connection.execute(cursor.update().values(cursor_data=changed)) + + @pytest.fixture async def database(request): """Every test gets its own schema; no global application DSN is consulted.""" @@ -145,16 +240,43 @@ async def database(request): try: async with engine.begin() as connection: await connection.run_sync(migrate, "0000_baseline.py") - if getattr(request, "param", None) in ("legacy", "forest_upgrade"): + if getattr(request, "param", None) in ( + "legacy", + "forest_upgrade", + "scan_upgrade", + "scan_upgrade_children_object", + "scan_upgrade_foreign_version", + "scan_upgrade_completed", + "scan_upgrade_default_children", + "scan_upgrade_too_many_children", + ): await connection.run_sync(seed_legacy) await connection.run_sync(migrate, "0001_canonical_records.py") await connection.run_sync(migrate, "0002_projection_publication.py") await connection.run_sync(migrate, "0003_mail_thread_index.py") await connection.run_sync(migrate, "0004_projection_generation.py") await connection.run_sync(migrate, "0005_capture_scan.py") - if getattr(request, "param", None) == "forest_upgrade": + if getattr(request, "param", None) in ( + "forest_upgrade", + "scan_upgrade", + "scan_upgrade_children_object", + "scan_upgrade_foreign_version", + "scan_upgrade_completed", + "scan_upgrade_default_children", + "scan_upgrade_too_many_children", + ): await connection.run_sync(seed_flat_visibility) await connection.run_sync(migrate, "0006_record_visibility.py") + if getattr(request, "param", None) in ( + "scan_upgrade", + "scan_upgrade_children_object", + "scan_upgrade_foreign_version", + "scan_upgrade_completed", + "scan_upgrade_default_children", + "scan_upgrade_too_many_children", + ): + await connection.run_sync(seed_flat_scan, request.param) + await connection.run_sync(migrate, "0007_scan_scope_owner.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_cycles.py b/backend/airweave/domains/entities/canonical/tests/test_cycles.py index 3b744a8b8f..cc148f2a91 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_cycles.py +++ b/backend/airweave/domains/entities/canonical/tests/test_cycles.py @@ -372,7 +372,8 @@ async def test_explicit_restart_preserves_originals_and_last_checkpoint(database async with database() as db: replacement = await service.restart_cycle(db, request) assert replacement.version.cycle_id != active.version.cycle_id - assert replacement.phase == "active" and replacement.root_writer_attempt_id is None + assert replacement.phase == "active" + assert "root_writer_attempt_id" not in replacement.model_dump() async with database() as db: with pytest.raises(CycleConflict): await service.restart_cycle(db, request) diff --git a/backend/airweave/domains/entities/canonical/tests/test_forest_scans.py b/backend/airweave/domains/entities/canonical/tests/test_forest_scans.py new file mode 100644 index 0000000000..38fa7dfd84 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_forest_scans.py @@ -0,0 +1,342 @@ +"""Real SQL frontier and identity proofs for nested source inventories.""" + +from datetime import datetime, timezone + +import pytest + +from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + parent_container_key, +) +from airweave.domains.entities.canonical.tests.test_cycles import complete, finish, page, scan + + +def test_source_owned_parent_digest_is_bounded_and_distinguishes_full_identity(): + left = RecordIdentity(record_type="block", container_id="left", native_id="same") + right = RecordIdentity(record_type="block", container_id="right", native_id="same") + assert parent_container_key(left) != parent_container_key(right) + for depth in range(96): + key = parent_container_key(left) + assert len(key) == 71 + left = RecordIdentity(record_type="block", container_id=key, native_id=str(depth)) + assert parent_container_key(left) == parent_container_key( + RecordIdentity.model_validate_json(left.model_dump_json()) + ) + + +def item(kind, native, parent=None): + return CaptureRecord( + identity=RecordIdentity( + record_type=kind, + native_id=native, + container_id=parent_container_key(parent) if parent else None, + ), + parent=parent, + payload={"id": native}, + observed_at=datetime.now(timezone.utc), + ) + + +async def work(database, service, fence, active): + async with database() as db: + return await service.next_scope_work(db, fence, active.version.cycle_id) + + +async def collect(database, service, fence, active, scope, *records): + state = await scan(database, service, fence, active, scope) + state = await page(database, service, fence, state, *records, final=True) + return await finish(database, service, fence, state) + + +async def test_nested_frontier_discovers_descendants_and_requires_every_scope(database, source): + service, fence = source + config = CycleConfiguration.from_source( + fingerprint="a" * 64, + record_types=("repository", "issue", "comment", "file"), + container_parents={"issue": "repository", "comment": "issue", "file": "repository"}, + ) + async with database() as db: + active = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + assert (await work(database, service, fence, active)).record_type == "repository" + repo = item("repository", "repo") + await collect(database, service, fence, active, CompletedScope(record_type="repository"), repo) + issue_work = await work(database, service, fence, active) + assert issue_work.record_type == "issue" and issue_work.parent.identity == repo.identity + issue = item("issue", "1", repo.identity) + issue_scope = CompletedScope( + record_type="issue", container_id=issue.identity.container_id, parent=repo.identity + ) + await collect(database, service, fence, active, issue_scope, issue) + next_item = await work(database, service, fence, active) + assert next_item.record_type == "comment" and next_item.parent.identity == issue.identity + comment = item("comment", "1", issue.identity) + await collect( + database, + service, + fence, + active, + CompletedScope( + record_type="comment", container_id=comment.identity.container_id, parent=issue.identity + ), + comment, + ) + assert (await work(database, service, fence, active)).record_type == "file" + await collect( + database, + service, + fence, + active, + CompletedScope( + record_type="file", + container_id=parent_container_key(repo.identity), + parent=repo.identity, + ), + ) + assert await work(database, service, fence, active) is None + assert (await complete(database, service, fence)).phase == "complete" + + +def test_single_source_topology_accepts_same_kind_and_upgrades_old_flat_snapshot(): + config = CycleConfiguration.from_source( + fingerprint="a" * 64, + record_types=("page", "block"), + container_parents={"block": ("page", "block")}, + ) + assert config.children_of("block") == ("block",) + assert config.root_record_types == ("page",) + old = CycleConfiguration( + fingerprint="a" * 64, root_record_type="page", child_record_types=("block",) + ) + assert old.parents == {"page": (None,), "block": ("page",)} + with pytest.raises(ValueError, match="reachable"): + CycleConfiguration(fingerprint="a" * 64, parents={"block": ("block",)}) + + +@pytest.mark.parametrize("database", ["scan_upgrade"], indirect=True) +async def test_upgrade_preserves_old_acknowledged_scan_and_cycle_restart_cas(database): + from uuid import uuid4 + + from sqlalchemy import select + + from airweave.domains.entities.canonical.cycle_models import RestartCycle + from airweave.domains.entities.canonical.cycle_store import CycleConflict + from airweave.domains.entities.canonical.service import CanonicalCaptureService + from airweave.domains.entities.canonical.store import CanonicalRecordStore + from airweave.models import SyncJob + from airweave.models.capture_scan import CaptureScan + + service = CanonicalCaptureService(CanonicalRecordStore()) + async with database() as db: + job = await db.scalar(select(SyncJob)) + job.status = "running" + await db.commit() + fence = await service.activate_writer( + db, job.organization_id, job.sync_id, job.id, attempt_id=uuid4(), attempt_number=1 + ) + async with database() as db: + cycle = await service.read_cycle(db, fence) + assert cycle.version.revision == 7 and cycle.configuration.parents == { + "block": (None,), + "leaf": ("block",), + } + scope = CompletedScope(record_type="leaf", container_id="active-root") + async with database() as db: + state = await service.read_scan(db, fence, scope) + assert state.version.revision == 11 + assert state.continuation.value == {"cursor": "synthetic-acknowledged-page"} + assert state.scope.parent == RecordIdentity(record_type="block", native_id="active-root") + async with database() as db: + restarted = await service.restart_cycle( + db, RestartCycle(fence=fence, expected=cycle.version, configuration=cycle.configuration) + ) + assert restarted.version.cycle_id != cycle.version.cycle_id + async with database() as db: + old = await db.scalar(select(CaptureScan)) + assert old.sweep_id == state.version.sweep_id and old.revision == 11 + assert old.continuation == state.continuation.value + async with database() as db: + with pytest.raises(CycleConflict): + await service.restart_cycle( + db, + RestartCycle( + fence=fence, expected=cycle.version, configuration=cycle.configuration + ), + ) + + +async def test_driver_finds_new_same_kind_descendants_behind_uuid_frontier( + database, source, monkeypatch +): + from unittest.mock import AsyncMock, MagicMock + from uuid import UUID, uuid4 + + from airweave.domains.entities.canonical.page_source import CapturePage + from airweave.domains.entities.canonical.scan_models import ScanContinuation + from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver + + service, fence = source + ids = iter([UUID(int=100), UUID(int=50), UUID(int=1)]) + monkeypatch.setattr("airweave.domains.entities.canonical.store.uuid4", lambda: next(ids)) + + class NestedSource: + canonical_record_types = ("block",) + canonical_container_parents = {"block": (None, "block")} + capture_cycle_configuration = CycleConfiguration.from_source( + fingerprint="a" * 64, + record_types=canonical_record_types, + container_parents=canonical_container_parents, + ) + + def __init__(self): + self.calls = [] + self.interrupt = True + + def child_scope(self, parent, record_type): + return CompletedScope( + record_type=record_type, + parent=parent.identity, + container_id=parent_container_key(parent.identity), + ) + + async def capture_page(self, scope, continuation, *, files, parent=None): + name = parent.identity.native_id if parent else None + self.calls.append(name) + if name == "leaf" and self.interrupt: + self.interrupt = False + raise RuntimeError("Synthetic interruption before final descendant scan") + next_name = {None: "root", "root": "child", "child": "leaf"}.get(name) + records = (item("block", next_name, scope.parent),) if next_name else () + return CapturePage(records=records, continuation=ScanContinuation(), final=True) + + async def confirm_absent(self, record): + raise AssertionError("No omitted records in this synthetic tree") + + connector = NestedSource() + driver = CanonicalScanDriver( + service, database, fence, connector, AsyncMock(), AsyncMock(), MagicMock() + ) + with pytest.raises(RuntimeError, match="Synthetic interruption"): + await driver.run() + async with database() as db: + 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, + ) + resumed = CanonicalScanDriver( + service, database, newer, connector, AsyncMock(), AsyncMock(), MagicMock() + ) + after = await resumed.run() + assert before.version.cycle_id == after.version.cycle_id + assert connector.calls == [None, "root", "child", "leaf"] * 2 + assert (await complete(database, service, newer)).phase == "complete" + + +async def test_inflight_page_rejected_after_exact_parent_epoch_changes(database, source): + from airweave.domains.entities.canonical.requests import CaptureBatch + from airweave.domains.entities.canonical.scan_store import ScanConflict + + service, fence = source + config = CycleConfiguration.from_source( + fingerprint="a" * 64, + record_types=("repository", "issue"), + container_parents={"issue": "repository"}, + ) + async with database() as db: + active = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + repo = item("repository", "repo") + await collect(database, service, fence, active, CompletedScope(record_type="repository"), repo) + issue = item("issue", "one", repo.identity) + scope = CompletedScope( + record_type="issue", parent=repo.identity, container_id=issue.identity.container_id + ) + inflight = await scan(database, service, fence, active, scope) + from airweave.domains.entities.canonical.cycle_store import CycleConflict + + with pytest.raises(CycleConflict, match="identity mapping changed"): + await scan( + database, service, fence, active, scope.model_copy(update={"container_id": "different"}) + ) + async with database() as db: + await service.capture( + db, + CaptureBatch( + fence=fence, + records=( + repo.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}), + ), + ), + ) + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(repo,))) + with pytest.raises(ScanConflict, match="owner changed"): + await page(database, service, fence, inflight, issue, final=True) + async with database() as db: + saved = await service.read_scan(db, fence, scope) + assert saved.version == inflight.version and saved.continuation == inflight.continuation + restarted = await scan( + database, service, fence, active, scope, expected=saved.version, restart=True + ) + assert restarted.parent_visibility_epoch != inflight.parent_visibility_epoch + await finish( + database, service, fence, await page(database, service, fence, restarted, issue, final=True) + ) + assert (await complete(database, service, fence)).phase == "complete" + + +def test_legacy_topology_root_collision_and_oversized_kind_are_rejected(): + with pytest.raises(ValueError, match="Ambiguous"): + CycleConfiguration( + fingerprint="a" * 64, root_record_type="block", child_record_types=("block",) + ) + with pytest.raises(ValueError): + CycleConfiguration(fingerprint="a" * 64, parents={"a" * 201: (None,)}) + + +@pytest.mark.parametrize( + "database", + [ + "scan_upgrade_children_object", + "scan_upgrade_foreign_version", + "scan_upgrade_completed", + "scan_upgrade_too_many_children", + ], + indirect=True, +) +async def test_upgrade_never_guesses_ownership_from_invalid_or_inactive_cycle(database): + from sqlalchemy import select + + from airweave.models.capture_scan import CaptureScan + + async with database() as db: + row = await db.scalar(select(CaptureScan)) + assert row.parent_record_id is None and row.parent_visibility_epoch is None + assert row.scope_key == '["leaf","active-root"]' + assert row.revision == 11 and row.continuation == {"cursor": "synthetic-acknowledged-page"} + + +@pytest.mark.parametrize("database", ["scan_upgrade_default_children"], indirect=True) +async def test_upgrade_accepts_historical_omitted_child_list_default(database): + from sqlalchemy import select + + from airweave.models.capture_scan import CaptureScan + from airweave.models.sync_cursor import SyncCursor + + async with database() as db: + root = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "block")) + cursor = await db.scalar(select(SyncCursor)) + assert ( + str(root.membership_attempt_id) + == cursor.cursor_data["canonical_cycle"]["root_writer_attempt_id"] + ) + assert root.revision == 9 and root.phase == "complete" + child = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "leaf")) + assert child.parent_record_id is None diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py index 4675090f82..634b897572 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py @@ -199,7 +199,7 @@ async def test_hidden_parent_retirement_restoration_requires_fresh_generation(da # 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)) + await capture(database, service, fence, observation(parent=parent, allow_reparent=True)) await capture( database, service, diff --git a/backend/airweave/domains/entities/canonical/tests/test_record_forest.py b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py index 644ba9f585..26b9d5f0a8 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_record_forest.py +++ b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py @@ -177,11 +177,18 @@ async def test_reparent_invalidates_descendants_and_failed_batch_rolls_back(data 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}) + 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 diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index bc747d9aaa..da78ba6ec4 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -440,7 +440,7 @@ async def test_scope_withdrawal_keeps_uncertain_reason_without_claiming_access_r instance, _, pipeline = runner(database, source, [ROOT, HISTORY]) original_page = pipeline.page_source.capture_page - async def unavailable(scope, continuation, *, files): + 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" @@ -457,3 +457,48 @@ async def unavailable(scope, continuation, *, files): 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 = 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" diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index a433ecb0fb..20e3f1581f 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -8,7 +8,11 @@ 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 +from airweave.domains.entities.canonical.cycle_models import ( + CaptureCycle, + CompleteCycle, + CycleConfiguration, +) from airweave.domains.entities.canonical.models import CaptureResult from airweave.domains.entities.canonical.page_source import CanonicalPageSource from airweave.domains.entities.canonical.requests import ( @@ -42,7 +46,7 @@ def __init__( event_bus: EventBus, record_types: tuple[str, ...], attempt: CaptureAttempt, - container_parents: dict[str, str] | None = None, + container_parents: dict[str, str | tuple[str | None, ...]] | None = None, page_source: CanonicalPageSource | None = None, files: FileService | None = None, ): @@ -54,23 +58,30 @@ def __init__( self._event_bus = event_bus self._record_types = frozenset(record_types) self._container_parents = dict(container_parents or {}) - for child, parent in self._container_parents.items(): - if child not in self._record_types 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("Only root-container visibility relationships are supported") 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 - if config.root_record_type not in self._record_types or any( - self._container_parents.get(child) != config.root_record_type - for child in config.child_record_types - ): + declared = CycleConfiguration.from_source( + fingerprint=config.fingerprint, + record_types=record_types, + container_parents=self._container_parents, + ) + 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() @@ -116,7 +127,6 @@ async def progress(result: CaptureResult, records: tuple[CaptureRecord, ...]) -> self._writer(), self.page_source, progress, - self._with_parent, check_limits, self.files, ).run() diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index acf9f6acb4..8bf38f795e 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -9,7 +9,7 @@ from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle from airweave.domains.entities.canonical.cycle_store import CycleConflict -from airweave.domains.entities.canonical.models import CaptureResult +from airweave.domains.entities.canonical.models import CaptureResult, SourceRecord from airweave.domains.entities.canonical.page_source import ( CanonicalPageSource, InvalidScanContinuation, @@ -17,11 +17,8 @@ ScopeRemovalReason, ) from airweave.domains.entities.canonical.requests import ( - CaptureBatch, CaptureRecord, CompletedScope, - RecordIdentity, - RemovedScope, WriterFence, ) from airweave.domains.entities.canonical.scan_models import ( @@ -44,14 +41,13 @@ def __init__( fence: WriterFence, source: CanonicalPageSource, progress: Callable[[CaptureResult, tuple[CaptureRecord, ...]], Awaitable[None]], - with_parent: Callable[[CaptureRecord], CaptureRecord], check_limits: Callable[[], Awaitable[None]], files: FileService, ): """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.progress, self.with_parent, self.check_limits = progress, with_parent, check_limits + self.progress, self.check_limits = progress, check_limits async def run(self) -> CaptureCycle: """Refresh membership first, then resume its currently visible child scopes.""" @@ -68,30 +64,36 @@ async def run(self) -> CaptureCycle: cycle = await self.service.begin_cycle( db, BeginCycle(fence=self.fence, configuration=configuration, expected=expected) ) - root = CompletedScope(record_type=configuration.root_record_type) - await self.scan(cycle, root, refresh_membership=True) - after: UUID | None = None while True: async with self.sessions() as db: - roots = await self.service.list_cycle_roots( - db, self.fence, cycle.version.cycle_id, after=after - ) - if not roots: + work = await self.service.next_scope_work(db, self.fence, cycle.version.cycle_id) + if work is None: break - for item in roots: - for child_type in configuration.child_record_types: - scope = CompletedScope(record_type=child_type, container_id=item.native_id) - try: - await self.scan(cycle, scope) - except ScopeAccessLost as error: - await self.withdraw_root( - configuration.root_record_type, - item.native_id, - configuration.child_record_types, - error.removal_reason, - ) - break - after = roots[-1].id + 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(configuration.children_of(work.record_type)), + ) + 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 + ) async with self.sessions() as db: current = await self.service.read_cycle(db, self.fence) if current is None: @@ -104,6 +106,8 @@ async def scan( scope: CompletedScope, *, refresh_membership: bool = False, + parent: SourceRecord | None = None, + parent_epoch: int | None = None, ) -> None: """Only a provider invalid-cursor response permits one explicit sweep restart.""" async with self.sessions() as db: @@ -111,8 +115,10 @@ async def scan( restart = bool( previous and previous.cycle_id == cycle.version.cycle_id - and refresh_membership - and cycle.root_writer_attempt_id != self.fence.attempt_id + and ( + (refresh_membership and previous.membership_attempt_id != self.fence.attempt_id) + or previous.parent_visibility_epoch != parent_epoch + ) ) async with self.sessions() as db: state = await self.service.begin_scan( @@ -130,7 +136,9 @@ async def scan( while state.phase == "collecting": await self.check_limits() try: - page = await self.source.capture_page(scope, state.continuation, files=self.files) + page = await self.source.capture_page( + scope, state.continuation, files=self.files, parent=parent + ) except InvalidScanContinuation: if restarts >= 1: raise @@ -157,7 +165,7 @@ async def scan( scope=scope, cycle_id=state.cycle_id, expected=state.version, - records=tuple(self.with_parent(r) for r in page.records), + records=tuple(self._parented(r, scope) for r in page.records), continuation=page.continuation, final=page.final, ), @@ -165,23 +173,28 @@ async def scan( state = result.state await self.progress(result.capture, page.records) if state.phase == "reconciling" and refresh_membership: - await self.confirm_omissions(cycle) + await self.confirm_omissions(state) await self.reconcile(state) - async def confirm_omissions(self, cycle: CaptureCycle) -> None: - """A provider error or accessible omission fails before any absence removal.""" + @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 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: - roots = await self.service.list_cycle_roots( - db, self.fence, cycle.version.cycle_id, after=after, missing=True - ) - if not roots: + records = await self.service.scan_missing(db, self.fence, state, after=after) + if not records: return - for item in roots: + for record in records: await self.check_limits() - await self.source.confirm_root_absent(item.native_id) - after = roots[-1].id + 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.""" @@ -201,46 +214,23 @@ async def reconcile(self, state: ScanState) -> None: state = result.state await self.progress(result.capture, ()) - async def withdraw_root( - self, - root_type: str, - native_id: str, - child_types: tuple[str, ...], - removal_reason: ScopeRemovalReason, + async def withdraw_parent( + self, parent: SourceRecord, expected_epoch: int, removal_reason: ScopeRemovalReason ) -> None: - """Withdraw the parent before bounded children, retaining the source-audited reason.""" - observed = datetime.now(timezone.utc) + """Withdraw exact owner first; recursive visibility closes before bounded cleanup.""" async with self.sessions() as db: - result = await self.service.capture( + result = await self.service.withdraw_scan_parent( db, - CaptureBatch( - fence=self.fence, - records=( - CaptureRecord( - identity=RecordIdentity(record_type=root_type, native_id=native_id), - payload={"id": native_id}, - kind="delete", - removal_reason=removal_reason, - observed_at=observed, - ), - ), - ), + self.fence, + parent, + expected_epoch=expected_epoch, + removal_reason=removal_reason, ) await self.progress(result, ()) - for child_type in child_types: - while True: - await self.check_limits() - async with self.sessions() as db: - result = await self.service.remove_scope( - db, - self.fence, - RemovedScope( - record_type=child_type, - container_id=native_id, - removal_reason=removal_reason, - observed_at=observed, - ), - ) - await self.progress(result.capture, ()) - if not result.has_more: - break + 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/models/capture_scan.py b/backend/airweave/models/capture_scan.py index 630c7308a1..fe57aab52c 100644 --- a/backend/airweave/models/capture_scan.py +++ b/backend/airweave/models/capture_scan.py @@ -3,7 +3,16 @@ from datetime import datetime from uuid import UUID -from sqlalchemy import BigInteger, CheckConstraint, DateTime, ForeignKey, 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 @@ -18,6 +27,11 @@ class CaptureScan(OrganizationBase): 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) + membership_attempt_id: Mapped[UUID | None] cycle_id: Mapped[UUID] sweep_id: Mapped[UUID] revision: Mapped[int] = mapped_column(BigInteger) @@ -28,6 +42,14 @@ class CaptureScan(OrganizationBase): 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" diff --git a/backend/airweave/platform/sources/linear.py b/backend/airweave/platform/sources/linear.py index dd16c93cbf..5879cd04ce 100644 --- a/backend/airweave/platform/sources/linear.py +++ b/backend/airweave/platform/sources/linear.py @@ -23,6 +23,7 @@ from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError from airweave.domains.browse_tree.types import NodeSelectionData 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, @@ -154,12 +155,12 @@ async def create( "queries": QUERIES, "kinds": cls.canonical_record_types, } - instance._cycle_configuration = CycleConfiguration( + instance._cycle_configuration = CycleConfiguration.from_source( fingerprint=hashlib.sha256( json.dumps(fingerprint, sort_keys=True).encode() ).hexdigest(), - root_record_type="issue", - child_record_types=("comment", "attachment"), + record_types=cls.canonical_record_types, + container_parents=cls.canonical_container_parents, ) await instance.validate() return instance @@ -243,8 +244,19 @@ def _check_issue(self, payload: JsonValue, expected: str | None = None) -> _Issu ) return issue + 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 capture_page( - self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService + self, + 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) @@ -383,8 +395,9 @@ def _single_issue(data: dict[str, JsonValue], native_id: str) -> dict[str, JsonV raise ValueError("Linear returned the wrong issue identity") return node - async def confirm_root_absent(self, native_id: str) -> None: + 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) diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index c266351234..a9238f8376 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -16,6 +16,7 @@ from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError from airweave.domains.browse_tree.types import NodeSelectionData 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, @@ -144,10 +145,10 @@ class SlackSource(BaseSource): canonical_record_types = ("channel", "message") canonical_container_parents = {"message": "channel"} - capture_cycle_configuration = CycleConfiguration( + capture_cycle_configuration = CycleConfiguration.from_source( fingerprint=hashlib.sha256(b"slack:whole-conversations:history-and-replies:v1").hexdigest(), - root_record_type="channel", - child_record_types=("message",), + record_types=canonical_record_types, + container_parents=canonical_container_parents, ) @classmethod @@ -216,8 +217,19 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: completeness="partial" if message.get("files") else "complete", ) + 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 capture_page( - self, scope: CompletedScope, continuation: ScanContinuation, *, files: FileService + 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.""" try: @@ -335,8 +347,9 @@ async def _capture_page_response( recent = (*recent[-15:], digest) return items, following, recent - async def confirm_root_absent(self, native_id: str) -> None: + async def confirm_absent(self, record: SourceRecord) -> None: """Require provider confirmation before hiding an omitted prior conversation.""" + native_id = record.identity.native_id try: await self._get("https://slack.com/api/conversations.info", {"channel": native_id}) except SlackApiError as exc: 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/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index e906b5c67b..a7c42c2419 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -299,6 +299,8 @@ async def main(): "0003_mail_thread_index.py", "0004_projection_generation.py", "0005_capture_scan.py", + "0006_record_visibility.py", + "0007_scan_scope_owner.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: diff --git a/backend/tests/live/capture_resume.py b/backend/tests/live/capture_resume.py index 3ef259c186..463a3c6c12 100644 --- a/backend/tests/live/capture_resume.py +++ b/backend/tests/live/capture_resume.py @@ -82,7 +82,18 @@ async def before_page(self, scope: CompletedScope, continuation: ScanContinuatio assert sync.writer_job_id == self.job_id and sync.writer_attempt_id == self.attempt_id assert job.organization_id == self.organization_id and job.sync_id == self.sync_id assert job.status == "running" and cycle.phase == "active" - assert cycle.root_writer_attempt_id == self.attempt_id + root = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == self.organization_id, + CaptureScan.sync_id == self.sync_id, + CaptureScan.cycle_id == cycle.version.cycle_id, + CaptureScan.record_type == "channel", + CaptureScan.parent_record_id.is_(None), + CaptureScan.container_id.is_(None), + ) + ) + assert root is not None and root.phase == "complete" + assert root.membership_attempt_id == self.attempt_id assert row.cycle_id == cycle.version.cycle_id and row.phase == "collecting" assert row.continuation == continuation.value target = ResumeTarget( diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 2c2714b02a..1d629c7be0 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -86,11 +86,11 @@ def __init__(self, source, counters, record_limit, resume_probe=None): self.canonical_container_parents = source.canonical_container_parents self.capture_cycle_configuration = source.capture_cycle_configuration - async def capture_page(self, scope, continuation, *, files: FileService): + async def capture_page(self, scope, continuation, *, files: FileService, parent=None): if self.resume_probe is not None: await self.resume_probe.before_page(scope, continuation) try: - page = await self.source.capture_page(scope, continuation, files=files) + page = await self.source.capture_page(scope, continuation, files=files, parent=parent) except InvalidScanContinuation: if self.resume_probe is not None: self.resume_probe.cursor_expired() @@ -100,8 +100,11 @@ async def capture_page(self, scope, continuation, *, files: FileService): raise BudgetExceeded("records") return page - async def confirm_root_absent(self, native_id): - await self.source.confirm_root_absent(native_id) + def child_scope(self, parent, record_type): + return self.source.child_scope(parent, record_type) + + async def confirm_absent(self, record): + await self.source.confirm_absent(record) async def bounded_observations(source, cursor, files, counters, record_limit): @@ -657,6 +660,8 @@ async def main(): "0003_mail_thread_index.py", "0004_projection_generation.py", "0005_capture_scan.py", + "0006_record_visibility.py", + "0007_scan_scope_owner.py", ): await connection.run_sync(harness.migrate, migration) organization_id, sync_id = uuid4(), uuid4() diff --git a/backend/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py index cb67fb270c..76e85b40ad 100644 --- a/backend/tests/unit/platform/sources/test_linear_capture.py +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -7,8 +7,9 @@ import pytest from pydantic import ValidationError +from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.page_source import ScopeAccessLost -from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.entities.canonical.requests import 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 ManagedAuthProvider @@ -199,8 +200,16 @@ async def test_omitted_readable_selected_root_blocks_reconciliation(): response({"issue": {**node(), "team": {"id": str(UUID(int=99))}}}), ) with pytest.raises(ValueError, match="still readable"): - await capture.confirm_root_absent(str(ISSUE)) - await capture.confirm_root_absent(str(ISSUE)) + await capture.confirm_absent( + SourceRecord.model_construct( + identity=RecordIdentity(record_type="issue", native_id=str(ISSUE)) + ) + ) + await capture.confirm_absent( + SourceRecord.model_construct( + identity=RecordIdentity(record_type="issue", native_id=str(ISSUE)) + ) + ) @pytest.mark.asyncio diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 3114578411..87c8c4a4b0 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -4,7 +4,8 @@ import pytest -from airweave.domains.entities.canonical.requests import CompletedScope +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.token_providers.static import StaticTokenProvider from airweave.platform.sources.slack import SlackApiError, SlackSource @@ -75,10 +76,16 @@ async def test_incomplete_pagination_fails(pages): async def test_prior_channel_loss_requires_explicit_provider_confirmation(): connector = source() connector._get = AsyncMock(side_effect=SlackApiError("channel_not_found")) - await connector.confirm_root_absent("C1") + await connector.confirm_absent( + SourceRecord.model_construct(identity=RecordIdentity(record_type="channel", native_id="C1")) + ) connector._get = AsyncMock(return_value={"channel": {"id": "C1"}}) with pytest.raises(ValueError, match="omitted"): - await connector.confirm_root_absent("C1") + await connector.confirm_absent( + SourceRecord.model_construct( + identity=RecordIdentity(record_type="channel", native_id="C1") + ) + ) @pytest.mark.asyncio From 4b705253c80f89459c3a9c67a396c9938dd07a6c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:54:38 -0700 Subject: [PATCH 064/266] feat: capture selected GitHub originals with durable nested recovery --- Makefile | 2 +- .../auth_provider/providers/composio.py | 1 + .../entities/canonical/github_projection.py | 146 ++ .../entities/canonical/projection_mappers.py | 4 + .../canonical/tests/test_github_recovery.py | 255 ++++ backend/airweave/platform/configs/config.py | 80 +- .../airweave/platform/entities/__init__.py | 2 + backend/airweave/platform/entities/github.py | 20 + backend/airweave/platform/sources/github.py | 1224 +++++++---------- .../airweave/platform/sources/github_http.py | 217 +++ .../platform/sources/github_models.py | 187 +++ .../entities/test_github_projection.py | 124 ++ .../platform/sources/test_github_capture.py | 593 ++++++++ 13 files changed, 2053 insertions(+), 802 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/github_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_github_recovery.py create mode 100644 backend/airweave/platform/sources/github_http.py create mode 100644 backend/airweave/platform/sources/github_models.py create mode 100644 backend/tests/unit/domains/entities/test_github_projection.py create mode 100644 backend/tests/unit/platform/sources/test_github_capture.py diff --git a/Makefile b/Makefile index e0d7f319a8..72c5a051a4 100644 --- a/Makefile +++ b/Makefile @@ -35,7 +35,7 @@ test-store: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/syncs/tests test-capture: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities/test_github_projection.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 7d5995202a..8bf2a4bdb0 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -117,6 +117,7 @@ async def get_auth_result( "google_drive": frozenset({"www.googleapis.com"}), "slack": frozenset({"slack.com"}), "linear": frozenset({"api.linear.app", "uploads.linear.app"}), + "github": frozenset({"api.github.com"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( diff --git a/backend/airweave/domains/entities/canonical/github_projection.py b/backend/airweave/domains/entities/canonical/github_projection.py new file mode 100644 index 0000000000..5581b81518 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/github_projection.py @@ -0,0 +1,146 @@ +"""Offline GitHub views; search never fetches uncaptured provider content.""" + +from pathlib import Path +from urllib.parse import quote + +from pydantic import BaseModel, ConfigDict, JsonValue + +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 +from airweave.platform.entities.github import GitHubCodeFileEntity, GitHubTextEntity +from airweave.platform.sources.github_models import ( + BranchPayload, + IssuePayload, + Repository, + RepositoryContext, + TreeEntry, +) +from airweave.platform.utils.file_extensions import get_language_for_extension + + +class _Text(BaseModel): + model_config = ConfigDict(extra="ignore") + title: str | None = None + name: str | None = None + body: str | None = None + description: str | None = None + patch: str | None = None + filename: str | None = None + state: str | None = None + + +class _Child(BaseModel): + model_config = ConfigDict(extra="forbid") + native: dict[str, JsonValue] + repository: RepositoryContext + issue_number: int + head_sha: str | None = None + base_sha: str | None = None + + +class _File(BaseModel): + model_config = ConfigDict(extra="forbid") + entry: TreeEntry + path: str + ref: str + commit_sha: str + tree_sha: str + repository: RepositoryContext + + +async def map_github( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> tuple[BaseEntity, ...]: + """Map only current captured originals and verified immutable bytes.""" + kind = record.identity.record_type + if kind == "file": + return await _file(record, storage, directory) + if kind == "repository": + repository = Repository.model_validate(record.payload) + if str(repository.id) != record.identity.native_id: + raise ValueError("GitHub repository identity differs from captured payload") + data, full_name = record.payload, repository.full_name + elif kind == "branch": + branch = BranchPayload.model_validate(record.payload) + data, full_name = branch.branch, branch.repository.full_name + elif kind == "issue": + issue = IssuePayload.model_validate(record.payload) + data, full_name = issue.issue, issue.repository.full_name + elif kind in {"comment", "review", "review_comment", "pull_file"}: + child = _Child.model_validate(record.payload) + data, full_name = child.native, child.repository.full_name + else: + raise ValueError("Unsupported GitHub original kind") + native = _Text.model_validate(data) + return ( + GitHubTextEntity( + original_id=str(record.id), + title=native.title + or native.name + or native.filename + or f"{kind} {record.identity.native_id}", + resource_kind=kind, + repository=full_name, + text=native.body or native.description or native.patch or "", + state=native.state, + content_coverage=record.completeness, + breadcrumbs=[], + created_at=record.source_created_at, + updated_at=record.source_updated_at, + ), + ) + + +async def _file( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> tuple[BaseEntity, ...]: + data = _File.model_validate(record.payload) + if data.path != record.identity.native_id: + raise ValueError("GitHub file path differs from captured identity") + if record.completeness != "complete": + return ( + GitHubTextEntity( + original_id=str(record.id), + title=data.path, + resource_kind="file", + repository=data.repository.full_name, + text="", + content_coverage=record.completeness, + breadcrumbs=[], + ), + ) + if len(record.blobs) != 1 or data.entry.type != "blob": + raise ValueError("Complete GitHub file requires one retained blob") + content = await read_blob(record, record.blobs[0], storage) + suffix = Path(data.path).suffix.lower() + if not suffix: + content.decode("utf-8") # Extensionless text is supported; arbitrary binary is not. + suffix = ".txt" + local = await write_blob(content, directory, suffix=suffix) + owner, repo = data.repository.full_name.split("/", 1) + url = ( + f"https://github.com/{data.repository.full_name}/blob/{data.commit_sha}/" + f"{quote(data.path, safe='/')}" + ) + return ( + GitHubCodeFileEntity( + full_path=data.path, + name=Path(data.path).name, + branch=data.ref, + sha=data.entry.sha, + repo_name=repo, + repo_owner=owner, + path_in_repo=data.path, + language=get_language_for_extension(suffix), + commit_id=data.commit_sha, + url=url, + html_url=url, + size=len(content), + file_type=suffix, + mime_type=record.blobs[0].media_type, + local_path=str(local), + breadcrumbs=[], + ), + ) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 0fefada5fb..c3c4c6ecc2 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -225,6 +225,10 @@ async def map_record( from airweave.domains.entities.canonical.gmail_projection import map_gmail entities = await map_gmail(record, storage, directory) + elif source_name == "github": + from airweave.domains.entities.canonical.github_projection import map_github + + entities = await map_github(record, storage, directory) elif source_name == "linear": from airweave.domains.entities.canonical.linear_projection import map_linear diff --git a/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py new file mode 100644 index 0000000000..34d6a16b3b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py @@ -0,0 +1,255 @@ +"""Real SQL lifecycle with native synthetic GitHub HTTP and actual hierarchy adapter.""" + +import hashlib +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock +from urllib.parse import parse_qs, urlsplit +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.requests import BlobReference +from airweave.domains.entities.canonical.service import CanonicalCaptureService +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.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 GitHubConfig +from airweave.platform.sources.github import GitHubSource + +REPO = {"id": 100, "owner": {"id": 200}, "full_name": "team/repo", "default_branch": "main"} +ISSUE = { + "id": 300, + "number": 7, + "repository_url": "https://api.github.com/repos/team/repo", + "title": "Synthetic issue", +} + + +class NativeGitHub: + """Fixture native pages, retaining actual endpoint/query semantics.""" + + def __init__(self): + self.code = False + self.blob_reads = [] + self.code_bytes = b"synthetic code" + self.blob_sha = hashlib.sha1(b"blob 14\0" + self.code_bytes).hexdigest() + self.comment_pages = [] + self.unavailable = False + + async def get(self, url, **kwargs): + parsed = urlsplit(url) + headers = {} + status = 200 + if self.unavailable: + payload, status = {}, 404 + elif parsed.path == "/repos/team/repo": + payload = REPO + elif parsed.path == "/repos/team/repo/issues": + payload = [ISSUE] + elif parsed.path == "/repos/team/repo/issues/7": + payload = ISSUE + elif parsed.path == "/repos/team/repo/issues/7/comments": + page = int(parse_qs(parsed.query)["page"][0]) + self.comment_pages.append(page) + payload = [{"id": 500 + page, "body": f"Synthetic page {page}"}] + if page == 1: + headers["Link"] = ( + "; rel="next"' + ) + else: + return self.code_response(url) + return httpx.Response( + status, json=payload, headers=headers, request=httpx.Request("GET", url) + ) + + def code_response(self, url): + parsed = urlsplit(url) + if parsed.path.endswith("/branches/main"): + payload = {"name": "main", "commit": {"sha": "1" * 40}} + elif "/git/commits/" in parsed.path: + payload = {"sha": "1" * 40, "tree": {"sha": "2" * 40}} + elif "/git/trees/" in parsed.path: + payload = { + "sha": "2" * 40, + "truncated": False, + "tree": [ + { + "path": name, + "mode": "100644", + "type": "blob", + "sha": self.blob_sha, + "size": len(self.code_bytes), + } + for name in ("first.py", "second.py") + ], + } + elif "/git/blobs/" in parsed.path: + self.blob_reads.append(parsed.path) + return httpx.Response(200, content=self.code_bytes, request=httpx.Request("GET", url)) + else: + raise AssertionError(f"Unexpected synthetic route: {parsed.path}") + return httpx.Response(200, json=payload, request=httpx.Request("GET", url)) + + +async def runner(database, source, native, *, attempt=1, service=None): + default, fence = source + ctx, _, runtime, bus = components(database, source) + client = AsyncMock() + client.get.side_effect = native.get + + @asynccontextmanager + async def stream(method, url, **kwargs): + original = await native.get(url, **kwargs) + streamed = httpx.Response( + original.status_code, + headers=original.headers, + stream=httpx.ByteStream(original.content), + ) + try: + yield streamed + finally: + await streamed.aclose() + + client.stream = MagicMock(side_effect=stream) + files = MagicMock() + + async def store_blob(content, *, media_type=None): + digest = hashlib.sha256(content).hexdigest() + return BlobReference( + key=f"canonical/{fence.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type=media_type, + ) + + files.store_canonical_blob = AsyncMock(side_effect=store_blob) + connector = await GitHubSource.create( + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=client, + config=GitHubConfig.model_validate( + { + "repositories": [{"repository_id": 100, "owner_id": 200, "full_name": "team/repo"}], + "include_code": native.code, + } + ), + ) + pipeline = CanonicalCapturePipeline( + service or default, + 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=files, + ) + runtime.source, runtime.canonical_capture = connector, pipeline + instance = orchestrator(ctx, pipeline, runtime, None, bus) + instance.stream = None + return instance + + +async def test_nested_comment_lost_ack_resumes_without_skipping_or_repeating_committed_page( + 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 == "comment": + raise ConnectionError("synthetic lost acknowledgment") + return result + + native = NativeGitHub() + first = await runner(database, source, native, service=LostAck(source[0].store)) + with pytest.raises(ConnectionError, match="acknowledgment"): + await run(first) + async with database() as db: + comment_scan = await db.scalar( + select(CaptureScan).where(CaptureScan.record_type == "comment") + ) + assert comment_scan.continuation["page"] == 2 + prior_sweep = comment_scan.sweep_id + assert (await db.scalar(select(SyncCursor))).cursor_data["canonical_cycle"][ + "phase" + ] == "active" + second = await runner(database, source, native, attempt=2) + await run(second) + assert native.comment_pages == [1, 2] + async with database() as db: + rows = (await db.scalars(select(Entity))).all() + assert {row.native_id for row in rows} == {"100", "300", "501", "502"} + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "comment")) + assert scan.sweep_id == prior_sweep and scan.phase == "complete" + assert (await db.scalar(select(SyncCursor))).cursor_data["canonical_cycle"][ + "phase" + ] == "complete" + + +async def test_repository_loss_on_retry_hides_all_nested_originals(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 == "comment": + raise ConnectionError("synthetic interruption") + return result + + native = NativeGitHub() + first = await runner(database, source, native, service=LostAck(source[0].store)) + with pytest.raises(ConnectionError): + await run(first) + native.unavailable = True + await run(await runner(database, source, native, attempt=2)) + async with database() as db: + rows = (await db.scalars(select(Entity))).all() + assert rows and all(row.deleted_at is not None for row in rows) + assert all(row.removal_reason == "scope_removed" for row in rows) + for row in rows: + current = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, row.id + ) + assert ( + current.content_access == "unavailable" + and current.payload == {} + and not current.blobs + ) + assert (await db.scalar(select(SyncCursor))).cursor_data["canonical_cycle"][ + "phase" + ] == "complete" + + +async def test_code_scope_resumes_after_committed_blob_without_repeating_first_file( + 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 == "file": + raise ConnectionError("synthetic file commit acknowledgment loss") + return result + + native = NativeGitHub() + native.code = True + with pytest.raises(ConnectionError, match="acknowledgment"): + await run(await runner(database, source, native, service=LostAck(source[0].store))) + async with database() as db: + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "file")) + assert scan.continuation["stack"][0]["offset"] == 1 + sweep = scan.sweep_id + await run(await runner(database, source, native, attempt=2)) + assert len(native.blob_reads) == 2 + async with database() as db: + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "file")) + assert scan.sweep_id == sweep and scan.phase == "complete" + files = ( + await db.scalars(select(Entity).where(Entity.entity_definition_short_name == "file")) + ).all() + assert {item.native_id for item in files} == {"first.py", "second.py"} diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 16dc016873..22cf473231 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -4,7 +4,7 @@ from typing import Literal, Optional from uuid import UUID -from pydantic import AwareDatetime, Field, field_validator, model_validator +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator from airweave.platform.configs._base import BaseConfig, RequiredTemplateConfig from airweave.platform.utils.ssrf import validate_host, validate_url @@ -176,52 +176,46 @@ 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." - ), - ) - 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." - ), + 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", ) - @field_validator("repo_name") + +class GitHubConfig(SourceConfig): + """Selected repositories captured through one owned account and durable page engine.""" + + model_config = ConfigDict(extra="forbid") + repositories: tuple[GitHubRepositorySelection, ...] = Field(min_length=1, max_length=500) + include_code: bool = True + include_conversations: bool = True + + @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): diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index d0269675f5..fa66e99fca 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -121,6 +121,7 @@ GitHubPullRequestEntity, GithubRepoEntity, GitHubRepositoryEntity, + GitHubTextEntity, ) from .gitlab import ( GitLabCodeFileEntity, @@ -469,6 +470,7 @@ ], "github": [ GitHubCodeFileEntity, + GitHubTextEntity, GithubContentEntity, GitHubDirectoryEntity, GitHubFileDeletionEntity, 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/sources/github.py b/backend/airweave/platform/sources/github.py index 1a8c4d0c6c..0fcfee2e36 100644 --- a/backend/airweave/platform/sources/github.py +++ b/backend/airweave/platform/sources/github.py @@ -1,51 +1,66 @@ -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.storage import FileSkippedException +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + parent_container_key, +) +from airweave.domains.entities.canonical.scan_models import ScanContinuation +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.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, +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 +70,467 @@ 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 + 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 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, + async def validate(self) -> None: + """Verify authentication independently from permission to every selected repository.""" + Identity.model_validate(await self.reader.object("/user")) + + 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, ) - 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) - - raise_for_status( - response, - source_short_name=self.short_name, - 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__, + @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/") ) - - 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__, + 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, + 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}") ) - - if files: - await files.save_bytes( - entity=file_entity, - content=content_text.encode("utf-8"), - filename_with_extension=file_path, - logger=self.logger, - ) - - 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: + 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")] - - 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 _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, + ) - async for comment_entity in self._fetch_pr_review_comments( - repo_name, pr_number, owner, repo, [repo_breadcrumb, pr_breadcrumb] - ): - yield comment_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"} + ) - if cursor and latest_updated_at and latest_updated_at != since: - cursor.update(last_pr_updated_at=latest_updated_at) + 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}" + ) + ) + except GitHubUnavailable: + return + if current.id != issue.id: + 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 +539,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/tests/unit/domains/entities/test_github_projection.py b/backend/tests/unit/domains/entities/test_github_projection.py new file mode 100644 index 0000000000..aaf630ef76 --- /dev/null +++ b/backend/tests/unit/domains/entities/test_github_projection.py @@ -0,0 +1,124 @@ +"""Offline search views must retain native text and use only owned blob storage.""" + +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.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity + + +def record(kind, payload, *, native_id="id"): + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity(record_type=kind, native_id=native_id), + revision=1, + payload=payload, + payload_schema_version=1, + capture_hash="a" * 64, + 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, + ) + + +CONTEXT = {"repository_id": 1, "owner_id": 2, "full_name": "team/repo"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["comment", "review", "review_comment"]) +async def test_native_conversation_text_and_exact_kind(kind): + item = record( + kind, + { + "native": {"id": 3, "body": "Native **Markdown**", "state": "APPROVED"}, + "repository": CONTEXT, + "issue_number": 7, + }, + ) + storage = AsyncMock() + async with map_record(item, "github", storage) as views: + assert views[0].text == "Native **Markdown**" + assert views[0].resource_kind == kind + assert views[0].state == "APPROVED" + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_code_uses_verified_owned_blob_and_temporary_file(): + content = b"print('original')\n" + item = record( + "file", + { + "entry": { + "path": "main.py", + "mode": "100644", + "type": "blob", + "sha": "a" * 40, + "size": len(content), + }, + "path": "src/main.py", + "ref": "main", + "commit_sha": "b" * 40, + "tree_sha": "c" * 40, + "repository": CONTEXT, + }, + native_id="src/main.py", + ) + digest = hashlib.sha256(content).hexdigest() + blob = BlobReference( + key=f"canonical/{item.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type="text/x-python", + ) + item = item.model_copy(update={"blobs": (blob,)}) + storage = AsyncMock() + storage.read_file.return_value = content + async with map_record(item, "github", storage) as views: + local = Path(views[0].local_path) + assert local.read_bytes() == content + assert views[0].commit_id == "b" * 40 + assert not local.exists() + storage.read_file.assert_awaited_once_with(blob.key, max_bytes=len(content)) + + +@pytest.mark.asyncio +async def test_missing_complete_bytes_keeps_projection_pending(): + item = record( + "file", + { + "entry": { + "path": "large", + "mode": "100644", + "type": "blob", + "sha": "a" * 40, + "size": 100, + }, + "path": "large", + "ref": "main", + "commit_sha": "b" * 40, + "tree_sha": "c" * 40, + "repository": CONTEXT, + }, + native_id="large", + ) + with pytest.raises(ValueError, match="retained blob"): + async with map_record(item, "github", AsyncMock()): + pass + partial = item.model_copy(update={"completeness": "partial"}) + async with map_record(partial, "github", AsyncMock()) as views: + assert views[0].content_coverage == "partial" and views[0].text == "" diff --git a/backend/tests/unit/platform/sources/test_github_capture.py b/backend/tests/unit/platform/sources/test_github_capture.py new file mode 100644 index 0000000000..cd53848f57 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_github_capture.py @@ -0,0 +1,593 @@ +"""Synthetic native GitHub fixtures; no provider credentials or writes.""" + +import hashlib +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import httpx +import pytest +from pydantic import ValidationError + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.page_source import InvalidScanContinuation, ScopeAccessLost +from airweave.domains.entities.canonical.requests import BlobReference, CompletedScope +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 ManagedAuthProvider +from airweave.platform.configs.config import GitHubConfig +from airweave.platform.sources.github import GitHubSource +from airweave.platform.sources.github_http import GitHubReader + +REPO = {"id": 100, "owner": {"id": 200}, "full_name": "team/repo", "default_branch": "main"} +ISSUE = { + "id": 300, + "number": 7, + "repository_url": "https://api.github.com/repos/team/repo", + "title": "Original", + "unknown": {"preserved": True}, +} +COMMIT, TREE = "1" * 40, "2" * 40 +BRANCH = {"name": "main", "commit": {"sha": COMMIT}} +CONFIG = {"repositories": [{"repository_id": 100, "owner_id": 200, "full_name": "team/repo"}]} + + +def response(value=None, *, status=200, headers=None, content=None): + return httpx.Response( + status, + json=value if content is None else None, + content=content, + headers=headers, + request=httpx.Request("GET", "https://api.github.com/test"), + ) + + +async def source(*responses, config=None): + client = AsyncMock() + client.get.side_effect = responses + + @asynccontextmanager + async def stream(method, url, **kwargs): + original = await client.get(url, **kwargs) + if original.is_stream_consumed: + original = httpx.Response( + original.status_code, + headers=original.headers, + stream=httpx.ByteStream(original.content), + ) + try: + yield original + finally: + await original.aclose() + + client.stream = MagicMock(side_effect=stream) + auth = ManagedAuthProvider( + api_key="synthetic", + connected_account_id="ca_synthetic", + allowed_hosts=frozenset({"api.github.com"}), + ) + capture = await GitHubSource.create( + auth=auth, + logger=MagicMock(), + http_client=client, + config=GitHubConfig.model_validate(config or CONFIG), + ) + return capture, client + + +def saved(record): + return SourceRecord( + id=UUID(int=1), + sync_id=UUID(int=2), + identity=record.identity, + parent=record.parent, + revision=1, + payload=record.payload, + payload_schema_version=record.payload_schema_version, + capture_hash="a" * 64, + content_hash=None, + completeness=record.completeness, + observed_at=datetime.now(timezone.utc), + source_created_at=None, + source_updated_at=None, + deleted_at=None, + removal_reason=None, + blobs=record.blobs, + indexed_revision=None, + indexed_pipeline_version=None, + ) + + +async def repository(capture): + return saved( + ( + await capture.capture_page( + CompletedScope(record_type="repository"), ScanContinuation(), files=MagicMock() + ) + ).records[0] + ) + + +@pytest.mark.asyncio +async def test_managed_auth_and_raw_root_fields(): + capture, client = await source(response({**REPO, "private": True, "extra": [1, 2]})) + root = await repository(capture) + assert root.payload["extra"] == [1, 2] + assert root.identity.native_id == "100" + assert "Authorization" not in client.get.call_args.kwargs["headers"] + assert client.get.call_args.kwargs["follow_redirects"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("changed", [{"id": 101}, {"owner": {"id": 201}}]) +async def test_name_reuse_or_owner_transfer_not_captured(changed): + capture, _ = await source(response({**REPO, **changed})) + result = await capture.capture_page( + CompletedScope(record_type="repository"), ScanContinuation(), files=MagicMock() + ) + assert result.final and not result.records + + +@pytest.mark.asyncio +async def test_issue_payload_and_native_pagination(): + next_link = ( + "; rel="next"' + ) + capture, _ = await source( + response(REPO), + response(REPO), + response([ISSUE], headers={"Link": next_link}), + response(REPO), + response([]), + ) + root = await repository(capture) + scope = capture.child_scope(root, "issue") + first = await capture.capture_page(scope, ScanContinuation(), files=MagicMock(), parent=root) + assert first.records[0].payload["issue"] == ISSUE + assert first.records[0].parent == root.identity + assert not first.final + final = await capture.capture_page(scope, first.continuation, files=MagicMock(), parent=root) + assert final.final + + +@pytest.mark.asyncio +async def test_child_route_revalidated_against_repository_name_reuse(): + capture, client = await source(response(REPO), response({**REPO, "id": 999})) + root = await repository(capture) + with pytest.raises(ScopeAccessLost) as error: + await capture.capture_page( + capture.child_scope(root, "issue"), ScanContinuation(), files=MagicMock(), parent=root + ) + assert error.value.removal_reason == "scope_removed" + assert client.get.call_count == 2 + + +@pytest.mark.asyncio +async def test_issue_number_reuse_does_not_read_another_issues_comments(): + capture, client = await source( + response(REPO), + response(REPO), + response([ISSUE]), + response(REPO), + response({**ISSUE, "id": 999}), + ) + root = await repository(capture) + issue = saved( + ( + await capture.capture_page( + capture.child_scope(root, "issue"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + with pytest.raises(ScopeAccessLost): + await capture.capture_page( + capture.child_scope(issue, "comment"), + ScanContinuation(), + files=MagicMock(), + parent=issue, + ) + assert client.get.call_count == 5 + + +@pytest.mark.asyncio +async def test_pinned_code_blob_retained_and_resume_descends_tree(): + content = b"print('hello')\n" + blob_sha = hashlib.sha1(f"blob {len(content)}\0".encode() + content).hexdigest() + child_tree = "3" * 40 + root_tree = { + "sha": TREE, + "truncated": False, + "tree": [{"path": "src", "mode": "040000", "type": "tree", "sha": child_tree}], + } + capture, client = await source( + response(REPO), + response(REPO), + response(BRANCH), + response(REPO), + response({"sha": COMMIT, "tree": {"sha": TREE}}), + response(root_tree), + response(REPO), + response( + { + "sha": child_tree, + "truncated": False, + "tree": [ + { + "path": "hello.py", + "mode": "100644", + "type": "blob", + "sha": blob_sha, + "size": len(content), + } + ], + } + ), + response(content=content), + response(REPO), + response(root_tree), + ) + root = await repository(capture) + branch = saved( + ( + await capture.capture_page( + capture.child_scope(root, "branch"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + scope = capture.child_scope(branch, "file") + files = MagicMock() + files.store_canonical_blob = AsyncMock( + return_value=BlobReference( + key="canonical/test/blob", + sha256=hashlib.sha256(content).hexdigest(), + size_bytes=len(content), + ) + ) + first = await capture.capture_page(scope, ScanContinuation(), files=files, parent=branch) + assert not first.final and not first.records + second = await capture.capture_page(scope, first.continuation, files=files, parent=branch) + assert second.records[0].identity.native_id == "src/hello.py" + assert second.records[0].payload["commit_sha"] == COMMIT + assert second.records[0].blobs + assert files.store_canonical_blob.call_args.args == (content,) + final = await capture.capture_page(scope, second.continuation, files=files, parent=branch) + assert final.final + assert sum("/git/commits/" in str(call.args[0]) for call in client.get.call_args_list) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tree_response,exception", + [ + (response({"sha": TREE, "truncated": True, "tree": []}), ValueError), + (response({}, status=404), InvalidScanContinuation), + ], +) +async def test_unfinished_tree_never_completes(tree_response, exception): + capture, _ = await source( + response(REPO), + response(REPO), + response(BRANCH), + response(REPO), + response({"sha": COMMIT, "tree": {"sha": TREE}}), + tree_response, + response(BRANCH), + ) + root = await repository(capture) + branch = saved( + ( + await capture.capture_page( + capture.child_scope(root, "branch"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + with pytest.raises(exception): + await capture.capture_page( + capture.child_scope(branch, "file"), + ScanContinuation(), + files=MagicMock(), + parent=branch, + ) + + +@pytest.mark.asyncio +async def test_missing_branch_does_not_withdraw_repository_issues(): + capture, _ = await source( + response(REPO), response(REPO), response({}, status=404), response(REPO) + ) + root = await repository(capture) + result = await capture.capture_page( + capture.child_scope(root, "branch"), ScanContinuation(), files=MagicMock(), parent=root + ) + assert result.final and not result.records + + +@pytest.mark.asyncio +async def test_external_redirect_rejected_before_credentials_forwarded(): + capture, client = await source( + response({}, status=301, headers={"Location": "https://evil.invalid/steal"}) + ) + with pytest.raises(SourceError, match="untrusted"): + await repository(capture) + assert client.get.call_count == 1 + + +@pytest.mark.asyncio +async def test_foreign_or_stuck_page_link_rejected(): + capture, _ = await source( + response(REPO), + response(REPO), + response( + [ISSUE], + headers={"Link": '; rel="next"'}, + ), + ) + root = await repository(capture) + with pytest.raises(ValueError, match="out-of-scope"): + await capture.capture_page( + capture.child_scope(root, "issue"), ScanContinuation(), files=MagicMock(), parent=root + ) + + +def test_rate_limit_never_classified_as_scope_loss(): + reader = GitHubReader(MagicMock(), MagicMock()) + with pytest.raises(SourceRateLimitError): + reader._status(response({"message": "API rate limit exceeded"}, status=403)) + + +def test_legacy_and_ambiguous_selections_rejected(): + with pytest.raises(ValidationError, match="Refresh GitHub"): + GitHubConfig(repo_name="team/repo") + with pytest.raises(ValidationError): + GitHubConfig.model_validate({"repositories": CONFIG["repositories"] * 2}) + + +@pytest.mark.asyncio +async def test_reviews_comments_and_pr_files_keep_original_fields(): + pr_issue = {**ISSUE, "pull_request": {"url": "https://api.github.com/repos/team/repo/pulls/7"}} + pr = { + "id": 400, + "number": 7, + "head": {"sha": "a" * 40}, + "base": {"sha": "b" * 40}, + "changed_files": 1, + } + review = { + "id": 500, + "state": "APPROVED", + "body": "Looks good", + "submitted_at": "2026-09-30T00:00:00Z", + } + file = {"filename": "README.md", "patch": "@@ ...", "status": "modified"} + capture, _ = await source( + response(REPO), + response(REPO), + response([pr_issue]), + response(pr), + response(REPO), + response(pr_issue), + response([review]), + response(REPO), + response(pr_issue), + response(pr), + response([file]), + response(pr), + ) + root = await repository(capture) + issue = saved( + ( + await capture.capture_page( + capture.child_scope(root, "issue"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + assert issue.payload["pull_request_detail"] == pr + result = await capture.capture_page( + capture.child_scope(issue, "review"), ScanContinuation(), files=MagicMock(), parent=issue + ) + assert result.records[0].payload["native"] == review + result = await capture.capture_page( + capture.child_scope(issue, "pull_file"), ScanContinuation(), files=MagicMock(), parent=issue + ) + assert result.records[0].payload["native"] == file + assert result.records[0].completeness == "metadata_only" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["head", "cap"]) +async def test_pr_diff_changed_or_truncated_does_not_complete(change): + pr_issue = {**ISSUE, "pull_request": {}} + pr = { + "id": 400, + "number": 7, + "head": {"sha": "a" * 40}, + "base": {"sha": "b" * 40}, + "changed_files": 3001 if change == "cap" else 1, + } + after = {**pr, "head": {"sha": "c" * 40}} if change == "head" else pr + capture, _ = await source( + response(REPO), + response(REPO), + response([pr_issue]), + response(pr), + response(REPO), + response(pr_issue), + response(pr), + response([{"filename": "one"}]), + response(after), + ) + root = await repository(capture) + issue = saved( + ( + await capture.capture_page( + capture.child_scope(root, "issue"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + with pytest.raises(InvalidScanContinuation if change == "head" else ValueError): + await capture.capture_page( + capture.child_scope(issue, "pull_file"), + ScanContinuation(), + files=MagicMock(), + parent=issue, + ) + + +@pytest.mark.asyncio +async def test_streamed_blob_enforces_incremental_bound(): + class Chunks(httpx.AsyncByteStream): + def __init__(self): + self.parts = iter((b"small", b"oversized")) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.parts) + except StopIteration: + raise AssertionError("Should stop consuming at the size limit") from None + + native = httpx.Response(200, stream=Chunks()) + capture, _ = await source(native) + with pytest.raises(ValueError, match="byte bound"): + await capture.reader.blob("/repos/team/repo/git/blobs/" + "a" * 40, max_bytes=6) + + +def test_corrupt_pinned_progress_rejected(): + from airweave.platform.sources.github_models import FileProgress + + with pytest.raises(ValidationError, match="complete snapshot"): + FileProgress(initialized=True) + with pytest.raises(ValidationError, match="cannot contain"): + FileProgress(commit_sha=COMMIT) + + +@pytest.mark.asyncio +async def test_composio_capability_binds_only_selected_github_account_and_api_host(): + from airweave.domains.auth_provider.exceptions import AuthProviderConfigError + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "synthetic"}, + config={"auth_config_id": "cfg", "account_id": "ca_synthetic"}, + ) + provider._get_with_auth = AsyncMock( + return_value={ + "toolkit": {"slug": "github"}, + "auth_config": {"id": "cfg"}, + "status": "ACTIVE", + } + ) + result = await provider.get_auth_result("github", []) + assert result.managed_auth.connected_account_id == "ca_synthetic" + assert result.managed_auth.allowed_hosts == frozenset({"api.github.com"}) + assert provider._get_with_auth.call_args.args[1].endswith("/ca_synthetic") + provider._get_with_auth.return_value["toolkit"] = {"slug": "slack"} + with pytest.raises(AuthProviderConfigError, match="toolkit"): + await provider.get_auth_result("github", []) + + +@pytest.mark.asyncio +async def test_bad_saved_snapshot_fails_before_provider_request(): + capture, client = await source(response(REPO), response(REPO), response(BRANCH)) + root = await repository(capture) + branch = saved( + ( + await capture.capture_page( + capture.child_scope(root, "branch"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + with pytest.raises(ValidationError): + await capture.capture_page( + capture.child_scope(branch, "file"), + ScanContinuation(value={"initialized": True}), + files=MagicMock(), + parent=branch, + ) + assert client.get.call_count == 3 + with pytest.raises(InvalidScanContinuation, match="branch changed"): + await capture.capture_page( + capture.child_scope(branch, "file"), + ScanContinuation( + value={"initialized": True, "ref": "main", "commit_sha": "f" * 40, "tree_sha": TREE} + ), + files=MagicMock(), + parent=branch, + ) + assert client.get.call_count == 3 + + +@pytest.mark.asyncio +async def test_linked_issue_upload_is_explicitly_partial(): + issue = {**ISSUE, "body": "![screenshot](https://github.com/user-attachments/assets/example)"} + capture, _ = await source(response(REPO), response(REPO), response([issue])) + root = await repository(capture) + page = await capture.capture_page( + capture.child_scope(root, "issue"), ScanContinuation(), files=MagicMock(), parent=root + ) + assert page.records[0].completeness == "partial" + assert page.records[0].payload["issue"] == issue + + +@pytest.mark.asyncio +async def test_native_times_feed_existing_canonical_projection_metadata(): + issue = {**ISSUE, "created_at": "2025-01-01T10:00:00Z", "updated_at": "2026-09-30T10:00:00Z"} + capture, _ = await source(response(REPO), response(REPO), response([issue])) + root = await repository(capture) + assert root.source_created_at is None and root.source_updated_at is None + page = await capture.capture_page( + capture.child_scope(root, "issue"), ScanContinuation(), files=MagicMock(), parent=root + ) + item = page.records[0] + assert item.source_created_at.isoformat() == "2025-01-01T10:00:00+00:00" + assert item.source_updated_at.isoformat() == "2026-09-30T10:00:00+00:00" + assert item.source_created_at != item.observed_at + + +@pytest.mark.asyncio +async def test_existing_direct_pat_auth_remains_supported(): + from airweave.domains.sources.token_providers.credential import DirectCredentialProvider + from airweave.platform.configs.auth import GitHubAuthConfig + + client = AsyncMock() + client.get.return_value = response(REPO) + auth = DirectCredentialProvider( + GitHubAuthConfig(personal_access_token="ghp_synthetic_for_test") + ) + reader = GitHubReader(client, auth) + assert await reader.object("/repos/team/repo") == REPO + assert ( + client.get.call_args.kwargs["headers"]["Authorization"] == "Bearer ghp_synthetic_for_test" + ) + + +def test_review_submission_is_not_invented_as_creation_time(): + item = GitHubSource._record( + CompletedScope(record_type="review"), + "1", + {"native": {"id": 1, "submitted_at": "2026-09-30T10:00:00Z"}}, + ) + assert item.source_created_at is None and item.source_updated_at is None + assert item.payload["native"]["submitted_at"] == "2026-09-30T10:00:00Z" From 37e652663d4a9d175ec8ff54f169c6ee13406bbe Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 07:56:50 -0700 Subject: [PATCH 065/266] fix: confirm GitHub issue transfers before withdrawing old scope --- backend/airweave/platform/sources/github.py | 14 +++++-- .../platform/sources/test_github_capture.py | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/backend/airweave/platform/sources/github.py b/backend/airweave/platform/sources/github.py index 0fcfee2e36..ce69f91dfc 100644 --- a/backend/airweave/platform/sources/github.py +++ b/backend/airweave/platform/sources/github.py @@ -251,9 +251,13 @@ async def _verify_parent_repository(self, parent: SourceRecord) -> None: 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}") + await self.reader.object( + f"/repos/{context.full_name}/issues/{issue.number}", redirects=True + ) ) - if current_issue.id != issue.id: + 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( @@ -502,12 +506,14 @@ async def confirm_absent(self, record: SourceRecord) -> None: try: current = Issue.model_validate( await self.reader.object( - f"/repos/{body.repository.full_name}/issues/{issue.number}" + f"/repos/{body.repository.full_name}/issues/{issue.number}", redirects=True ) ) except GitHubUnavailable: return - if current.id != issue.id: + if current.id != issue.id or current.repository_url != ( + f"https://api.github.com/repos/{body.repository.full_name}" + ): return else: raise ValueError("GitHub absence confirmation is only valid for inventory owners") diff --git a/backend/tests/unit/platform/sources/test_github_capture.py b/backend/tests/unit/platform/sources/test_github_capture.py index cd53848f57..3d7986349f 100644 --- a/backend/tests/unit/platform/sources/test_github_capture.py +++ b/backend/tests/unit/platform/sources/test_github_capture.py @@ -591,3 +591,40 @@ def test_review_submission_is_not_invented_as_creation_time(): ) assert item.source_created_at is None and item.source_updated_at is None assert item.payload["native"]["submitted_at"] == "2026-09-30T10:00:00Z" + + +@pytest.mark.asyncio +async def test_transferred_issue_withdraws_old_scope_without_reading_new_comments(): + target = "https://api.github.com/repos/other/repo/issues/9" + moved = {**ISSUE, "number": 9, "repository_url": "https://api.github.com/repos/other/repo"} + capture, client = await source( + response(REPO), + response(REPO), + response([ISSUE]), + response(REPO), + response({}, status=301, headers={"Location": target}), + response(moved), + response({}, status=301, headers={"Location": target}), + response(moved), + ) + root = await repository(capture) + issue = saved( + ( + await capture.capture_page( + capture.child_scope(root, "issue"), + ScanContinuation(), + files=MagicMock(), + parent=root, + ) + ).records[0] + ) + with pytest.raises(ScopeAccessLost) as error: + await capture.capture_page( + capture.child_scope(issue, "comment"), + ScanContinuation(), + files=MagicMock(), + parent=issue, + ) + assert error.value.removal_reason == "scope_removed" + await capture.confirm_absent(issue) + assert all("/comments" not in call.args[0] for call in client.get.call_args_list) From 250828851c5bac235af72fd17dc1a2094a6efdf3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:30:02 -0700 Subject: [PATCH 066/266] Distinguish discovery policies from exhaustive scope reconciliation --- .../domains/entities/canonical/coverage.py | 56 ++++ .../entities/canonical/cycle_models.py | 14 +- .../domains/entities/canonical/page_source.py | 11 +- .../domains/entities/canonical/query.py | 6 +- .../entities/canonical/query_models.py | 2 + .../domains/entities/canonical/scan_models.py | 18 +- .../domains/entities/canonical/scan_store.py | 74 ++++- .../domains/entities/canonical/service.py | 6 + .../canonical/tests/test_owned_search.py | 38 +++ .../canonical/tests/test_scan_policies.py | 253 ++++++++++++++++++ backend/airweave/domains/search/owned.py | 3 + .../airweave/domains/search/owned_models.py | 2 + .../domains/sync_pipeline/canonical_scan.py | 34 ++- 13 files changed, 506 insertions(+), 11 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/coverage.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_scan_policies.py diff --git a/backend/airweave/domains/entities/canonical/coverage.py b/backend/airweave/domains/entities/canonical/coverage.py new file mode 100644 index 0000000000..efd1927cc6 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/coverage.py @@ -0,0 +1,56 @@ +"""Describe persisted capture promises independently of indexing counts or job success.""" + +from typing import Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, ValidationError +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.domains.entities.canonical.cycle_models import ( + CYCLE_KEY, + CaptureCycle, + CompletionPolicy, +) +from airweave.models.sync_cursor import SyncCursor + + +class CaptureCoverage(BaseModel): + """No cycle means unknown; completion only certifies the declared scope policies.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + phase: Literal["active", "complete"] + policies: dict[str, CompletionPolicy] + discovery: Literal["incomplete", "pending", "scope_enumeration_complete"] + + +async def capture_coverage( + db: AsyncSession, organization_id: UUID, sync_ids: tuple[UUID, ...] +) -> dict[UUID, CaptureCoverage]: + """Read only existing tenant-scoped cycle metadata, never create an authority.""" + rows = await db.scalars( + select(SyncCursor).where( + SyncCursor.organization_id == organization_id, SyncCursor.sync_id.in_(sync_ids) + ) + ) + result = {} + for row in rows: + raw = row.cursor_data.get(CYCLE_KEY) + if raw is None: + continue + try: + cycle = CaptureCycle.model_validate(raw) + except ValidationError: + continue # Unknown/corrupt capture state cannot certify coverage. + policies = {kind: cycle.configuration.policy(kind) for kind in cycle.configuration.parents} + discovery = ( + "incomplete" + if any(p != "exhaustive" for p in policies.values()) + else "pending" + if cycle.phase == "active" + else "scope_enumeration_complete" + ) + result[row.sync_id] = CaptureCoverage( + phase=cycle.phase, policies=policies, discovery=discovery + ) + return result diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 2353ba1627..27795a6b16 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -8,6 +8,8 @@ from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.requests import RecordKind, WriterFence +CompletionPolicy = Literal["exhaustive", "discovery_only", "discovery_with_validation"] + CYCLE_KEY = "canonical_cycle" @@ -17,6 +19,11 @@ class CycleConfiguration(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") parents: dict[RecordKind, tuple[RecordKind | None, ...]] + completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) + + def policy(self, kind: str) -> CompletionPolicy: + """Omitted declarations preserve historical exhaustive semantics.""" + return self.completion_policies.get(kind, "exhaustive") @classmethod def from_source( @@ -25,6 +32,7 @@ def from_source( fingerprint: str, record_types: tuple[str, ...], container_parents: dict[str, str | tuple[str | None, ...]], + completion_policies: dict[RecordKind, CompletionPolicy] | None = None, ) -> "CycleConfiguration": """Snapshot the existing source declaration; no second topology registry.""" if len(set(record_types)) != len(record_types): @@ -35,7 +43,9 @@ def from_source( for kind in record_types: declared = container_parents.get(kind, (None,)) parents[kind] = (declared,) if isinstance(declared, str) else declared - return cls(fingerprint=fingerprint, parents=parents) + return cls( + fingerprint=fingerprint, parents=parents, completion_policies=completion_policies or {} + ) @model_validator(mode="before") @classmethod @@ -56,6 +66,8 @@ def upgrade_flat(cls, value): @model_validator(mode="after") def reachable_types(self) -> "CycleConfiguration": """Type cycles are valid; every declared kind must be reachable from a root.""" + if set(self.completion_policies) - set(self.parents): + raise ValueError("Completion policy has an undeclared record kind") if not self.parents or any(not kind or not values for kind, values in self.parents.items()): raise ValueError("Capture topology must declare kinds and their parents") if any(len(set(values)) != len(values) for values in self.parents.values()): diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index 36252787c6..0a288c3330 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -1,4 +1,4 @@ -"""Opted-in whole-scope pages; the existing capture pipeline owns durable progress.""" +"""Opted-in scope pages; the existing capture pipeline owns durable progress.""" from typing import Protocol, runtime_checkable @@ -68,3 +68,12 @@ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: async def confirm_absent(self, record: SourceRecord) -> None: """Raise unless an omitted inventory record is confirmed outside accessible scope.""" ... + + +@runtime_checkable +class KnownObjectSource(Protocol): + """Only discovery-with-validation sources implement exact current reads.""" + + async def refresh_known(self, record: SourceRecord, *, files: FileService) -> CaptureRecord: + """Return exact fresh state or explicit unavailability; errors never mean absence.""" + ... diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py index 07895e1690..645e3a6b6a 100644 --- a/backend/airweave/domains/entities/canonical/query.py +++ b/backend/airweave/domains/entities/canonical/query.py @@ -7,6 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.blob_materializer import BlobIntegrityError, read_blob +from airweave.domains.entities.canonical.coverage import capture_coverage from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import ( MailThreadCursor, @@ -116,7 +117,10 @@ async def list_records( after_id=page[-1].id, ) ) - return RecordPage(records=page, next_cursor=next_cursor, has_more=more) + coverage = await capture_coverage(db, organization_id, (sync_id,)) + return RecordPage( + records=page, next_cursor=next_cursor, has_more=more, capture=coverage.get(sync_id) + ) async def changes( self, diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index ea36e14e9f..2d64cae6c2 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, ConfigDict, Field +from airweave.domains.entities.canonical.coverage import CaptureCoverage from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord @@ -31,6 +32,7 @@ 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 diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index 5da784983e..bda883c8ee 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -1,4 +1,4 @@ -"""Whole-scope page commits; partial ranges must not use absence reconciliation.""" +"""Declared-scope page commits; incomplete discovery never owns absence reconciliation.""" import json from typing import Literal @@ -6,7 +6,8 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, field_validator -from airweave.domains.entities.canonical.models import CaptureResult +from airweave.domains.entities.canonical.cycle_models import CompletionPolicy +from airweave.domains.entities.canonical.models import CaptureResult, SourceRecord from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, WriterFence @@ -47,7 +48,7 @@ class BeginScan(BaseModel): class ScanState(BaseModel): - """Complete means successful whole-scope enumeration and absence reconciliation.""" + """Complete means the declared policy finished, not necessarily exhaustive discovery.""" model_config = ConfigDict(extra="forbid", frozen=True) scope: CompletedScope @@ -60,6 +61,7 @@ class ScanState(BaseModel): completed_at: AwareDatetime | None parent_visibility_epoch: int | None = None membership_attempt_id: UUID | None = None + completion_policy: CompletionPolicy = "exhaustive" class CommitScanPage(BaseModel): @@ -94,3 +96,13 @@ class ScanResult(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) state: ScanState capture: CaptureResult + + +class CommitOmission(BaseModel): + """Fresh exact read after discovery; never broad absence evidence.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + fence: WriterFence + state: ScanState + expected_record: SourceRecord + observation: CaptureRecord diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 6eefe3bce4..9ed83e4443 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -28,6 +28,7 @@ ) from airweave.domains.entities.canonical.scan_models import ( BeginScan, + CommitOmission, CommitScanPage, ReconcileScan, ScanContinuation, @@ -42,6 +43,7 @@ ) from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity +from airweave.models.sync_cursor import SyncCursor class ScanConflict(CanonicalStoreError): @@ -121,7 +123,15 @@ async def _state(self, db: AsyncSession, row: CaptureScan) -> ScanState: parent = await db.get(Entity, row.parent_record_id) if row.parent_record_id else None if row.parent_record_id is not None and parent is None: raise CycleConflict("Scope parent disappeared; explicitly restart the cycle") - return scan_state(row, parent) + cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == row.sync_id)) + cycle = cycle_state(cursor) if cursor else None + return scan_state(row, parent).model_copy( + update={ + "completion_policy": cycle.configuration.policy(row.record_type) + if cycle + else "exhaustive" + } + ) async def read( self, db: AsyncSession, fence: WriterFence, scope: CompletedScope @@ -313,7 +323,10 @@ async def missing( parent = await scope_owner(db, fence, cycle, state.scope) if ( row.phase != "reconciling" - or row.membership_attempt_id != fence.attempt_id + or ( + cycle.configuration.children_of(row.record_type) + and row.membership_attempt_id != fence.attempt_id + ) or row.parent_visibility_epoch != (parent.visibility_epoch if parent else None) ): raise CycleConflict("Omission checks require completed current-attempt inventory") @@ -331,9 +344,10 @@ async def missing( identity.container_id if identity else None ), Entity.record_revision > 0, - Entity.deleted_at.is_(None), Entity.last_seen_run_id.is_distinct_from(row.sweep_id), ] + if cycle.configuration.policy(row.record_type) != "discovery_with_validation": + predicates.append(Entity.deleted_at.is_(None)) if after is not None: predicates.append(Entity.id > after) records = ( @@ -341,6 +355,44 @@ async def missing( ).all() return tuple(source_record(record) for record in records) + async def omission(self, db: AsyncSession, request: CommitOmission) -> ScanResult: + """Commit exact known-object validation with its sweep acknowledgement.""" + sync = await self.records._fenced_sync(db, request.fence) + _, cycle = await attest_cycle(db, request.fence, request.state.cycle_id) + if ( + cycle.configuration.policy(request.state.scope.record_type) + != "discovery_with_validation" + ): + raise ScanConflict("This scope does not accept known-object validation") + candidates = await self.missing(db, request.fence, request.state) + previous = request.expected_record + current = next((record for record in candidates if record.id == previous.id), None) + if ( + current is None + or current.revision != previous.revision + or current.identity != previous.identity + ): + raise ScanConflict("Known object changed or is outside the next validation batch") + observation = request.observation + if observation.removal_reason == "absent": + raise ScanConflict("Discovery requires explicit unavailability, never absence") + if observation.identity != current.identity or observation.parent != current.parent: + raise ScanConflict("Known-object validation changed identity or owner") + row = self._expect( + await self._row(db, request.fence, request.state.scope), + request.state.version, + request.state.cycle_id, + ) + captured = await self.records._capture_locked( + db, + sync, + CaptureBatch(fence=request.fence, records=(observation,)), + seen_id=row.sweep_id, + ) + row.revision += 1 + await db.flush() + return ScanResult(state=await self._state(db, row), capture=captured) + async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: """A final page is necessary; absence completion is durable and bounded.""" sync = await self.records._fenced_sync(db, request.fence) @@ -359,6 +411,22 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul raise CycleConflict("Membership belongs to an earlier writer attempt") if row.phase != "reconciling": raise ScanConflict("Only fully collected scans can reconcile absence") + policy = cycle.configuration.policy(row.record_type) + if policy != "exhaustive": + if policy == "discovery_with_validation" and await self.missing( + db, request.fence, await self._state(db, row) + ): + raise ScanConflict("Known objects still require exact validation") + row.phase = "complete" + row.completed_at = datetime.now(timezone.utc) + row.revision += 1 + await db.flush() + return ScanResult( + state=await self._state(db, row), + capture=CaptureResult( + changes=(), sequence=sync.observed_change_sequence, unchanged=0 + ), + ) result = await self.records._reconcile_scope_locked( db, sync, diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index e9d434521c..5a48ca8283 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -33,6 +33,7 @@ ) from airweave.domains.entities.canonical.scan_models import ( BeginScan, + CommitOmission, CommitScanPage, ReconcileScan, ScanResult, @@ -163,6 +164,11 @@ async def commit_scan_page(self, db: AsyncSession, request: CommitScanPage) -> S async with UnitOfWork(db): return await self.scans.page(db, request) + async def commit_omission(self, db: AsyncSession, request: CommitOmission) -> ScanResult: + """Accept exact provider validation without network I/O in the transaction.""" + async with UnitOfWork(db): + return await self.scans.omission(db, request) + async def reconcile_scan(self, db: AsyncSession, request: ReconcileScan) -> ScanResult: """Commit one bounded whole-scope absence reconciliation step.""" async with UnitOfWork(db): diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 388727c3b5..44db15469e 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -444,3 +444,41 @@ async def test_dates_and_types_prefilter_before_candidate_limit(indexed, http_se assert by_field["airweave_system_metadata.source_created_known"] == 1 assert by_field["airweave_system_metadata.source_created_us"] == 1767225600000001 assert by_field["airweave_system_metadata.source_updated_known"] == 1 + + +async def test_capture_discovery_coverage_is_separate_from_retrieval( + database, source, indexed, http_search +): + from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration + + service, fence = source + async with database() as db: + await service.begin_cycle( + db, + BeginCycle( + fence=fence, + configuration=CycleConfiguration( + fingerprint="d" * 64, + parents={"event": (None,)}, + completion_policies={"event": "discovery_only"}, + ), + ), + ) + client, vector, _, _, _ = http_search + vector.seed_results(SearchResults(results=[])) + response = await client.post( + "/sync/search", + json={ + "query": "budget", + "sync_ids": [str(fence.sync_id)], + "mode": "keyword", + }, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["sources"][0]["capture"] == { + "phase": "active", + "policies": {"event": "discovery_only"}, + "discovery": "incomplete", + } + assert not body["retrieval_incomplete"] 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..34fd304cce --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py @@ -0,0 +1,253 @@ +"""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, 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 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/search/owned.py b/backend/airweave/domains/search/owned.py index cebb66b458..e2c4403a8a 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.api.context import ApiContext +from airweave.domains.entities.canonical.coverage import capture_coverage from airweave.domains.entities.canonical.projection_models import ProjectionLocator from airweave.domains.entities.canonical.projection_store import publication_matches from airweave.domains.entities.canonical.requests import RecordIdentity @@ -396,9 +397,11 @@ async def _coverage( .group_by(Entity.sync_id) ) counts = {sync: (active, pending) for sync, active, pending 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], ) diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py index cc0f13ed0c..48facdb70f 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -5,6 +5,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator +from airweave.domains.entities.canonical.coverage import CaptureCoverage from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.search.types import RetrievalStrategy @@ -62,6 +63,7 @@ class OwnedSearchCoverage(BaseModel): sync_id: UUID active_records: int + capture: CaptureCoverage | None = None pending_records: int diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 8bf38f795e..21e2b32ee5 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -13,6 +13,7 @@ from airweave.domains.entities.canonical.page_source import ( CanonicalPageSource, InvalidScanContinuation, + KnownObjectSource, ScopeAccessLost, ScopeRemovalReason, ) @@ -23,6 +24,7 @@ ) from airweave.domains.entities.canonical.scan_models import ( BeginScan, + CommitOmission, CommitScanPage, ReconcileScan, ScanState, @@ -172,8 +174,11 @@ async def scan( ) state = result.state await self.progress(result.capture, page.records) - if state.phase == "reconciling" and refresh_membership: - await self.confirm_omissions(state) + if state.phase == "reconciling": + if state.completion_policy == "discovery_with_validation": + state = await self.refresh_known(state) + elif state.completion_policy == "exhaustive" and refresh_membership: + await self.confirm_omissions(state) await self.reconcile(state) @staticmethod @@ -183,6 +188,31 @@ def _parented(record: CaptureRecord, scope: CompletedScope) -> CaptureRecord: 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 From 10c6549e3f7ad2504415fdc5e22ed0d2723cb03e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:30:49 -0700 Subject: [PATCH 067/266] Prove completed discovery cycles remain explicitly incomplete --- .../canonical/tests/test_scan_policies.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py b/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py index 34fd304cce..054b874129 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py +++ b/backend/airweave/domains/entities/canonical/tests/test_scan_policies.py @@ -6,7 +6,11 @@ import pytest from airweave.domains.entities.canonical.coverage import capture_coverage -from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration +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 ( @@ -109,6 +113,15 @@ async def test_incomplete_discovery_preserves_known_record_and_reports_policy( ), ) + 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 From fd8a77710ac9a3a410c06408af94802d57baaab5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:34:26 -0700 Subject: [PATCH 068/266] Resume Wispr meeting capture through durable discovery scopes --- .../entities/canonical/projection_mappers.py | 2 + .../canonical/tests/test_wispr_recovery.py | 138 +++++++++ backend/airweave/platform/sources/wispr.py | 293 +++++++++++++----- backend/tests/live/provider_lifecycle.py | 6 +- .../platform/sources/test_wispr_capture.py | 182 ++++++++++- 5 files changed, 519 insertions(+), 102 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index c3c4c6ecc2..d5510b17b8 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -174,6 +174,8 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: + if record.identity.record_type == "meeting_listing": + return () if record.identity.record_type != "meeting": raise ProjectionMappingError("Unsupported Wispr record type") responses = record.payload.get("responses") 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..15ad646ca1 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py @@ -0,0 +1,138 @@ +"""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.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): + 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_MEETINGS": + return {"meetings": rows, "has_more": False} + native = arguments["meeting_id"] + body_calls.append(native) + if len(body_calls) == fail_at: + raise 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() + ) + + +async def test_failed_body_resumes_without_repeating_completed_siblings(database, source): + 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) + with pytest.raises(ValueError, match="Synthetic opaque"): + 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 diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 340757adde..264f2145ff 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -2,20 +2,26 @@ from __future__ import annotations +import hashlib +import json import re -from contextlib import aclosing from datetime import datetime, timezone -from typing import AsyncGenerator from urllib.parse import quote -from pydantic import BaseModel, ConfigDict, JsonValue +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool from airweave.core.logging import ContextualLogger -from airweave.domains.browse_tree.types import NodeSelectionData -from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity +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.token_providers.protocol import ManagedToolAuthProvider from airweave.domains.storage.file_service import FileService -from airweave.domains.syncs.cursors.cursor import SyncCursor from airweave.platform.decorators import source from airweave.platform.http_client.airweave_client import AirweaveHttpClient from airweave.platform.http_client.composio_transport import ComposioTransport @@ -40,10 +46,44 @@ class MeetingPage(BaseModel): """Wispr's real search response pagination fields.""" model_config = ConfigDict(extra="ignore") - meetings: list[dict[str, JsonValue]] - has_more: bool + meetings: list[dict[str, JsonValue]] = Field(max_length=200) + has_more: StrictBool next_cursor: str | None = None - truncated: bool = False + truncated: StrictBool = False + + +class _Listing(BaseModel): + model_config = ConfigDict(extra="ignore") + id: str = Field(min_length=1) + + +class _MeetingVersion(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( @@ -60,8 +100,13 @@ class WisprSource(BaseSource): Captures remain explicitly partial, even after every exposed text range is fetched. """ - canonical_record_types = ("meeting",) - canonical_container_parents: dict[str, str] = {} + canonical_record_types = ("meeting_listing", "meeting") + canonical_container_parents = {"meeting": "meeting_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( @@ -78,6 +123,25 @@ async def create( 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": 2, + "account": auth.connected_account_id, + "inventory": "meeting_listing", + "body": "meeting", + "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"}, + ) return instance async def _post(self, path: str, body: dict) -> dict: @@ -124,60 +188,6 @@ async def validate(self) -> None: await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", {"limit": 1}) ) - async def _list_window( - self, - since: datetime | None, - until: datetime | None, - ) -> tuple[list[dict[str, JsonValue]], bool]: - arguments: dict[str, JsonValue] = {"limit": 200} - if since is not None: - arguments["since"] = since.isoformat() - if until is not None: - arguments["until"] = until.isoformat() - records: list[dict[str, JsonValue]] = [] - seen_cursors: set[str] = set() - while True: - page = MeetingPage.model_validate( - await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", arguments) - ) - records.extend(page.meetings) - if page.truncated: - return records, True - if not page.has_more: - return records, False - if not page.next_cursor or page.next_cursor in seen_cursors: - raise ValueError("Wispr returned incomplete or repeated pagination") - if len(records) >= 1000: - return records, True - seen_cursors.add(page.next_cursor) - arguments["cursor"] = page.next_cursor - - async def _list_all(self) -> AsyncGenerator[dict[str, JsonValue], None]: - windows: list[tuple[datetime | None, datetime | None, int]] = [(None, None, 0)] - seen_ids: set[str] = set() - while windows: - since, until, depth = windows.pop() - records, truncated = await self._list_window(since, until) - if truncated: - starts = [ - datetime.fromisoformat(str(row["start"]).replace("Z", "+00:00")) - for row in records - ] - if not starts or min(starts) == max(starts) or depth >= 32: - raise ValueError("Wispr query cap cannot be partitioned safely") - split = min(starts) + (max(starts) - min(starts)) / 2 - if (since is not None and split <= since) or (until is not None and split >= until): - raise ValueError("Wispr date partition made no progress") - windows.extend([(since, split, depth + 1), (split, until, depth + 1)]) - continue - for record in records: - identity = record.get("id") - if not isinstance(identity, str) or not identity: - raise ValueError("Wispr returned a meeting without identity") - if identity not in seen_ids: - seen_ids.add(identity) - yield record - @staticmethod def _next_offset(text: JsonValue, field: str) -> int | None: if not isinstance(text, str): @@ -222,24 +232,147 @@ async def _meeting(self, identity: str) -> dict[str, JsonValue]: offsets = following raise ValueError("Wispr meeting exceeds the bounded capture page limit") - async def generate_observations( + def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: + """Keep each body attached to its exact native listing observation.""" + if parent.identity.record_type != "meeting_listing" or record_type != "meeting": + raise ValueError("Unsupported Wispr child scope") + # Preserve existing provider-global body identity, including null container. + return CompletedScope(record_type="meeting", 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, *, - cursor: SyncCursor | None = None, - files: FileService | None = None, - node_selections: list[NodeSelectionData] | None = None, - ) -> AsyncGenerator[CaptureRecord, None]: - """Capture full exposed ranges without claiming deletion or snapshot completeness.""" - if node_selections: - raise ValueError("Wispr capture does not support selected meeting scopes") - async with aclosing(self._list_all()) as meetings: - async for listing in meetings: - identity = str(listing["id"]) - payload = await self._meeting(identity) - payload["listing"] = listing - yield CaptureRecord( + 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. + """ + if scope.record_type == "meeting_listing" and scope.parent is None and parent is None: + return await self._listing_page(continuation) + if ( + scope.record_type != "meeting" + or parent is None + or parent.identity.record_type != "meeting_listing" + 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._meeting(identity) + payload["listing"] = parent.payload + modified = _MeetingVersion.model_validate(payload["responses"][0]["response"]).modified_at + return CapturePage( + records=( + CaptureRecord( identity=RecordIdentity(record_type="meeting", 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, 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() + page = MeetingPage.model_validate( + await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", arguments) + ) + count = progress.count + len(page.meetings) + following = page.next_cursor if page.has_more else None + starts = [_MeetingStart.model_validate(row).start for row in page.meetings] + dates_complete = progress.dates_complete and all(value is not None for value in starts) + known_starts = [value for value in starts if value is not None] + [ + d for d in (progress.earliest, progress.latest) if d is not None + ] + earliest = min(known_starts) if known_starts else None + latest = max(known_starts) if known_starts 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 + identities = [_Listing.model_validate(row).id for row in page.meetings] + if len(identities) != len(set(identities)): + raise ValueError("Wispr listing page repeats a meeting identity") + records = tuple( + CaptureRecord( + identity=RecordIdentity(record_type="meeting_listing", native_id=identity), + payload=row, + completeness="metadata_only", + observed_at=datetime.now(timezone.utc), + ) + for identity, row in zip(identities, page.meetings, strict=True) + ) + return CapturePage( + records=records, + continuation=ScanContinuation(value=next_progress.model_dump(mode="json")), + final=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/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 1d629c7be0..c1b4e0a5f9 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -182,13 +182,9 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert saved["canonical_page_token"] assert counters["started"] == counters["completed"] == (0 if loaded else 1) assert bool(counters["resumed_changes_requests"]) == loaded - elif name == "slack": + elif name in {"slack", "wispr"}: assert saved["canonical_cycle"]["phase"] == "complete" assert saved["canonical_cycle"]["completed_job_id"] - elif name == "wispr": - assert saved is None - assert counters["started"] == counters["completed"] == 0 - return else: assert saved["canonical_query"] == manifest["query"] assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index 89f253b9c8..e304827f1b 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -83,35 +83,183 @@ async def test_changed_meeting_fails_instead_of_mixing_versions(monkeypatch): @pytest.mark.asyncio -async def test_query_cap_partitions_by_meeting_start_not_modified(monkeypatch): +async def test_explicit_tool_failure_never_becomes_partial_success(monkeypatch): + connector = await source() + monkeypatch.setattr( + connector, + "_post", + AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": "failed"}]), + ) + with pytest.raises(ValueError, match="tool execution failed; capture is incomplete"): + await connector._meeting("m") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"meetings": [], "has_more": False, "truncated": True}, + {"meetings": [], "has_more": True, "next_cursor": None}, + {"meetings": [{"id": "m"}, {"id": "m"}], "has_more": False}, + ], +) +async def test_listing_failure_never_commits_partial_final_page(monkeypatch, response): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + + connector = await source() + monkeypatch.setattr(connector, "_execute", AsyncMock(return_value=response)) + with pytest.raises(ValueError): + await connector.capture_page( + CompletedScope(record_type="meeting_listing"), ScanContinuation(), files=MagicMock() + ) + + +@pytest.mark.asyncio +async def test_listing_preserves_native_row_and_continuation(monkeypatch): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + + connector = await source() + row = {"id": "m", "title": "Meeting", "unknown_native": True, "start": "2026-01-01T00:00:00Z"} + execute = AsyncMock(return_value={"meetings": [row], "has_more": True, "next_cursor": "next"}) + monkeypatch.setattr(connector, "_execute", execute) + page = await connector.capture_page( + CompletedScope(record_type="meeting_listing"), ScanContinuation(), files=MagicMock() + ) + assert page.records[0].payload == row + assert page.records[0].completeness == "metadata_only" + assert page.records[0].source_created_at is None + assert page.continuation.value["cursor"] == "next" + assert page.final is False + with pytest.raises(ValueError, match="repeated pagination"): + await connector.capture_page( + CompletedScope(record_type="meeting_listing"), page.continuation, files=MagicMock() + ) + + +@pytest.mark.asyncio +async def test_listing_projection_is_explicitly_empty(): + from types import SimpleNamespace + + from airweave.domains.entities.canonical.projection_mappers import _wispr + + assert _wispr(SimpleNamespace(identity=SimpleNamespace(record_type="meeting_listing"))) == () + + +@pytest.mark.asyncio +async def test_partitioned_listing_resumes_both_half_open_windows(monkeypatch): + from airweave.domains.entities.canonical.page_source import CanonicalPageSource + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + connector = await source() - row1 = {"id": "1", "start": "2026-01-01T00:00:00Z", "modified_at": "2026-09-01T00:00:00Z"} - row2 = {"id": "2", "start": "2026-01-03T00:00:00Z"} - listing = AsyncMock(side_effect=[([row1, row2], True), ([row2], False), ([row1], False)]) - monkeypatch.setattr(connector, "_list_window", listing) - assert {row["id"] async for row in connector._list_all()} == {"1", "2"} - assert listing.call_args_list[1].args[0].isoformat() == "2026-01-02T00:00:00+00:00" + assert isinstance(connector, CanonicalPageSource) + rows = [ + {"id": "a", "start": "2026-01-01T00:00:00Z"}, + {"id": "b", "start": "2026-01-03T00:00:00Z"}, + ] + execute = AsyncMock( + side_effect=[ + {"meetings": rows, "has_more": True, "truncated": True}, + {"meetings": rows[:1], "has_more": False}, + {"meetings": rows[1:], "has_more": False}, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + scope = CompletedScope(record_type="meeting_listing") + first = await connector.capture_page(scope, ScanContinuation(), files=MagicMock()) + assert not first.final + # Round-trip persisted progress into a newly constructed source. + resumed = await source() + monkeypatch.setattr(resumed, "_execute", execute) + second = await resumed.capture_page( + scope, + ScanContinuation.model_validate_json(first.continuation.model_dump_json()), + files=MagicMock(), + ) + assert not second.final + third = await resumed.capture_page(scope, second.continuation, files=MagicMock()) + assert third.final + assert execute.call_args_list[1].args[1]["until"] == "2026-01-02T00:00:00+00:00" + assert execute.call_args_list[2].args[1]["since"] == "2026-01-02T00:00:00+00:00" + assert third.records[0].identity.native_id == "b" @pytest.mark.asyncio -async def test_unpartitionable_cap_fails(monkeypatch): +@pytest.mark.parametrize( + "rows", + [ + [{"id": "a"}], + [ + {"id": "a", "start": "2026-01-01T00:00:00Z"}, + {"id": "b", "start": "2026-01-01T00:00:00Z"}, + ], + ], +) +async def test_capped_unsplittable_listing_fails_without_final_page(monkeypatch, rows): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + connector = await source() monkeypatch.setattr( connector, - "_list_window", - AsyncMock(return_value=([{"start": "2026-01-01T00:00:00Z"}], True)), + "_execute", + AsyncMock( + return_value={ + "meetings": rows, + "has_more": False, + "truncated": True, + } + ), ) - with pytest.raises(ValueError, match="cap"): - _ = [row async for row in connector._list_all()] + with pytest.raises(ValueError, match="cannot be partitioned safely"): + await connector.capture_page( + CompletedScope(record_type="meeting_listing"), ScanContinuation(), files=MagicMock() + ) @pytest.mark.asyncio -async def test_explicit_tool_failure_never_becomes_partial_success(monkeypatch): +async def test_uncapped_listing_does_not_invent_missing_dates(monkeypatch): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + connector = await source() monkeypatch.setattr( connector, - "_post", - AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": "failed"}]), + "_execute", + AsyncMock( + return_value={ + "meetings": [{"id": "a"}], + "has_more": False, + } + ), ) - with pytest.raises(ValueError, match="tool execution failed; capture is incomplete"): - await connector._meeting("m") + page = await connector.capture_page( + CompletedScope(record_type="meeting_listing"), ScanContinuation(), files=MagicMock() + ) + assert page.final + assert page.records[0].source_created_at is None + assert page.records[0].source_updated_at is None + + +@pytest.mark.asyncio +async def test_short_listing_pages_are_not_limited_to_six_cursors(monkeypatch): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + + connector = await source() + execute = AsyncMock( + side_effect=[ + {"meetings": [{"id": str(i)}], "has_more": i < 8, "next_cursor": str(i + 1)} + for i in range(9) + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + continuation = ScanContinuation() + for _ in range(9): + page = await connector.capture_page( + CompletedScope(record_type="meeting_listing"), continuation, files=MagicMock() + ) + continuation = page.continuation + assert page.final From f8104da49d71903bf377aed23fb02ffdc3d7910f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:36:07 -0700 Subject: [PATCH 069/266] Correct Wispr trial documentation after page-source migration --- backend/tests/live/README.md | 71 +++++++++++++++++++++++++----------- 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 949e7dd3f2..ecc953d5ca 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -76,25 +76,24 @@ 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 sample and current Wispr sampler - -The earlier Slack sample below used the old observation adapter. Slack now uses -durable pages and the sampler explicitly directs it to `provider_lifecycle.py`; -there is no second capture implementation. `LIVE_PROVIDERS=wispr` remains supported, -with `LIVE_WISPR_ACCOUNT_ID` and `LIVE_WISPR_USER_ID` for a previously connected -account. Slack's lifecycle verifies `auth.test` and the matching `users.info` -email against `LIVE_EXPECTED_EMAIL`. Wispr 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. A Wispr-only -run does not require the email variable. - -The historical Slack and current Wispr sampler stop after three distinct records, or -25 observations. No file/blob presence is required. Repeated native identities -are reduced to their latest raw observation in this bounded sample before -capture and replay; Slack history/replies can return the same parent twice. -Declared channel/container parent relationships are retained. Every record's -native JSON is compared after durable readback; identical replay must add no -journal changes. Source markers/checkpoints remain deliberately uncommitted. +## 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 @@ -220,9 +219,11 @@ The earlier smaller-budget failure evidence remains unchanged. Sanitized result: ### Slack and Wispr full lifecycle attempts These historical attempts used the real canonical orchestrator and production -`cursor=None`. Slack's current durable-page trial is described below. Wispr still -has no deletion scope or durable provider checkpoint; its enumeration is not a -delta-resume guarantee. +`cursor=None`. Current adapters use canonical durable page checkpoints. Slack's +interruption trial is described below. Wispr now resumes individual meeting scopes, +but has no provider delta cursor or deletion evidence; its discovery-only traversal +is not an exhaustive-source claim. Its current lifecycle runs first/second cycles; +the deliberate same-job interruption probe is still Slack-specific. Actual network requests are metered, including Wispr session/account operations. The first Slack full-accessible-history attempt stopped on a rate limit after @@ -291,3 +292,29 @@ Private schema, downloaded files, and resume evidence are removed by the parent. resume using synthetic Slack responses and isolated PostgreSQL. Its budget tests use fake child processes to check decreasing limits and unknown killed-child usage. These are harness tests, not evidence of live Slack coverage or throughput. + +### Planned Wispr same-job recovery proof (not yet runnable) + +The current `LIVE_LIFECYCLE_PROVIDER=wispr PYTHONPATH=. .venv/bin/python +tests/live/provider_lifecycle.py` invokes the real page pipeline, but does **not** +prove failure recovery: `run_trials` restarts a failed Wispr run zero times, and +`PageResumeProbe` matches only Slack message scopes with pending replies. Its +current Wispr limits are150 requests/1,000 observations/300 seconds **per process**, +not an aggregate interruption budget. Do not call that command a resume trial. + +A narrowly extended existing probe should interrupt before the next body request +after three meeting child scopes are SQL-verified complete. Persist only private +scan/cycle identity, revision and completed-sibling identity evidence; retain no +meeting content in diagnostic output. A fresh child uses the same job with writer +attempt2, refreshes listing membership, then commits one previously incomplete +meeting. Verify the three prior child revisions remain unchanged and intercept +GET arguments in memory to assert zero calls for those completed identities. End +after this proof, without requiring a complete traversal or starting a new cycle. + +Proposed aggregate guard:20 HTTP requests including account/session/listing/body +calls across both children,600 observed records and180 seconds total; no retries. +A body needing many ranges, changed inventory, rate failure or exhausted budget +may make the proof inconclusive; stop immediately on a rate signal. Reuse the parent's remaining-budget accounting, +private PostgreSQL schema and cleanup on every exit. Add offline process/SQL tests +for the Wispr selector and early-failure path before requesting a live run. No new +command switch or executable trial mode has been implemented for this proposal. From 7438ee6b0a401a753bbfc125a89a8dd425d98ccc Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:40:29 -0700 Subject: [PATCH 070/266] Add bounded cross-process Wispr recovery proof mode --- .../tests/test_wispr_resume_probe.py | 125 ++++++++++++++++++ backend/tests/live/README.md | 53 ++++---- backend/tests/live/provider_lifecycle.py | 39 +++++- backend/tests/live/wispr_resume.py | 108 +++++++++++++++ 4 files changed, 296 insertions(+), 29 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py create mode 100644 backend/tests/live/wispr_resume.py 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..448b631213 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py @@ -0,0 +1,125 @@ +"""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 + +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) + 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) + await driver(service, sessions, fence, BoundedPageSource(source, counters, 600, probe)).run() +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"] == 4 + 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"] == 2 + 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" diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index ecc953d5ca..d90f951511 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -293,28 +293,31 @@ resume using synthetic Slack responses and isolated PostgreSQL. Its budget tests use fake child processes to check decreasing limits and unknown killed-child usage. These are harness tests, not evidence of live Slack coverage or throughput. -### Planned Wispr same-job recovery proof (not yet runnable) - -The current `LIVE_LIFECYCLE_PROVIDER=wispr PYTHONPATH=. .venv/bin/python -tests/live/provider_lifecycle.py` invokes the real page pipeline, but does **not** -prove failure recovery: `run_trials` restarts a failed Wispr run zero times, and -`PageResumeProbe` matches only Slack message scopes with pending replies. Its -current Wispr limits are150 requests/1,000 observations/300 seconds **per process**, -not an aggregate interruption budget. Do not call that command a resume trial. - -A narrowly extended existing probe should interrupt before the next body request -after three meeting child scopes are SQL-verified complete. Persist only private -scan/cycle identity, revision and completed-sibling identity evidence; retain no -meeting content in diagnostic output. A fresh child uses the same job with writer -attempt2, refreshes listing membership, then commits one previously incomplete -meeting. Verify the three prior child revisions remain unchanged and intercept -GET arguments in memory to assert zero calls for those completed identities. End -after this proof, without requiring a complete traversal or starting a new cycle. - -Proposed aggregate guard:20 HTTP requests including account/session/listing/body -calls across both children,600 observed records and180 seconds total; no retries. -A body needing many ranges, changed inventory, rate failure or exhausted budget -may make the proof inconclusive; stop immediately on a rate signal. Reuse the parent's remaining-budget accounting, -private PostgreSQL schema and cleanup on every exit. Add offline process/SQL tests -for the Wispr selector and early-failure path before requesting a live run. No new -command switch or executable trial mode has been implemented for this proposal. +### Opt-in Wispr same-job recovery proof + +From `backend`, with the same private PostgreSQL and existing verified account environment: + +```sh +LIVE_LIFECYCLE_PROVIDER=wispr LIVE_WISPR_RESUME=1 PYTHONPATH=. .venv/bin/python tests/live/provider_lifecycle.py +``` + +Required secret environment names are `COMPOSIO_API_KEY`, `LIVE_WISPR_ACCOUNT_ID` +and `LIVE_WISPR_USER_ID`; PostgreSQL uses `CANONICAL_TEST_DATABASE_URL`. Do not +print these values. This mode shares a hard20 HTTP request,600 observed record, +180 second budget across both processes, including account/session/listing calls. +No retries or new account grants occur. Errors, including rate failures, stop the trial. + +The first process exits before another body call after SQL verifies three completed +meeting children. The second uses the same job/cycle with a new writer attempt, +refreshes listing membership and commits one additional body. Before each body call, +the typed probe rejects a completed sibling identity. It verifies prior scan revisions +remain unchanged and then exits intentionally. Private evidence contains only scan +identities and native-ID SHA256 digests, never meeting content; output contains counts +and booleans. Parent cleanup removes the private schema/files on all normal or failed +exits. `wispr_recovery_verified` is distinct from full traversal or deletion completeness. + +A source with fewer than five discoverable meetings, many body ranges, changed inventory, +rate failure or exhausted budget can leave the proof inconclusive. Without the opt-in +flag, Wispr retains its ordinary first/second-cycle trial; that is not a recovery proof. +Offline PostgreSQL/subprocess tests verify the actual process exits and fresh retry, +failed-body stop, and aggregate budget. No live recovery result is claimed here. diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index c1b4e0a5f9..f25bf25765 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -218,8 +218,10 @@ def prepare_resume_probe(manifest, sessions, organization_id, sync_id, job_id, a if manifest.get("resume_stage") not in {"interrupt", "resume"}: return None from capture_resume import PageResumeProbe + from wispr_resume import WisprResumeProbe - return PageResumeProbe( + probe_type = WisprResumeProbe if manifest.get("wispr_resume") else PageResumeProbe + return probe_type( sessions, organization_id, sync_id, @@ -568,8 +570,13 @@ async def run_trials(manifest, evidence_reader=None): path = Path(manifest["root"]) / "manifest.json" results = [] name = manifest["provider"] - resume_trial = name == "slack" - stages = ("interrupt", "resume", "new_cycle") if resume_trial else ("first", "second") + wispr_resume = manifest.get("wispr_resume", False) + resume_trial = name == "slack" or wispr_resume + stages = ( + (("interrupt", "resume") if wispr_resume else ("interrupt", "resume", "new_cycle")) + if resume_trial + else ("first", "second") + ) total_requests, total_records = manifest["request_limit"], manifest["record_limit"] deadline = time.monotonic() + manifest["timeout"] if resume_trial else None interrupted_job_id = str(uuid4()) @@ -613,6 +620,13 @@ async def run_trials(manifest, evidence_reader=None): if stage == "interrupt": if code != 75 or not result.get("intentional_interruption"): break + elif ( + wispr_resume + and stage == "resume" + and code == 76 + and result.get("wispr_recovery_verified") + ): + break elif code: break return results @@ -636,6 +650,9 @@ async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") if name not in {"gmail", "google_calendar", "google_drive", "slack", "wispr"}: raise ValueError("Unsupported lifecycle provider") + wispr_resume = os.environ.get("LIVE_WISPR_RESUME") == "1" + if wispr_resume and name != "wispr": + raise ValueError("Wispr resume requires Wispr provider") free_disk_bytes = check_free_disk(name) url = harness.test_database_url() schema = "canonical_live_" + uuid4().hex @@ -709,6 +726,8 @@ async def main(): "sync_id": str(sync_id), "query": f"after:{end - 7 * 86400} before:{end}", } + if wispr_resume: + manifest.update(wispr_resume=True, request_limit=20, record_limit=600, timeout=180) if name == "google_calendar": now = datetime.now(timezone.utc).replace(microsecond=0) manifest["calendar_config"] = { @@ -756,6 +775,11 @@ async def resume_evidence(): "schema_removed": True, "blob_directory_removed": True, "resume_verified": bool(resume_verified), + "wispr_recovery_verified": bool( + wispr_resume + and len(results) == 2 + and results[1].get("wispr_recovery_verified", False) + ), "resumed_page_committed": any( item.get("resumed_page_committed", False) for item in results ), @@ -770,7 +794,14 @@ async def resume_evidence(): } ) ) - success = resume_verified if name == "slack" else len(results) == 2 + wispr_verified = ( + wispr_resume and len(results) == 2 and results[1].get("wispr_recovery_verified", False) + ) + success = ( + wispr_verified + if wispr_resume + else (resume_verified if name == "slack" else len(results) == 2) + ) return 0 if success and not any(item["failed"] for item in results) else 1 diff --git a/backend/tests/live/wispr_resume.py b/backend/tests/live/wispr_resume.py new file mode 100644 index 0000000000..d87bf11dc3 --- /dev/null +++ b/backend/tests/live/wispr_resume.py @@ -0,0 +1,108 @@ +"""Test-only committed-body interruption; no provider content in evidence or output.""" + +import hashlib +import json +import os +from uuid import UUID + +from pydantic import BaseModel, ConfigDict +from sqlalchemy import select + +from airweave.models.capture_scan import CaptureScan +from airweave.models.entity import Entity +from airweave.models.sync import Sync +from airweave.models.sync_job import SyncJob + + +class CompletedBody(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + scan_id: UUID + revision: int + native_digest: str + + +class WisprTarget(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + bodies: tuple[CompletedBody, ...] + + +class WisprResumeProbe: + """Stop before another call, only after SQL proves prior bodies committed.""" + + resumed_saved_page = False + saved_cursor_expired = False + + def __init__( + self, sessions, organization_id, sync_id, job_id, attempt_id, target_path, mode, counters + ): + self.sessions = sessions + self.organization_id, self.sync_id = organization_id, sync_id + self.job_id, self.attempt_id = job_id, attempt_id + self.target_path, self.mode, self.counters = target_path, mode, counters + self.target = ( + WisprTarget.model_validate_json(target_path.read_text()) if mode == "resume" else None + ) + + async def before_page(self, scope, continuation): + if scope.record_type != "meeting": + return + async with self.sessions() as db: + sync, job = await db.get(Sync, self.sync_id), await db.get(SyncJob, self.job_id) + assert sync.organization_id == job.organization_id == self.organization_id + assert sync.writer_job_id == job.id and sync.writer_attempt_id == self.attempt_id + assert job.status == "running" + rows = ( + await db.execute( + select(CaptureScan, Entity.native_id) + .join(Entity, Entity.id == CaptureScan.parent_record_id) + .where( + CaptureScan.organization_id == self.organization_id, + CaptureScan.sync_id == self.sync_id, + CaptureScan.record_type == "meeting", + CaptureScan.phase == "complete", + ) + ) + ).all() + bodies = tuple( + CompletedBody( + scan_id=row.id, + revision=row.revision, + native_digest=hashlib.sha256(native.encode()).hexdigest(), + ) + for row, native in rows + ) + cycles = {row.cycle_id for row, _ in rows} + if self.target is not None: + assert cycles == {self.target.cycle_id} + assert all(body in bodies for body in self.target.bodies), "Completed sibling changed" + digest = hashlib.sha256(scope.parent.native_id.encode()).hexdigest() + assert digest not in {body.native_digest for body in self.target.bodies}, ( + "Completed body repeated" + ) + if len(bodies) == len(self.target.bodies) + 1: + self._stop(True) + elif len(bodies) == 3: + assert len(cycles) == 1 + self.target_path.write_text( + WisprTarget(cycle_id=next(iter(cycles)), bodies=bodies).model_dump_json() + ) + os.chmod(self.target_path, 0o600) + self._stop(False) + + def _stop(self, verified): + print( + json.dumps( + { + "failed": False, + "counters_complete": True, + "intentional_interruption": not verified, + "wispr_recovery_verified": verified, + "full_scope_completed": False, + "completed_bodies": 4 if verified else 3, + **self.counters, + } + ), + flush=True, + ) + os._exit(76 if verified else 75) From 0dce24a6459e39926e00806edd63d827a1ed4fd0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:45:09 -0700 Subject: [PATCH 071/266] Admit verified discovery roots atomically with canonical scan pages --- .../domains/entities/canonical/page_source.py | 5 + .../domains/entities/canonical/scan_models.py | 5 + .../domains/entities/canonical/scan_store.py | 28 ++ .../domains/entities/canonical/store.py | 14 +- .../tests/test_discovered_records.py | 281 ++++++++++++++++++ .../domains/sync_pipeline/canonical_scan.py | 3 +- 6 files changed, 332 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_discovered_records.py diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index 0a288c3330..a62526c60e 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -20,6 +20,11 @@ class CapturePage(BaseModel): 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", + ) continuation: ScanContinuation final: bool = False diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index bda883c8ee..3848e49e96 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -73,6 +73,11 @@ class CommitScanPage(BaseModel): cycle_id: UUID expected: ScanVersion 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", + ) continuation: ScanContinuation final: bool = False diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 9ed83e4443..c83c69447f 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -253,12 +253,40 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: expected_parent = source_record(parent).identity if parent is not None else None if any(record.parent != expected_parent for record in request.records): raise ScanConflict("Page records must retain their exact declared parent identity") + combined = (*request.records, *request.discovered_records) + identities = { + (record.identity.record_type, record.identity.entity_key) for record in combined + } + if len(combined) > 500 or len(identities) != len(combined): + raise ScanConflict("A page must contain at most 500 distinct original identities") + if any( + record.identity.record_type not in cycle.configuration.root_record_types + or cycle.configuration.policy(record.identity.record_type) == "exhaustive" + or record.parent is not None + or record.identity.container_id is not None + or record.kind != "upsert" + or record.allow_reparent + or not record.payload + for record in request.discovered_records + ): + raise ScanConflict("Discovered originals must be verified independent declared roots") captured = await self.records._capture_locked( db, sync, CaptureBatch(fence=request.fence, records=request.records), seen_id=row.sweep_id, ) + discovered = await self.records._capture_locked( + db, + sync, + CaptureBatch(fence=request.fence, records=request.discovered_records), + mark_seen=False, + ) + captured = CaptureResult( + changes=(*captured.changes, *discovered.changes), + sequence=discovered.sequence, + unchanged=captured.unchanged + discovered.unchanged, + ) row.continuation = request.continuation.value row.revision += 1 if request.final: diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 0034c3adbf..0af5036e15 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -337,7 +337,13 @@ async def _parent_attestation( return parent.visibility_epoch async def _capture_locked( - self, db: AsyncSession, sync: Sync, batch: CaptureBatch, *, seen_id: UUID | None = None + self, + db: AsyncSession, + sync: Sync, + batch: CaptureBatch, + *, + seen_id: UUID | None = None, + mark_seen: bool = True, ) -> CaptureResult: changes = [] unchanged = 0 @@ -395,7 +401,8 @@ async def _capture_locked( and not attestation_changed and not (observation.kind == "upsert" and not was_available) ): - entity.last_seen_run_id = seen_id or batch.fence.attempt_id + if mark_seen: + entity.last_seen_run_id = seen_id or batch.fence.attempt_id entity.observed_at = observation.observed_at unchanged += 1 continue @@ -438,7 +445,8 @@ async def _capture_locked( entity.removal_reason = observation.removal_reason entity.completeness = observation.completeness entity.blob_references = [blob.model_dump(mode="json") for blob in observation.blobs] - entity.last_seen_run_id = seen_id or batch.fence.attempt_id + if mark_seen: + entity.last_seen_run_id = seen_id or batch.fence.attempt_id entity.projection_error = None # Flush before reading default/index state and before another observation of this ID. await db.flush() diff --git a/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py b/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py new file mode 100644 index 0000000000..6625910a9b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py @@ -0,0 +1,281 @@ +"""Side discovery shares the page transaction, never another scope's sightings.""" + +from datetime import datetime, timezone +from uuid import uuid4 + +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 +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.store import StaleWriter +from airweave.domains.entities.canonical.tests.helpers import observation +from airweave.models.entity import Entity +from airweave.models.entity_change import EntityChange + + +def original(kind, native): + return observation(native, identity=RecordIdentity(record_type=kind, native_id=native)) + + +async def begin(database, service, fence, cycle, scope): + async with database() as db: + return await service.begin_scan( + db, + BeginScan( + fence=fence, + cycle_id=cycle.version.cycle_id, + scope=scope, + fingerprint=cycle.configuration.fingerprint, + ), + ) + + +async def page(database, service, fence, state, records=(), discovered=()): + 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, + discovered_records=discovered, + continuation=ScanContinuation(value={"cursor": "after"}), + final=True, + ), + ) + + +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), + ), + ) + + +async def setup(database, service, fence): + config = CycleConfiguration( + fingerprint="f" * 64, + parents={"page": (None,), "database": (None,), "block": ("page",)}, + completion_policies={"page": "discovery_only", "database": "discovery_only"}, + ) + async with database() as db: + cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + state = await begin(database, service, fence, cycle, CompletedScope(record_type="page")) + captured = await page(database, service, fence, state, (original("page", "old"),)) + await finish(database, service, fence, captured.state) + old = captured.capture.changes[0].record + state = await begin(database, service, fence, cycle, CompletedScope(record_type="database")) + return cycle, state, old + + +async def test_completed_root_inventory_admits_new_frontier_without_changing_sightings( + database, source +): + service, fence = source + cycle, state, old = await setup(database, service, fence) + async with database() as db: + prior = (await db.get(Entity, old.id)).last_seen_run_id + result = await page( + database, + service, + fence, + state, + (original("database", "db"),), + (original("page", "old"), original("page", "new")), + ) + assert len(result.capture.changes) == 2 and result.capture.unchanged == 1 + await finish(database, service, fence, result.state) + child = await begin( + database, service, fence, cycle, CompletedScope(record_type="block", parent=old.identity) + ) + child = await page(database, service, fence, child) + await finish(database, service, fence, child.state) + async with database() as db: + assert (await db.get(Entity, old.id)).last_seen_run_id == prior + new = await db.scalar(select(Entity).where(Entity.native_id == "new")) + assert new.last_seen_run_id is None + work = await service.next_scope_work(db, fence, cycle.version.cycle_id) + assert work.record_type == "block" and work.parent.identity.native_id == "new" + with pytest.raises(ScanConflict): + await page(database, service, fence, state, discovered=(original("page", "again"),)) + + +@pytest.mark.parametrize( + "invalid", ["duplicate", "parent", "undeclared", "empty", "oversized", "promotion"] +) +async def test_rejected_side_discovery_preserves_page_cursor_and_journal(database, source, invalid): + service, fence = source + _, state, old = await setup(database, service, fence) + main = original("database", "db") + discovered = original("page", "new") + if invalid == "duplicate": + discovered = main + if invalid == "parent": + discovered = discovered.model_copy(update={"parent": old.identity}) + if invalid == "undeclared": + discovered = original("block", "new") + if invalid == "empty": + discovered = discovered.model_copy(update={"payload": {}}) + if invalid == "promotion": + discovered = discovered.model_copy(update={"allow_reparent": True}) + side = ( + tuple(original("page", str(i)) for i in range(500)) + if invalid == "oversized" + else (discovered,) + ) + with pytest.raises(ScanConflict): + await page(database, service, fence, state, (main,), side) + async with database() as db: + current = await service.read_scan(db, fence, state.scope) + assert current.version == state.version and current.continuation == state.continuation + assert await db.scalar(select(Entity).where(Entity.native_id == "db")) is None + assert await db.scalar(select(EntityChange).where(EntityChange.sequence > 1)) is None + + +async def test_new_writer_fences_entire_discovery_commit(database, source): + service, fence = source + _, state, _ = await setup(database, service, 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, + ) + with pytest.raises(StaleWriter): + await page( + database, + service, + fence, + state, + (original("database", "db"),), + (original("page", "new"),), + ) + async with database() as db: + assert await db.scalar(select(Entity).where(Entity.native_id == "new")) is None + assert (await service.read_scan(db, newer, state.scope)).version == state.version + + +async def test_existing_child_cannot_be_promoted_and_rolls_back_scoped_write(database, source): + from airweave.domains.entities.canonical.store import CanonicalStoreError + from airweave.domains.entities.canonical.tests.helpers import capture + + service, fence = source + _, state, old = await setup(database, service, fence) + # A prior connector mapping stored this native object beneath another page. + child = original("page", "child").model_copy(update={"parent": old.identity}) + await capture(database, service, fence, child) + with pytest.raises(CanonicalStoreError, match="parent"): + await page( + database, + service, + fence, + state, + (original("database", "db"),), + (original("page", "child"),), + ) + async with database() as db: + assert await db.scalar(select(Entity).where(Entity.native_id == "db")) is None + assert (await service.read_scan(db, fence, state.scope)).version == state.version + assert await db.scalar(select(EntityChange).where(EntityChange.sequence > 2)) is None + + +async def test_parent_withdrawn_after_io_blocks_scoped_and_discovered_records(database, source): + from airweave.domains.entities.canonical.cycle_store import CycleConflict + from airweave.domains.entities.canonical.tests.helpers import capture + + service, fence = source + cycle, root, old = await setup(database, service, fence) + final = await page(database, service, fence, root) + await finish(database, service, fence, final.state) + state = await begin( + database, service, fence, cycle, CompletedScope(record_type="block", parent=old.identity) + ) + await capture( + database, + service, + fence, + original("page", "old").model_copy( + update={"kind": "delete", "removal_reason": "access_revoked"} + ), + ) + body = original("block", "body").model_copy(update={"parent": old.identity}) + with pytest.raises(CycleConflict): + await page(database, service, fence, state, (body,), (original("page", "new"),)) + async with database() as db: + assert await db.scalar(select(Entity).where(Entity.native_id.in_(["new", "body"]))) is None + assert (await service.read_scan(db, fence, state.scope)).version == state.version + + +async def test_exhaustive_target_root_rejects_side_admission(database, source): + service, fence = source + config = CycleConfiguration( + fingerprint="f" * 64, + parents={"page": (None,), "database": (None,)}, + completion_policies={"database": "discovery_only"}, + ) + async with database() as db: + cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + state = await begin(database, service, fence, cycle, CompletedScope(record_type="database")) + with pytest.raises(ScanConflict, match="independent declared roots"): + await page(database, service, fence, state, discovered=(original("page", "new"),)) + + +async def test_side_revival_invalidates_previously_completed_child_scope(database, source): + from airweave.domains.entities.canonical.tests.helpers import capture + + service, fence = source + cycle, root, old = await setup(database, service, fence) + final = await page(database, service, fence, root) + final = await finish(database, service, fence, final.state) + child = await begin( + database, service, fence, cycle, CompletedScope(record_type="block", parent=old.identity) + ) + child = await page(database, service, fence, child) + await finish(database, service, fence, child.state) + await capture( + database, + service, + fence, + original("page", "old").model_copy( + update={"kind": "delete", "removal_reason": "scope_removed"} + ), + ) + async with database() as db: + restarted = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=root.scope, + cycle_id=root.cycle_id, + fingerprint=root.fingerprint, + expected=final.state.version, + restart=True, + ), + ) + revived = await page(database, service, fence, restarted, discovered=(original("page", "old"),)) + await finish(database, service, fence, revived.state) + async with database() as db: + work = await service.next_scope_work(db, fence, cycle.version.cycle_id) + assert work.parent.id == old.id + assert work.parent_visibility_epoch > child.state.parent_visibility_epoch + assert work.record_type == "block" diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 21e2b32ee5..02ca0aa103 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -168,12 +168,13 @@ async def scan( cycle_id=state.cycle_id, expected=state.version, records=tuple(self._parented(r, scope) for r in page.records), + discovered_records=page.discovered_records, continuation=page.continuation, final=page.final, ), ) state = result.state - await self.progress(result.capture, page.records) + await self.progress(result.capture, (*page.records, *page.discovered_records)) if state.phase == "reconciling": if state.completion_policy == "discovery_with_validation": state = await self.refresh_known(state) From dff605a662bb20b7969e4f9d21a4cff314f3dc63 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:46:30 -0700 Subject: [PATCH 072/266] Diagnose Wispr setup failures without retaining private errors --- .../canonical/tests/test_wispr_diagnostics.py | 85 +++++++++++++++++++ backend/tests/live/README.md | 2 + backend/tests/live/provider_lifecycle.py | 37 +++++++- backend/tests/live/provider_sample.py | 15 +++- backend/tests/live/wispr_diagnostics.py | 69 +++++++++++++++ 5 files changed, 204 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py create mode 100644 backend/tests/live/wispr_diagnostics.py 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..d9b1fb669e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py @@ -0,0 +1,85 @@ +"""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 + +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 as exc: + assert str(exc) == "Wispr tool execution failed; capture is incomplete" + 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/tests/live/README.md b/backend/tests/live/README.md index d90f951511..c15684b671 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -305,6 +305,8 @@ Required secret environment names are `COMPOSIO_API_KEY`, `LIVE_WISPR_ACCOUNT_ID and `LIVE_WISPR_USER_ID`; PostgreSQL uses `CANONICAL_TEST_DATABASE_URL`. Do not print these values. This mode shares a hard20 HTTP request,600 observed record, 180 second budget across both processes, including account/session/listing calls. +Set `LIVE_WISPR_REQUEST_LIMIT=17` to reduce the aggregate request budget; values +from1 through20 are validated before setup/network. The default remains20. No retries or new account grants occur. Errors, including rate failures, stop the trial. The first process exits before another body call after SQL verifies three completed diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index f25bf25765..eb0386a95d 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -249,6 +249,11 @@ async def child(manifest): is_calendar = name == "google_calendar" counters["sync_token_requests"] = 0 counters["resumed_changes_requests"] = 0 + from functools import partial + + from wispr_diagnostics import WisprDiagnostics, observe_request + + wispr_diagnostics = WisprDiagnostics() last_operation = "identity" rate_limits = [] resume_probe = None @@ -345,7 +350,14 @@ async def db_context(): "wispr": "LIVE_WISPR_ACCOUNT_ID", }[name] connection = ( - wispr_source(os.environ[account_variable], key, fence, request_hook=request_hook) + wispr_source( + os.environ[account_variable], + key, + fence, + request_hook=partial(observe_request, wispr_diagnostics, request_hook), + response_hook=wispr_diagnostics.response, + envelope_hook=wispr_diagnostics.envelope, + ) if name == "wispr" else rest_source( name, @@ -512,7 +524,8 @@ async def db_context(): "saved_cursor_expired": resume_probe.saved_cursor_expired if resume_probe else False, "error_type": type(error).__name__, "safe_reason": safe_failure_reason(error), - "last_operation": last_operation, + "last_operation": wispr_diagnostics.operation if name == "wispr" else last_operation, + "wispr_diagnostics": wispr_diagnostics.model_dump() if name == "wispr" else None, **counters, "rate_limits": rate_limits, "checkpoint_unchanged": saved == (previous or None), @@ -646,6 +659,17 @@ def trial_counter_summary(results): } +def wispr_request_limit(value: str) -> int: + """Validate an explicitly reduced trial budget before setup or network access.""" + try: + limit = int(value) + except ValueError: + raise ValueError("Wispr request limit must be an integer from 1 to 20") from None + if not 1 <= limit <= 20: + raise ValueError("Wispr request limit must be an integer from 1 to 20") + return limit + + async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") if name not in {"gmail", "google_calendar", "google_drive", "slack", "wispr"}: @@ -653,6 +677,11 @@ async def main(): wispr_resume = os.environ.get("LIVE_WISPR_RESUME") == "1" if wispr_resume and name != "wispr": raise ValueError("Wispr resume requires Wispr provider") + request_limit = ( + wispr_request_limit(os.environ.get("LIVE_WISPR_REQUEST_LIMIT", "20")) + if wispr_resume + else None + ) free_disk_bytes = check_free_disk(name) url = harness.test_database_url() schema = "canonical_live_" + uuid4().hex @@ -727,7 +756,9 @@ async def main(): "query": f"after:{end - 7 * 86400} before:{end}", } if wispr_resume: - manifest.update(wispr_resume=True, request_limit=20, record_limit=600, timeout=180) + manifest.update( + wispr_resume=True, request_limit=request_limit, record_limit=600, timeout=180 + ) if name == "google_calendar": now = datetime.now(timezone.utc).replace(microsecond=0) manifest["calendar_config"] = { diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index 3eb62922fa..222ff69cb5 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -95,10 +95,14 @@ async def rest_source( @asynccontextmanager -async def wispr_source(account, key, fence, *, request_hook=None): +async def wispr_source( + account, key, fence, *, request_hook=None, response_hook=None, envelope_hook=None +): """Verify exact Composio account binding; Wispr exposes no identity profile here.""" expected_user = os.environ["LIVE_WISPR_USER_ID"] hooks = {"request": [request_hook]} if request_hook is not None else {} + if response_hook is not None: + hooks["response"] = [response_hook] async with httpx.AsyncClient(timeout=180, event_hooks=hooks) as client: metadata_response = await client.get( "https://backend.composio.dev/api/v3/connected_accounts/" + quote(account, safe=""), @@ -122,5 +126,14 @@ async def wispr_source(account, key, fence, *, request_hook=None): client, fence.organization_id, "wispr", feature_flag_enabled=False ), ) + if envelope_hook is not None: + original_post = source._post + + async def observed_post(path, body): + result = await original_post(path, body) + envelope_hook(result) + return result + + source._post = observed_post await source.validate() yield source, "composio_account_principal_only" diff --git a/backend/tests/live/wispr_diagnostics.py b/backend/tests/live/wispr_diagnostics.py new file mode 100644 index 0000000000..5e62f36d80 --- /dev/null +++ b/backend/tests/live/wispr_diagnostics.py @@ -0,0 +1,69 @@ +"""Allowlisted diagnostic metadata; never retain provider error strings or arguments.""" + +import json +from typing import Literal + +from pydantic import BaseModel + + +class WisprDiagnostics(BaseModel): + operation: Literal["metadata", "session", "search", "get", "unknown"] = "unknown" + last_http_status: int | None = None + tool_error_present: bool | None = None + tool_error_kind: Literal["string", "object", "array", "boolean", "number", "none"] | None = None + rate_signal_detected: bool = False + + def request(self, request): + self.last_http_status = None + self.tool_error_present = self.tool_error_kind = None + self.rate_signal_detected = False + path = request.url.path + if "/connected_accounts/" in path: + self.operation = "metadata" + elif path.endswith("/tool_router/session"): + self.operation = "session" + elif path.endswith("/execute"): + try: + value = json.loads(request.content) + except (ValueError, UnicodeDecodeError): + value = None + slug = value.get("tool_slug") if isinstance(value, dict) else None + self.operation = { + "WISPR_FLOW_MCP_SEARCH_MEETINGS": "search", + "WISPR_FLOW_MCP_GET_MEETING": "get", + }.get(slug, "unknown") + else: + self.operation = "unknown" + + async def response(self, response): + self.last_http_status = response.status_code + self.rate_signal_detected = response.status_code == 429 + + def envelope(self, value): + if self.operation not in {"search", "get"}: + return + error = value.get("error") + self.tool_error_present = error is not None + if error is None: + self.tool_error_kind = "none" + elif isinstance(error, str): + self.tool_error_kind = "string" + lowered = error.lower() + self.rate_signal_detected |= "rate limit" in lowered or "too many requests" in lowered + elif isinstance(error, dict): + self.tool_error_kind = "object" + self.rate_signal_detected |= any( + type(error.get(key)) is int and error[key] == 429 + for key in ("code", "status", "status_code") + ) + elif isinstance(error, list): + self.tool_error_kind = "array" + elif isinstance(error, bool): + self.tool_error_kind = "boolean" + else: + self.tool_error_kind = "number" + + +async def observe_request(diagnostics, meter, request): + await meter(request) + diagnostics.request(request) From 21958fa0c6737128c4ac670bb1158f74e5aba741 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:46:41 -0700 Subject: [PATCH 073/266] Allow account-bound Notion API proxy transport --- .../auth_provider/providers/composio.py | 1 + .../platform/http_client/test_notion_proxy.py | 136 ++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 backend/tests/unit/platform/http_client/test_notion_proxy.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 8bf2a4bdb0..fb4c61877b 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -118,6 +118,7 @@ async def get_auth_result( "slack": frozenset({"slack.com"}), "linear": frozenset({"api.linear.app", "uploads.linear.app"}), "github": frozenset({"api.github.com"}), + "notion": frozenset({"api.notion.com"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( diff --git a/backend/tests/unit/platform/http_client/test_notion_proxy.py b/backend/tests/unit/platform/http_client/test_notion_proxy.py new file mode 100644 index 0000000000..6b20c981bc --- /dev/null +++ b/backend/tests/unit/platform/http_client/test_notion_proxy.py @@ -0,0 +1,136 @@ +"""Notion uses the existing account-bound proxy without exposing provider tokens.""" + +import json +from unittest.mock import AsyncMock + +import httpx +import pytest + +from airweave.domains.auth_provider.exceptions import AuthProviderConfigError +from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider +from airweave.platform.http_client.composio_transport import ComposioProxyError, ComposioTransport + + +@pytest.mark.asyncio +async def test_notion_binding_preserves_version_body_and_native_error(monkeypatch): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "project-secret"}, + config={"account_id": "ca_notion", "auth_config_id": "ac_notion"}, + ) + fetch = AsyncMock( + return_value={ + "toolkit": {"slug": "notion"}, + "status": "ACTIVE", + "auth_config": {"id": "ac_notion"}, + } + ) + monkeypatch.setattr(provider, "_get_with_auth", fetch) + result = await provider.get_auth_result("notion", ["access_token"]) + assert result.credentials is None + auth = result.managed_auth + assert auth.allowed_hosts == frozenset({"api.notion.com"}) + assert fetch.call_args.args[1].endswith("/connected_accounts/ca_notion") + body = {"filter": {"value": "page", "property": "object"}, "start_cursor": "opaque"} + error = { + "object": "error", + "status": 400, + "code": "validation_error", + "message": "Invalid cursor", + } + + async def proxy(request): + payload = json.loads(request.content) + assert payload["connected_account_id"] == "ca_notion" + assert payload["endpoint"] == "https://api.notion.com/v1/search" + assert payload["method"] == "POST" + assert payload["body"] == body + headers = {item["name"]: item["value"] for item in payload["parameters"]} + assert headers["notion-version"] == "2026-03-11" + assert headers["content-type"] == "application/json" + assert not {"authorization", "cookie", "x-api-key"}.intersection(headers) + assert "provider-secret" not in request.content.decode() + assert request.headers["x-api-key"] == "project-secret" + return httpx.Response(200, json={"status": 400, "data": error}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key=auth.api_key.get_secret_value(), + connected_account_id=auth.connected_account_id, + allowed_hosts=auth.allowed_hosts, + ) + ) as client: + response = await client.post( + "https://api.notion.com/v1/search", + json=body, + headers={ + "Notion-Version": "2026-03-11", + "Authorization": "Bearer provider-secret", + "Cookie": "provider-secret", + "x-api-key": "provider-secret", + }, + ) + assert response.status_code == 400 + assert response.json() == error + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "https://prod-files-secure.s3.us-west-2.amazonaws.com/file", + "https://api.notion.com.evil.example/v1/search", + "https://www.notion.so/page", + "http://api.notion.com/v1/search", + "https://api.notion.com:444/v1/search", + "https://user@api.notion.com/v1/search", + ], +) +async def test_notion_rejects_unapproved_destinations_before_proxy(url): + proxy = AsyncMock() + async with httpx.AsyncClient( + transport=ComposioTransport( + client=proxy, + api_key="secret", + connected_account_id="ca_notion", + allowed_hosts={"api.notion.com"}, + ) + ) as client: + with pytest.raises(ComposioProxyError): + await client.get(url) + proxy.stream.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "account,toolkit,config_id,reason", + [ + (None, "notion", "ac_notion", "explicit Composio account"), + ("ca_notion", "github", "ac_notion", "toolkit"), + ("ca_notion", "notion", "other", "auth config"), + ], +) +async def test_notion_rejects_missing_account_or_mismatched_binding( + monkeypatch, + account, + toolkit, + config_id, + reason, +): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "secret"}, + config={"account_id": account, "auth_config_id": "ac_notion"}, + ) + fetch = AsyncMock( + return_value={ + "toolkit": {"slug": toolkit}, + "status": "ACTIVE", + "auth_config": {"id": config_id}, + } + ) + monkeypatch.setattr(provider, "_get_with_auth", fetch) + with pytest.raises(AuthProviderConfigError, match=reason): + await provider.get_auth_result("notion", ["access_token"]) + if account is None: + fetch.assert_not_called() From 2479e0eeb9d8c9f04db79dc7abf2526a446824e4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:48:48 -0700 Subject: [PATCH 074/266] Verify managed account and explicit user identities --- .../auth_provider/providers/composio.py | 8 ++ .../tests/test_composio_binding.py | 110 ++++++++++++++++++ .../http_client/test_composio_transport.py | 4 +- .../platform/http_client/test_notion_proxy.py | 2 + .../platform/sources/test_github_capture.py | 1 + .../platform/sources/test_linear_capture.py | 1 + 6 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/domains/auth_provider/tests/test_composio_binding.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index fb4c61877b..3409abe9d0 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -135,6 +135,14 @@ async def get_auth_result( "https://backend.composio.dev/api/v3/connected_accounts/" + quote(self.account_id, safe=""), ) + if account.get("id") != self.account_id: + raise AuthProviderConfigError( + "Connected account identity does not match", provider_name="composio" + ) + if self.user_id is not None and account.get("user_id") != self.user_id: + raise AuthProviderConfigError( + "Connected account user does not match", provider_name="composio" + ) if account.get("toolkit", {}).get("slug", "").lower() != self._get_composio_slug( source_short_name ): diff --git a/backend/airweave/domains/auth_provider/tests/test_composio_binding.py b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py new file mode 100644 index 0000000000..fc32a2ede8 --- /dev/null +++ b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py @@ -0,0 +1,110 @@ +"""Managed access requires the selected account and any explicit user binding.""" + +from unittest.mock import AsyncMock + +import pytest + +from airweave.domains.auth_provider.exceptions import AuthProviderConfigError +from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source,toolkit", + [ + ("gmail", "gmail"), + ("google_calendar", "googlecalendar"), + ("google_drive", "googledrive"), + ("slack", "slack"), + ("linear", "linear"), + ("github", "github"), + ("notion", "notion"), + ("wispr", "wispr_flow_mcp"), + ], +) +@pytest.mark.parametrize( + "field,value", + [ + ("id", "wrong-account"), + ("id", None), + ("user_id", "wrong-user"), + ("user_id", None), + ], +) +async def test_managed_binding_rejects_wrong_or_missing_identity( + monkeypatch, + source, + toolkit, + field, + value, +): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "secret"}, + config={"account_id": "selected", "user_id": "owner"}, + ) + account = { + "id": "selected", + "user_id": "owner", + "toolkit": {"slug": toolkit}, + "status": "ACTIVE", + } + if value is None: + account.pop(field) + else: + account[field] = value + monkeypatch.setattr(provider, "_get_with_auth", AsyncMock(return_value=account)) + with pytest.raises(AuthProviderConfigError, match="identity" if field == "id" else "user"): + await provider.get_auth_result(source, []) + assert provider._last_credential_blob is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "source,user_id", + [ + ("notion", None), + ("notion", "owner"), + ("wispr", "owner"), + ], +) +async def test_managed_binding_accepts_exact_identity_without_requiring_optional_user( + monkeypatch, + source, + user_id, +): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "secret"}, + config={"account_id": "selected", "user_id": user_id}, + ) + account = { + "id": "selected", + "user_id": "owner", + "toolkit": {"slug": "wispr_flow_mcp" if source == "wispr" else source}, + "status": "ACTIVE", + } + monkeypatch.setattr(provider, "_get_with_auth", AsyncMock(return_value=account)) + result = await provider.get_auth_result(source, []) + assert result.credentials is None + assert result.managed_auth.connected_account_id == "selected" + + +@pytest.mark.asyncio +async def test_wispr_still_requires_configured_user(monkeypatch): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "secret"}, + config={"account_id": "selected"}, + ) + monkeypatch.setattr( + provider, + "_get_with_auth", + AsyncMock( + return_value={ + "id": "selected", + "user_id": "owner", + "toolkit": {"slug": "wispr_flow_mcp"}, + "status": "ACTIVE", + } + ), + ) + with pytest.raises(AuthProviderConfigError, match="explicit Composio user"): + await provider.get_auth_result("wispr", []) diff --git a/backend/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py index 506b6546c3..67cd97e0c5 100644 --- a/backend/tests/unit/platform/http_client/test_composio_transport.py +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -241,7 +241,9 @@ async def test_managed_provider_account_binding(monkeypatch): provider = await ComposioAuthProvider.create( credentials={"api_key": "key"}, config={"account_id": "ca_bound"} ) - fetch = AsyncMock(return_value={"toolkit": {"slug": "gmail"}, "status": "ACTIVE"}) + fetch = AsyncMock( + return_value={"id": "ca_bound", "toolkit": {"slug": "gmail"}, "status": "ACTIVE"} + ) monkeypatch.setattr(provider, "_get_with_auth", fetch) result = await provider.get_auth_result("gmail", ["access_token"]) assert result.credentials is None diff --git a/backend/tests/unit/platform/http_client/test_notion_proxy.py b/backend/tests/unit/platform/http_client/test_notion_proxy.py index 6b20c981bc..38fecb8405 100644 --- a/backend/tests/unit/platform/http_client/test_notion_proxy.py +++ b/backend/tests/unit/platform/http_client/test_notion_proxy.py @@ -19,6 +19,7 @@ async def test_notion_binding_preserves_version_body_and_native_error(monkeypatc ) fetch = AsyncMock( return_value={ + "id": "ca_notion", "toolkit": {"slug": "notion"}, "status": "ACTIVE", "auth_config": {"id": "ac_notion"}, @@ -124,6 +125,7 @@ async def test_notion_rejects_missing_account_or_mismatched_binding( ) fetch = AsyncMock( return_value={ + "id": account, "toolkit": {"slug": toolkit}, "status": "ACTIVE", "auth_config": {"id": config_id}, diff --git a/backend/tests/unit/platform/sources/test_github_capture.py b/backend/tests/unit/platform/sources/test_github_capture.py index 3d7986349f..a08984a375 100644 --- a/backend/tests/unit/platform/sources/test_github_capture.py +++ b/backend/tests/unit/platform/sources/test_github_capture.py @@ -491,6 +491,7 @@ async def test_composio_capability_binds_only_selected_github_account_and_api_ho ) provider._get_with_auth = AsyncMock( return_value={ + "id": "ca_synthetic", "toolkit": {"slug": "github"}, "auth_config": {"id": "cfg"}, "status": "ACTIVE", diff --git a/backend/tests/unit/platform/sources/test_linear_capture.py b/backend/tests/unit/platform/sources/test_linear_capture.py index 76e85b40ad..48b42060a1 100644 --- a/backend/tests/unit/platform/sources/test_linear_capture.py +++ b/backend/tests/unit/platform/sources/test_linear_capture.py @@ -282,6 +282,7 @@ async def test_managed_linear_binding_only_allows_graphql_origin(monkeypatch): ) fetch = AsyncMock( return_value={ + "id": "ca_test", "toolkit": {"slug": "linear"}, "status": "ACTIVE", "auth_config": {"id": "ac_test"}, From cf14b4e9302e4b4622d27d210c852eb4ded8ce4d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:53:08 -0700 Subject: [PATCH 075/266] Preserve discovery policy when validating page pipeline topology --- .../tests/test_wispr_resume_probe.py | 35 ++++++++++++++++++- .../sync_pipeline/canonical_capture.py | 1 + 2 files changed, 35 insertions(+), 1 deletion(-) 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 index 448b631213..ba04eda148 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py @@ -19,6 +19,11 @@ 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"], @@ -43,7 +48,19 @@ async def counted(*args): 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) - await driver(service, sessions, fence, BoundedPageSource(source, counters, 600, probe)).run() + 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()) """ @@ -123,3 +140,19 @@ async def execute(path, timeout): 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 diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index 20e3f1581f..a87ddfd7ab 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -67,6 +67,7 @@ def __init__( fingerprint=config.fingerprint, record_types=record_types, container_parents=self._container_parents, + completion_policies=config.completion_policies, ) if config != declared: raise ValueError("Page source cycle must match its declared container topology") From 19776c5321eb363f75ffbb8522367cb002a21ef3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:55:10 -0700 Subject: [PATCH 076/266] Record bounded Wispr capture and incomplete live recovery proof --- ...r-recovery-budget-incomplete-20260930.json | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 backend/tests/live/evidence/wispr-recovery-budget-incomplete-20260930.json diff --git a/backend/tests/live/evidence/wispr-recovery-budget-incomplete-20260930.json b/backend/tests/live/evidence/wispr-recovery-budget-incomplete-20260930.json new file mode 100644 index 0000000000..c1fe90cd4f --- /dev/null +++ b/backend/tests/live/evidence/wispr-recovery-budget-incomplete-20260930.json @@ -0,0 +1,42 @@ +{ + "date": "2026-09-30", + "commit": "cf14b4e", + "provider": "wispr", + "method": "real provider reads through actual pipeline, intentional process interruption; disposable PostgreSQL and files", + "earlier_preflight_requests": 6, + "trial_request_limit": 14, + "aggregate_trial_requests": 14, + "aggregate_batch_requests": 20, + "runs": [ + { + "stage": "interrupt", + "provider_requests": 12, + "records_observed": 256, + "completed_meeting_bodies": 3, + "intentional_interruption": true, + "full_scope_completed": false + }, + { + "stage": "resume", + "provider_requests": 2, + "records_observed": 0, + "failed": true, + "error_type": "BudgetExceeded", + "safe_reason": "provider_requests", + "last_operation": "session", + "last_http_status": 201, + "checkpoint_unchanged": true, + "job_status": "running", + "full_scope_completed": false + } + ], + "wispr_recovery_verified": false, + "schema_removed": true, + "independent_live_schema_count": 0, + "blob_directory_removed": true, + "limitations": [ + "Resume budget expired during setup before a new body could be read.", + "No complete discovery, deletion detection or full-sync proof.", + "No hosted authentication or product owner attestation." + ] +} From 231890137c70e0218fd52a164847c1e53e45ae7e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 08:56:12 -0700 Subject: [PATCH 077/266] Record bounded Wispr interruption result and remaining proof budget --- backend/tests/live/README.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index c15684b671..03478fa7f2 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -323,3 +323,25 @@ rate failure or exhausted budget can leave the proof inconclusive. Without the o flag, Wispr retains its ordinary first/second-cycle trial; that is not a recovery proof. Offline PostgreSQL/subprocess tests verify the actual process exits and fresh retry, failed-body stop, and aggregate budget. No live recovery result is claimed here. + +Live checkpoint, 2026-09-30 (`cf14b4e`): the pipeline constructor had discarded +`completion_policies` while checking topology; the fix preserves the source policy +and still rejects incorrect topology. Full orchestrator/subprocess tests now cover +this integration, rather than only the scan driver. + +The subsequent14-request trial used12 requests for253 listing observations and +three committed meeting bodies, then intentionally exited75. The fresh process +used its remaining2 requests during account/session setup (last response HTTP201) +and hit the request guard before validation/body capture. Thus storage and intentional +interruption are live-verified; recovery remains **unverified**. Together with the +six prior preflight requests, this exhausted the20-request batch. Private schemas, +files and code archive were removed; independent schema count was zero. See +[evidence](evidence/wispr-recovery-budget-incomplete-20260930.json). + +A future independently approved trial with a fresh20-request budget is plausibly +sufficient, not guaranteed: measured first-stage12 + expected fresh setup3 + listing +refresh2 leaves3 GET/range calls for one additional meeting. A one-call body needs18 +requests total; a body requiring4 calls needs21 and must stop under the current cap. +Inventory and range sizes can change, and prior rate failures provide no reliable +reset interval. Keep the current20 maximum and180-second deadline; do not automatically +raise limits or rerun after a failure. No further live requests were made for this note. From 78e7f374ea5f3fc765b88ac540a7bf27c8fe73c6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:01:00 -0700 Subject: [PATCH 078/266] Record verified live Wispr fresh-process recovery --- .../evidence/wispr-recovery-20260930.json | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 backend/tests/live/evidence/wispr-recovery-20260930.json diff --git a/backend/tests/live/evidence/wispr-recovery-20260930.json b/backend/tests/live/evidence/wispr-recovery-20260930.json new file mode 100644 index 0000000000..3821498f04 --- /dev/null +++ b/backend/tests/live/evidence/wispr-recovery-20260930.json @@ -0,0 +1,42 @@ +{ + "date": "2026-09-30", + "commit": "2318901", + "provider": "wispr", + "method": "read-only real provider, actual orchestrator/pipeline and disposable PostgreSQL; intentional process loss then fresh same-job writer attempt", + "request_limit": 20, + "timeout_seconds": 180, + "runs": [ + { + "stage": "interrupt", + "failed": false, + "provider_requests": 10, + "records_observed": 256, + "completed_bodies": 3, + "intentional_interruption": true, + "full_scope_completed": false + }, + { + "stage": "resume", + "failed": false, + "provider_requests": 8, + "records_observed": 254, + "completed_bodies": 4, + "wispr_recovery_verified": true, + "full_scope_completed": false + } + ], + "aggregate_provider_requests": 18, + "aggregate_records_observed": 510, + "wispr_recovery_verified": true, + "schema_removed": true, + "independent_live_schema_count": 0, + "blob_directory_removed": true, + "proof": "Three completed body scopes remain at the same cycle/revisions; the fresh process rejects revisiting their native identities and commits one new body.", + "limitations": [ + "No complete-source traversal or subsequent-cycle proof.", + "No exhaustive discovery or deletion guarantee.", + "Failed individual meetings may repeat their text ranges.", + "Captured bodies omit provider raw editor data and remain partial.", + "No hosted auth or Almanac owner-binding attestation." + ] +} From 5b6c061cf037c61fc50a8527998294d467aad088 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:01:54 -0700 Subject: [PATCH 079/266] Capture native Notion roots and blocks with explicit discovery limits --- .../canonical/tests/test_notion_recovery.py | 251 +++ backend/airweave/platform/sources/notion.py | 1655 +++++------------ .../platform/sources/test_notion_capture.py | 333 ++++ 3 files changed, 1044 insertions(+), 1195 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py create mode 100644 backend/tests/unit/platform/sources/test_notion_capture.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py new file mode 100644 index 0000000000..d2500d1149 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py @@ -0,0 +1,251 @@ +"""Real Notion pipeline and PostgreSQL recovery with synthetic native HTTP only.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID, 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, orchestrator +from airweave.domains.entities.canonical.tests.test_slack_recovery import run, saved +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.platform.configs.config import NotionConfig +from airweave.platform.sources.notion import NotionSource + +PAGE, DATA, DATABASE, BLOCK1, BLOCK2 = (str(UUID(int=i)) for i in range(1, 6)) +STAMP = "2026-09-30T00:00:00Z" + + +def native(kind, identity, **fields): + return { + "object": kind, + "id": identity, + "created_time": STAMP, + "last_edited_time": STAMP, + "in_trash": False, + "parent": {"type": "workspace", "workspace": True}, + **fields, + } + + +def listing(rows=(), cursor=None): + return { + "object": "list", + "results": list(rows), + "next_cursor": cursor, + "has_more": cursor is not None, + } + + +class NativeFixture: + """Endpoint router only; production source, pipeline, driver and SQL stay real.""" + + def __init__(self): + self.interrupt = True + self.calls = [] + self.page = native("page", PAGE, properties={"unknown": {"kept": True}}) + self.data = native( + "data_source", DATA, parent={"type": "database_id", "database_id": DATABASE} + ) + self.database = native("database", DATABASE, data_sources=[{"id": DATA}]) + self.blocks = [ + native( + "block", + item, + type="paragraph", + has_children=False, + parent={"type": "page_id", "page_id": PAGE}, + paragraph={"rich_text": [{"plain_text": "retained"}]}, + ) + for item in (BLOCK1, BLOCK2) + ] + + async def request(self, url, *, headers, json=None, params=None): + assert headers["Notion-Version"] == "2026-03-11" + path = url.removeprefix("https://api.notion.com/v1/") + self.calls.append((path, json, params)) + if path == "search": + payload = listing([self.data] if json["filter"]["value"] == "data_source" else []) + elif path == f"data_sources/{DATA}": + payload = self.data + elif path == f"databases/{DATABASE}": + payload = self.database + elif path == f"data_sources/{DATA}/query": + payload = listing([self.page]) + elif path == f"pages/{PAGE}": + payload = self.page + elif path == f"blocks/{PAGE}/children": + if params.get("start_cursor"): + if self.interrupt: + raise ConnectionError("synthetic interruption after committed block page") + payload = listing([self.blocks[1]]) + else: + payload = listing([self.blocks[0]], "block-page-2") + else: + raise AssertionError(f"Unexpected synthetic request: {path}") + return httpx.Response(200, json=payload, request=httpx.Request("GET", url)) + + +async def runner(database, source, fixture, attempt=1): + service, fence = source + client = AsyncMock() + client.get.side_effect = fixture.request + client.post.side_effect = fixture.request + connector = await NotionSource.create( + auth=StaticTokenProvider("fixture"), + logger=MagicMock(), + http_client=client, + config=NotionConfig(), + ) + 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, connector + + +async def test_discovered_roots_reach_frontier_and_interrupted_membership_refreshes( + database, source +): + fixture = NativeFixture() + first, _ = await runner(database, source, fixture) + with pytest.raises(ConnectionError, match="synthetic interruption"): + await run(first) + rows, checkpoint, scans = await saved(database) + assert "canonical_checkpoint" not in checkpoint + assert all( + row.parent_record_type is None + for row in rows + if row.entity_definition_short_name != "block" + ) + scan = next( + row for row in scans if row.record_type == "block" and row.continuation.get("cursor") + ) + assert scan.continuation["cursor"] == "block-page-2" and scan.phase == "collecting" + before = len(fixture.calls) + fixture.interrupt = False + second, _ = await runner(database, source, fixture, attempt=2) + await run(second) + later = fixture.calls[before:] + block_calls = [item for item in later if item[0] == f"blocks/{PAGE}/children"] + # Block inventories are themselves membership authorities. A new attempt + # refreshes the whole exact inventory; durable originals survive the failure. + assert len(block_calls) == 2 + assert "start_cursor" not in block_calls[0][2] + assert block_calls[1][2]["start_cursor"] == "block-page-2" + rows, checkpoint, _ = await saved(database) + assert {row.native_id for row in rows} == {PAGE, DATA, DATABASE, BLOCK1, BLOCK2} + assert all(row.deleted_at is None for row in rows) + assert all(row.completeness == "partial" for row in rows) + assert checkpoint["canonical_cycle"]["phase"] == "complete" + assert ( + checkpoint["canonical_cycle"]["configuration"]["completion_policies"]["page"] + == "discovery_with_validation" + ) + # No query-scope ancestry was assigned to independently visible pages/data sources. + assert all( + row.parent_record_type is None + for row in rows + if row.entity_definition_short_name in {"page", "data_source", "database"} + ) + + +async def test_root_omission_exact_refresh_preserves_visible_page_and_current_native_payload( + database, source +): + fixture = NativeFixture() + fixture.interrupt = False + instance, connector = await runner(database, source, fixture) + older = {**fixture.page, "properties": {"title": "old"}} + service, fence = source + async with database() as db: + await service.capture( + db, CaptureBatch(fence=fence, records=(connector._record("page", older),)) + ) + await run(instance) + async with database() as db: + page = await db.scalar(select(Entity).where(Entity.native_id == PAGE)) + assert page.source_payload == fixture.page and page.deleted_at is None + assert page.parent_record_type is None + assert any(call[0] == f"pages/{PAGE}" for call in fixture.calls) + + +async def test_known_root_loss_redacts_descendants_even_when_search_returns_no_page( + database, source +): + fixture = NativeFixture() + service, fence = source + instance, connector = await runner(database, source, fixture) + root = connector._record("page", fixture.page) + child = connector._record("block", fixture.blocks[0], parent=root.identity) + async with database() as db: + captured = await service.capture(db, CaptureBatch(fence=fence, records=(root, child))) + + async def missing(url, **kwargs): + path = url.removeprefix("https://api.notion.com/v1/") + if path == "search": + return httpx.Response(200, json=listing(), request=httpx.Request("POST", url)) + assert path == f"pages/{PAGE}" + return httpx.Response( + 404, json={"code": "object_not_found"}, request=httpx.Request("GET", url) + ) + + connector.http_client.get.side_effect = missing + connector.http_client.post.side_effect = missing + await run(instance) + async with database() as db: + for change in captured.changes: + record = await service.store.read( + db, fence.organization_id, fence.sync_id, change.record.id + ) + assert record.content_access == "unavailable" and record.payload == {} + rows = (await db.scalars(select(Entity))).all() + assert all(row.removal_reason == "scope_removed" for row in rows) + + +async def test_prior_block_deleted_between_sweeps_reconciles_after_exact_confirmation( + database, source +): + fixture = NativeFixture() + fixture.interrupt = False + instance, connector = await runner(database, source, fixture) + service, fence = source + root = connector._record("page", fixture.page) + child = connector._record("block", fixture.blocks[0], parent=root.identity) + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(root, child))) + calls = [] + + async def request(url, **kwargs): + path = url.removeprefix("https://api.notion.com/v1/") + calls.append(path) + if path == f"blocks/{PAGE}/children": + return httpx.Response(200, json=listing(), request=httpx.Request("GET", url)) + if path == f"blocks/{BLOCK1}": + return httpx.Response( + 404, json={"code": "object_not_found"}, request=httpx.Request("GET", url) + ) + return await fixture.request(url, **kwargs) + + connector.http_client.get.side_effect = request + await run(instance) + rows, checkpoint, _ = await saved(database) + missing = next(row for row in rows if row.native_id == BLOCK1) + assert missing.deleted_at is not None + assert f"blocks/{BLOCK1}" in calls + assert checkpoint["canonical_cycle"]["phase"] == "complete" diff --git a/backend/airweave/platform/sources/notion.py b/backend/airweave/platform/sources/notion.py index 9643bff6d6..54f46411cf 100644 --- a/backend/airweave/platform/sources/notion.py +++ b/backend/airweave/platform/sources/notion.py @@ -1,36 +1,49 @@ -"""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 +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 SourceError +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 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.http_client.retry_helpers import ( retry_if_rate_limit_or_timeout, @@ -40,6 +53,85 @@ 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 = 1 +_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 _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 _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 +149,396 @@ 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") + # 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"), + } + + @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 - # ------------------------------------------------------------------ - - @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) -> 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) - - 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 - ) - - raise_for_status( - response, - source_short_name=self.short_name, - token_provider_kind=self.auth.provider_kind, + """Build a native managed/direct-auth source without acquiring additional grants.""" + instance = cls(auth=auth, logger=logger, http_client=http_client) + fingerprint = f"notion:{API_VERSION}:{FIELD_SET_VERSION}:partial-native-roots-blocks" + 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 response.json() + 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, 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 - ) - + 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 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 - ) - + 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() - - # ------------------------------------------------------------------ - # 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 + result = response.json() + if not isinstance(result, dict): + raise ValueError("Notion returned a non-object response") + return result - 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}" + async def validate(self) -> None: + """Verify identity independently of content discovery capabilities.""" + payload = await self._request("users/me") + _Id.model_validate(payload) + + 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, ) - 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}" + @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), ) - # ------------------------------------------------------------------ - # 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 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) + ) - async def _create_comprehensive_page_entity( + 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) + + 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) + scope: CompletedScope, + continuation: ScanContinuation, + *, + files: FileService, + parent: SourceRecord | None = None, + ) -> CapturePage: + """Return bounded observations; SQL commits them together with this continuation.""" + 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, + 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, + ) - 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", {}) + @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") + if len(progress.cursor_hashes) >= 512: + raise ValueError( + "Notion pagination exceeded the bounded cursor history; capture is incomplete" + ) + return page.next_cursor, (*progress.cursor_hashes, 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" ) - - 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"], + 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, ) - files_with_breadcrumbs = [] - for file_entity in content_result["files"]: - file_entity.breadcrumbs = breadcrumbs.copy() - files_with_breadcrumbs.append(file_entity) - - 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}" - ) - 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)}" + 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" ) - 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 - ) - - 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( - 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 = _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" ) - 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=[] + capped = progress.capped or bool( + page.request_status and page.request_status.type == "incomplete" ) - 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", [])) - 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, + next_progress = progress.model_copy( + update={ + "cursor": cursor, + "cursor_hashes": cursor_hashes, + "last_created": last, + "capped": capped, + } ) - - # ------------------------------------------------------------------ - # 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, - ) - - 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}" + 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", ) - 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 - # ------------------------------------------------------------------ + 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" + ) + 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, + ) - 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/tests/unit/platform/sources/test_notion_capture.py b/backend/tests/unit/platform/sources/test_notion_capture.py new file mode 100644 index 0000000000..9145974685 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_notion_capture.py @@ -0,0 +1,333 @@ +"""Native Notion acquisition fixtures; no provider calls or OAuth grants.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import httpx +import pytest +from pydantic import ValidationError + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.page_source import ScopeAccessLost +from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity +from airweave.domains.entities.canonical.scan_models import ScanContinuation +from airweave.domains.sources.exceptions import SourceEntityForbiddenError, SourceError +from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider +from airweave.platform.configs.config import NotionConfig +from airweave.platform.sources.notion import API_VERSION, NotionSource + +PAGE, DATA, DATABASE, BLOCK, CHILD = (str(UUID(int=i)) for i in range(1, 6)) +STAMP = "2026-09-30T00:00:00Z" + + +def native(kind="page", identity=PAGE, **fields): + return { + "object": kind, + "id": identity, + "created_time": STAMP, + "last_edited_time": STAMP, + "in_trash": False, + "parent": {"type": "workspace", "workspace": True}, + "unknown": {"retained": True}, + **fields, + } + + +def listing(rows=(), cursor=None, status=None): + result = { + "object": "list", + "results": list(rows), + "has_more": cursor is not None, + "next_cursor": cursor, + } + if status: + result["request_status"] = {"type": "incomplete", "incomplete_reason": status} + return result + + +def response(payload, status=200): + return httpx.Response( + status, json=payload, request=httpx.Request("GET", "https://api.notion.com/v1/test") + ) + + +async def source(*, gets=(), posts=()): + client = AsyncMock() + client.get.side_effect = [response(item) if isinstance(item, dict) else item for item in gets] + client.post.side_effect = [response(item) if isinstance(item, dict) else item for item in posts] + result = await NotionSource.create( + auth=ManagedAuthProvider( + api_key="fixture", + connected_account_id="fixture", + allowed_hosts=frozenset({"api.notion.com"}), + ), + logger=MagicMock(), + http_client=client, + config=NotionConfig(), + ) + return result, client + + +def parent(kind="page", identity=PAGE, **fields): + return SourceRecord.model_construct( + identity=RecordIdentity(record_type=kind, native_id=identity), + payload=native(kind, identity, **fields), + parent=None, + ) + + +@pytest.mark.asyncio +async def test_search_uses_exact_current_root_not_search_snapshot_and_keeps_unknown_fields(): + original = native(title=[{"plain_text": "Current"}]) + connector, client = await source(gets=[original], posts=[listing([native(title=[])])]) + page = await connector.capture_page( + CompletedScope(record_type="page"), ScanContinuation(), files=MagicMock() + ) + assert page.final and page.records[0].payload == original + assert page.records[0].parent is None and page.records[0].completeness == "partial" + assert page.records[0].source_created_at.isoformat() == "2026-09-30T00:00:00+00:00" + assert client.get.call_args.kwargs["headers"] == {"Notion-Version": API_VERSION} + assert client.post.call_args.kwargs["json"]["filter"]["value"] == "page" + assert connector.capture_cycle_configuration.policy("page") == "discovery_with_validation" + + +@pytest.mark.asyncio +async def test_search_rejects_wrong_exact_id_and_capped_discovery(): + connector, _ = await source(gets=[native(identity=CHILD)], posts=[listing([native()])]) + with pytest.raises(ValueError, match="wrong exact"): + await connector.capture_page( + CompletedScope(record_type="page"), ScanContinuation(), files=MagicMock() + ) + connector, _ = await source(posts=[listing(status="query_result_limit_reached")]) + with pytest.raises(SourceError, match="capped"): + await connector.capture_page( + CompletedScope(record_type="page"), ScanContinuation(), files=MagicMock() + ) + + +@pytest.mark.asyncio +async def test_omission_refresh_distinguishes_native_unavailability_from_capability_failure(): + connector, _ = await source(gets=[response({"code": "object_not_found"}, 404)]) + result = await connector.refresh_known(parent(), files=MagicMock()) + assert result.kind == "delete" and result.removal_reason == "scope_removed" + connector, _ = await source(gets=[response({"code": "restricted_resource"}, 403)]) + with pytest.raises(SourceEntityForbiddenError): + await connector.refresh_known(parent(), files=MagicMock()) + connector, _ = await source(gets=[native(in_trash=True)]) + result = await connector.refresh_known(parent(), files=MagicMock()) + assert result.removal_reason == "provider_deleted" and result.payload["in_trash"] is True + + +@pytest.mark.asyncio +async def test_query_rows_are_independently_retrieved_roots_and_cap_windows_resume(): + row = native(created_time="2026-09-30T01:00:00Z") + connector, client = await source( + gets=[row, row], posts=[listing([row], status="query_result_limit_reached"), listing([row])] + ) + owner = parent("data_source", DATA) + scope = connector.child_scope(owner, "page") + first = await connector.capture_page(scope, ScanContinuation(), files=MagicMock(), parent=owner) + assert not first.final and first.records == () + assert first.discovered_records[0].parent is None + second = await connector.capture_page( + scope, first.continuation, files=MagicMock(), parent=owner + ) + assert second.final + assert client.post.call_args.kwargs["json"]["filter"]["created_time"]["on_or_after"].startswith( + "2026-09-30T01:" + ) + + +@pytest.mark.asyncio +async def test_query_limit_with_same_timestamp_fails_without_claiming_completion(): + connector, _ = await source( + gets=[native()], posts=[listing([native()], status="query_result_limit_reached")] + ) + owner = parent("data_source", DATA) + with pytest.raises(SourceError, match="cannot be split"): + await connector.capture_page( + connector.child_scope(owner, "page"), + ScanContinuation(value={"window_start": STAMP}), + files=MagicMock(), + parent=owner, + ) + + +@pytest.mark.asyncio +async def test_database_query_scope_does_not_parent_discovered_data_source(): + database = native("database", DATABASE, data_sources=[{"id": DATA, "name": "Rows"}]) + data = native("data_source", DATA, parent={"type": "database_id", "database_id": DATABASE}) + connector, _ = await source(gets=[database, data]) + owner = parent("database", DATABASE) + result = await connector.capture_page( + connector.child_scope(owner, "data_source"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + assert result.final and not result.records + assert result.discovered_records[0].identity.record_type == "data_source" + assert result.discovered_records[0].parent is None + + +@pytest.mark.asyncio +async def test_direct_shared_data_source_does_not_depend_on_reading_database_parent(): + connector, _ = await source(gets=[response({"code": "object_not_found"}, 404)]) + owner = parent("data_source", DATA, parent={"type": "database_id", "database_id": DATABASE}) + result = await connector.capture_page( + connector.child_scope(owner, "database"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + assert result.final and result.discovered_records == () + + +@pytest.mark.asyncio +async def test_child_page_reference_retains_block_and_independent_page_without_double_crawl(): + block = native( + "block", + CHILD, + type="child_page", + has_children=True, + parent={"type": "page_id", "page_id": PAGE}, + ) + connector, _ = await source(gets=[listing([block]), native(identity=CHILD)]) + owner = parent() + result = await connector.capture_page( + connector.child_scope(owner, "block"), ScanContinuation(), files=MagicMock(), parent=owner + ) + assert result.records[0].parent == owner.identity + assert result.records[0].allow_reparent + assert result.discovered_records[0].parent is None + reference = parent( + "block", + CHILD, + **{key: value for key, value in block.items() if key not in {"object", "id"}}, + ) + empty = await connector.capture_page( + connector.child_scope(reference, "block"), + ScanContinuation(), + files=MagicMock(), + parent=reference, + ) + assert empty.final and not empty.records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["wrong_parent", "oversized", "missing_cursor", "repeated_cursor"]) +async def test_invalid_block_pages_fail_before_capture(case): + block = native( + "block", + BLOCK, + type="paragraph", + has_children=False, + parent={"type": "page_id", "page_id": CHILD if case == "wrong_parent" else PAGE}, + ) + data = listing([block] * (26 if case == "oversized" else 1)) + progress = ScanContinuation() + if case in {"missing_cursor", "repeated_cursor"}: + data["has_more"] = True + if case == "repeated_cursor": + data["next_cursor"] = "same" + progress = ScanContinuation(value={"cursor": "same"}) + connector, _ = await source(gets=[data]) + owner = parent() + with pytest.raises((ValueError, ValidationError)): + await connector.capture_page( + connector.child_scope(owner, "block"), progress, files=MagicMock(), parent=owner + ) + + +@pytest.mark.asyncio +async def test_exact_owner_unavailability_withdraws_scope_not_generic_errors(): + connector, _ = await source(gets=[response({"code": "object_not_found"}, 404)]) + owner = parent() + with pytest.raises(ScopeAccessLost) as caught: + await connector.capture_page( + connector.child_scope(owner, "block"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + assert caught.value.removal_reason == "scope_removed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("state", ["deleted", "moved", "trashed", "same_owner"]) +async def test_omitted_block_requires_exact_loss_or_move_evidence(state): + owner = RecordIdentity(record_type="page", native_id=PAGE) + block = native( + "block", + BLOCK, + type="paragraph", + has_children=False, + in_trash=state == "trashed", + parent={"type": "page_id", "page_id": CHILD if state == "moved" else PAGE}, + ) + reply = response({"code": "object_not_found"}, 404) if state == "deleted" else block + connector, _ = await source(gets=[reply]) + known = parent("block", BLOCK).model_copy(update={"parent": owner}) + if state == "same_owner": + with pytest.raises(ValueError, match="still readable"): + await connector.confirm_absent(known) + else: + await connector.confirm_absent(known) + + +@pytest.mark.asyncio +async def test_database_inventory_mutation_uses_existing_typed_restart(): + from airweave.domains.entities.canonical.page_source import InvalidScanContinuation + + database = native("database", DATABASE, data_sources=[{"id": DATA}]) + connector, _ = await source(gets=[database]) + owner = parent("database", DATABASE) + with pytest.raises(InvalidScanContinuation): + await connector.capture_page( + connector.child_scope(owner, "data_source"), + ScanContinuation(value={"cursor": "previous-list-digest", "offset": 25}), + files=MagicMock(), + parent=owner, + ) + + +@pytest.mark.asyncio +async def test_trashed_database_stops_enumeration_without_withdrawing_independent_children(): + connector, client = await source( + gets=[native("database", DATABASE, in_trash=True, data_sources=[{"id": DATA}])] + ) + owner = parent("database", DATABASE) + with pytest.raises(ScopeAccessLost) as caught: + await connector.capture_page( + connector.child_scope(owner, "data_source"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + assert caught.value.removal_reason == "scope_removed" + assert client.get.call_count == 1 + + +@pytest.mark.asyncio +async def test_unknown_query_incompleteness_never_uses_limit_recovery(): + connector, _ = await source(posts=[listing(status="unrecognized_future_reason")]) + owner = parent("data_source", DATA) + with pytest.raises(SourceError, match="unsupported reason"): + await connector.capture_page( + connector.child_scope(owner, "page"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + + +@pytest.mark.asyncio +async def test_nonadjacent_cursor_cycle_is_rejected_from_durable_continuation(): + connector, _ = await source( + posts=[listing(cursor="A"), listing(cursor="B"), listing(cursor="A")] + ) + scope = CompletedScope(record_type="page") + first = await connector.capture_page(scope, ScanContinuation(), files=MagicMock()) + second = await connector.capture_page(scope, first.continuation, files=MagicMock()) + with pytest.raises(ValueError, match="cursor cycle"): + await connector.capture_page(scope, second.continuation, files=MagicMock()) From 3790883b63c84f64f1e18d60b1caffaf006af5f7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:02:06 -0700 Subject: [PATCH 080/266] Document verified live Wispr body recovery without full-sync claims --- backend/tests/live/README.md | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 03478fa7f2..c4d92246e1 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -322,7 +322,7 @@ A source with fewer than five discoverable meetings, many body ranges, changed i rate failure or exhausted budget can leave the proof inconclusive. Without the opt-in flag, Wispr retains its ordinary first/second-cycle trial; that is not a recovery proof. Offline PostgreSQL/subprocess tests verify the actual process exits and fresh retry, -failed-body stop, and aggregate budget. No live recovery result is claimed here. +failed-body stop, and aggregate budget. The subsequent live result is recorded below. Live checkpoint, 2026-09-30 (`cf14b4e`): the pipeline constructor had discarded `completion_policies` while checking topology; the fix preserves the source policy @@ -345,3 +345,22 @@ requests total; a body requiring4 calls needs21 and must stop under the current Inventory and range sizes can change, and prior rate failures provide no reliable reset interval. Keep the current20 maximum and180-second deadline; do not automatically raise limits or rerun after a failure. No further live requests were made for this note. + + +**Latest live result, 2026-09-30: recovery verified.** A separate fresh20-request, +180-second trial59141 ran immutable `2318901`. First stage:10 HTTP requests, +256 observations, three completed meeting bodies, intentional exit75. Fresh +same-job process:8 HTTP requests,254 observations, one additional body committed, +intentional exit76 with `wispr_recovery_verified=true`. Aggregate:18 requests and +510 observations. These observations include repeated listing inventory, not510 +unique meetings. The prior three body scopes retained the same cycle/revisions; +the probe rejected any attempt to fetch their identities again. + +This proves per-meeting recovery through the actual source/orchestrator/store, +not a full-source traversal, new-cycle refresh, exhaustive discovery, deletion +handling or Almanac owner attestation. A failed meeting may repeat its own ranges; +all bodies remain partial because provider raw editor data is unavailable. +Parent independently verified zero live schemas and private-live directories and +removed the immutable code archive. The earlier20-request failed batch remains +historical evidence; this successful trial had its own20-request cap. Safe results: +[evidence](evidence/wispr-recovery-20260930.json), committed `78e7f37`. From 6fb78b35efc4f747f37d69342ef3bd56bed22558 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:04:15 -0700 Subject: [PATCH 081/266] Keep a rolling Notion cursor history without limiting pagination --- backend/airweave/platform/sources/notion.py | 8 ++-- .../platform/sources/test_notion_capture.py | 43 +++++++++++++++++++ 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/backend/airweave/platform/sources/notion.py b/backend/airweave/platform/sources/notion.py index 54f46411cf..1d447c9890 100644 --- a/backend/airweave/platform/sources/notion.py +++ b/backend/airweave/platform/sources/notion.py @@ -360,11 +360,9 @@ def _following(page: _List, progress: _Progress) -> tuple[str | None, tuple[str, digest = hashlib.sha256(page.next_cursor.encode()).hexdigest() if digest in progress.cursor_hashes: raise ValueError("Notion pagination entered a cursor cycle") - if len(progress.cursor_hashes) >= 512: - raise ValueError( - "Notion pagination exceeded the bounded cursor history; capture is incomplete" - ) - return page.next_cursor, (*progress.cursor_hashes, digest) + # 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 diff --git a/backend/tests/unit/platform/sources/test_notion_capture.py b/backend/tests/unit/platform/sources/test_notion_capture.py index 9145974685..3d674ee578 100644 --- a/backend/tests/unit/platform/sources/test_notion_capture.py +++ b/backend/tests/unit/platform/sources/test_notion_capture.py @@ -1,5 +1,6 @@ """Native Notion acquisition fixtures; no provider calls or OAuth grants.""" +import hashlib from unittest.mock import AsyncMock, MagicMock from uuid import UUID @@ -331,3 +332,45 @@ async def test_nonadjacent_cursor_cycle_is_rejected_from_durable_continuation(): second = await connector.capture_page(scope, first.continuation, files=MagicMock()) with pytest.raises(ValueError, match="cursor cycle"): await connector.capture_page(scope, second.continuation, files=MagicMock()) + + +@pytest.mark.asyncio +async def test_large_search_advances_with_bounded_recent_cursor_history(): + connector, _ = await source(posts=[listing(cursor=f"page-{i}") for i in range(520)]) + scope = CompletedScope(record_type="page") + continuation = ScanContinuation() + for _ in range(520): + page = await connector.capture_page(scope, continuation, files=MagicMock()) + assert not page.final + # Exercise the durable JSON boundary on every page, not only in-memory state. + continuation = ScanContinuation.model_validate_json(page.continuation.model_dump_json()) + assert len(continuation.value["cursor_hashes"]) <= 512 + hashes = continuation.value["cursor_hashes"] + assert continuation.value["cursor"] == "page-519" + assert hashes[0] == hashlib.sha256(b"page-8").hexdigest() + assert hashes[-1] == hashlib.sha256(b"page-519").hexdigest() + + +@pytest.mark.asyncio +async def test_query_window_resets_recent_cursor_history(): + row = native(created_time="2026-09-30T01:00:00Z") + connector, _ = await source( + gets=[row], + posts=[ + listing(cursor="reusable"), + listing([row], status="query_result_limit_reached"), + listing(cursor="reusable"), + ], + ) + owner = parent("data_source", DATA) + scope = connector.child_scope(owner, "page") + first = await connector.capture_page(scope, ScanContinuation(), files=MagicMock(), parent=owner) + window = await connector.capture_page( + scope, first.continuation, files=MagicMock(), parent=owner + ) + assert not window.final and window.continuation.value["cursor_hashes"] == [] + assert window.continuation.value["cursor"] is None + following = await connector.capture_page( + scope, window.continuation, files=MagicMock(), parent=owner + ) + assert following.continuation.value["cursor"] == "reusable" From b37c2455d492aafe0780fd06e1e17d033d32a097 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:13:29 -0700 Subject: [PATCH 082/266] Project retained Notion metadata and block text for search --- .../entities/canonical/notion_projection.py | 258 ++++++++++++++ .../entities/canonical/projection_mappers.py | 4 + .../canonical/tests/test_notion_projection.py | 314 ++++++++++++++++++ .../airweave/platform/entities/__init__.py | 2 + backend/airweave/platform/entities/notion.py | 24 ++ 5 files changed, 602 insertions(+) create mode 100644 backend/airweave/domains/entities/canonical/notion_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_notion_projection.py 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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index d5510b17b8..e371b33dd2 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -235,6 +235,10 @@ async def map_record( from airweave.domains.entities.canonical.linear_projection import map_linear entities = await map_linear(record, storage, directory) + elif source_name == "notion": + from airweave.domains.entities.canonical.notion_projection import map_notion + + entities = map_notion(record) elif source_name == "google_drive": entities = await _drive(record, storage, directory) elif source_name == "google_calendar": diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py new file mode 100644 index 0000000000..c871845a10 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py @@ -0,0 +1,314 @@ +"""Retained native Notion text, explicit limits, and existing publication authority.""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID, uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import ( + ProjectionMappingError, + map_record, +) +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.projector import CanonicalProjector +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared +from airweave.models import Entity, Sync + +PAGE, BLOCK = str(UUID(int=1)), str(UUID(int=2)) +STAMP = "2026-09-30T00:00:00Z" + + +def rich(text, kind="text"): + return {"type": kind, "plain_text": text} + + +def original(kind="page", **fields): + native = BLOCK if kind == "block" else PAGE + payload = { + "object": kind, + "id": native, + "created_time": STAMP, + "last_edited_time": STAMP, + "in_trash": False, + **fields, + } + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity(record_type=kind, native_id=native), + parent=RecordIdentity(record_type="page", native_id=PAGE) if kind == "block" else None, + revision=1, + payload=payload, + payload_schema_version=1, + capture_hash="hash", + 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 block(kind, content, **fields): + return original( + "block", + type=kind, + has_children=False, + parent={"type": "page_id", "page_id": PAGE}, + **{kind: content}, + **fields, + ) + + +async def mapped(record): + storage = AsyncMock() + async with map_record(record, "notion", storage) as entities: + assert len(entities) == 1 + entity = entities[0] + assert not storage.mock_calls + assert record.completeness == "partial" + assert entity.native_id == record.identity.native_id + return entity + + +@pytest.mark.asyncio +async def test_root_metadata_is_not_an_assembled_page_body(): + record = original( + properties={ + "Name": {"type": "title", "title": [rich("Launch "), rich("@Mira", "mention")]} + }, + url=f"https://app.notion.com/p/{PAGE}", + unknown={"private": "not prose"}, + ) + entity = await mapped(record) + assert entity.title == "Launch @Mira" and entity.text == "" + assert entity.web_url == record.payload["url"] + assert entity.created_at.tzinfo is not None + assert "Metadata only" in entity.content_coverage + for kind in ("database", "data_source"): + fields = {"title": [rich("Plans")], "description": [rich("Roadmap")]} + if kind == "data_source": + fields["properties"] = { + "Name": {"type": "title", "title": {}}, + "Stage": {"type": "status", "status": {}}, + } + entity = await mapped(original(kind, **fields)) + assert entity.title == "Plans" and "Roadmap" in entity.text + assert ("Stage (status)" in entity.text) == (kind == "data_source") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind,content,expected", + [ + ( + "paragraph", + {"rich_text": [rich("Hi "), rich("@Mira", "mention"), rich(" x²", "equation")]}, + "Hi @Mira x²", + ), + ("heading_2", {"rich_text": [rich("Milestones")]}, "Milestones"), + ("to_do", {"checked": True, "rich_text": [rich("Ship")]}, "[x] Ship"), + ( + "code", + { + "language": "python", + "rich_text": [rich("print('hello')\n")], + "caption": [rich("Example")], + }, + "python\nprint('hello')\n\nExample", + ), + ("table_row", {"cells": [[rich("Owner")], [rich("Mira")]]}, "Owner | Mira"), + ("equation", {"expression": "E=mc^2"}, "E=mc^2"), + ], +) +async def test_supported_native_block_text(kind, content, expected): + entity = await mapped(block(kind, content)) + assert entity.text == expected + assert entity.web_url == "" + assert "This block only" in entity.content_coverage + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind,content", + [ + ("table", {"table_width": 2, "has_column_header": True, "has_row_header": False}), + ("column", {}), + ("divider", {}), + ("synced_block", {"synced_from": None}), + ], +) +async def test_known_structure_has_explicit_empty_body(kind, content): + entity = await mapped(block(kind, content)) + assert entity.title and entity.text == "" + assert any(word in entity.content_coverage for word in ("only", "separate")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["child_page", "child_database"]) +async def test_native_reference_label_does_not_duplicate_target_body(kind): + entity = await mapped(block(kind, {"title": "Budget"})) + assert entity.title == "Budget" and entity.text == "" + assert entity.web_url == "" and "Reference label only" in entity.content_coverage + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "record", + [ + block("future_type", {"secret": "never serialize me"}), + block("synced_block", {"synced_from": {"type": "block_id", "block_id": PAGE}}), + block("paragraph", {"rich_text": [{"type": "future", "plain_text": "unknown"}]}), + block( + "paragraph", + {"rich_text": [{"type": "text", "text": {"content": "not retained plain text"}}]}, + ), + ], +) +async def test_unsupported_or_malformed_content_never_becomes_empty_success(record): + with pytest.raises(ValueError): + await mapped(record) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "javascript:alert(1)", + "https://evil.example/page", + "https://app.notion.com.evil.example/page", + "https://user@app.notion.com/page", + "http://notion.so/page", + ], +) +async def test_untrusted_navigation_remains_pending(url): + with pytest.raises(ValueError): + await mapped(original(properties={"title": {"type": "title", "title": []}}, url=url)) + + +@pytest.mark.asyncio +async def test_identity_and_authority_mismatches_are_rejected(): + record = block("paragraph", {"rich_text": [rich("private")]}) + cases = [ + record.model_copy(update={"parent": None}), + record.model_copy( + update={"identity": record.identity.model_copy(update={"native_id": PAGE})} + ), + record.model_copy(update={"content_access": "parent_unavailable"}), + record.model_copy(update={"deleted_at": datetime.now(timezone.utc)}), + record.model_copy(update={"payload": {**record.payload, "in_trash": True}}), + ] + for case in cases: + with pytest.raises(ProjectionMappingError): + await mapped(case) + + +def test_original_search_view_is_in_existing_entity_registry(): + from airweave.domains.entities.registry import EntityDefinitionRegistry + from airweave.platform.entities.notion import NotionOriginalEntity + + registry = EntityDefinitionRegistry() + registry.build() + assert registry.get_short_name_by_class(NotionOriginalEntity) == "notion_original_entity" + assert (NotionOriginalEntity.model_fields["content_coverage"].json_schema_extra or {}).get( + "embeddable" + ) is not True + + +@pytest.mark.integration +async def test_notion_publication_uses_current_revision_and_ancestor_authority(database, source): + service, fence = source + root = original(properties={"title": {"type": "title", "title": [rich("Page")]}}) + child = block("paragraph", {"rich_text": [rich("First version")]}) + root_observation = observation( + identity=root.identity, payload=root.payload, completeness="partial" + ) + child_observation = observation( + identity=child.identity, parent=root.identity, payload=child.payload, completeness="partial" + ) + await capture(database, service, fence, root_observation, child_observation) + store = CanonicalProjectionStore() + async with database() as db: + pending = await store.pending(db, fence.organization_id, fence.sync_id) + work = next(item for item in pending if item.record.identity.record_type == "block") + assert (await mapped(work.record)).text == "First version" + generation = uuid4() + async with database() as db: + assert await publish_prepared(store, db, work, generation, 1) + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=0, + ) + updated = child_observation.model_copy( + update={"payload": {**child.payload, "paragraph": {"rich_text": [rich("Changed")]}}} + ) + await capture(database, service, fence, updated) + 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(), 1) + pending = await store.pending(db, fence.organization_id, fence.sync_id) + fresh = next(item for item in pending if item.record.identity.record_type == "block") + await capture( + database, + service, + fence, + root_observation.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}), + ) + async with database() as db: + assert not await publish_prepared(store, db, fresh, uuid4(), 1) + + +@pytest.mark.integration +async def test_unknown_notion_block_stays_pending_in_actual_projector(database, source): + service, fence = source + root = original(properties={"title": {"type": "title", "title": []}}) + child = block("future_type", {"body": "not raw JSON prose"}) + await capture( + database, + service, + fence, + observation(identity=root.identity, payload=root.payload, completeness="partial"), + observation( + identity=child.identity, + parent=root.identity, + payload=child.payload, + completeness="partial", + ), + ) + store = CanonicalProjectionStore() + async with database() as db: + pending = await store.pending(db, fence.organization_id, fence.sync_id) + root_work = next(item for item in pending if item.record.identity.record_type == "page") + assert await publish_prepared(store, db, root_work, uuid4(), 1) + processor, destination = MagicMock(), MagicMock() + projector = CanonicalProjector(store, database, processor, AsyncMock()) + result = await projector.batch( + fence.organization_id, fence.sync_id, "notion", destination, MagicMock() + ) + assert result.failed == 1 and result.published == 0 + processor.process.assert_not_called() + destination.feed_prepared.assert_not_called() + async with database() as db: + pending = await store.pending(db, fence.organization_id, fence.sync_id) + assert len(pending) == 1 + row = await db.get(Entity, pending[0].record.id) + assert row.projection_error == "ProjectionMappingError" and row.indexed_revision is None diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index fa66e99fca..811f37f5fd 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -198,6 +198,7 @@ from .notion import ( NotionDatabaseEntity, NotionFileEntity, + NotionOriginalEntity, NotionPageEntity, NotionPropertyEntity, ) @@ -570,6 +571,7 @@ "notion": [ NotionDatabaseEntity, NotionFileEntity, + NotionOriginalEntity, NotionPageEntity, NotionPropertyEntity, ], diff --git a/backend/airweave/platform/entities/notion.py b/backend/airweave/platform/entities/notion.py index 752b116269..06d911481d 100644 --- a/backend/airweave/platform/entities/notion.py +++ b/backend/airweave/platform/entities/notion.py @@ -556,3 +556,27 @@ 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 UUID.", 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 "" From 2bb84cfe51ce2b14ede2e5ce5e7cbab416419b1f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:24:16 -0700 Subject: [PATCH 083/266] Retain complete native Notion property responses atomically per property --- .../canonical/tests/test_notion_recovery.py | 152 ++++++++++- backend/airweave/platform/sources/notion.py | 238 +++++++++++++++++- .../platform/sources/test_notion_capture.py | 176 +++++++++++++ 3 files changed, 558 insertions(+), 8 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py index d2500d1149..d463093478 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py @@ -48,7 +48,7 @@ class NativeFixture: def __init__(self): self.interrupt = True self.calls = [] - self.page = native("page", PAGE, properties={"unknown": {"kept": True}}) + self.page = native("page", PAGE, properties={}) self.data = native( "data_source", DATA, parent={"type": "database_id", "database_id": DATABASE} ) @@ -91,7 +91,7 @@ async def request(self, url, *, headers, json=None, params=None): return httpx.Response(200, json=payload, request=httpx.Request("GET", url)) -async def runner(database, source, fixture, attempt=1): +async def runner(database, source, fixture, attempt=1, files=None): service, fence = source client = AsyncMock() client.get.side_effect = fixture.request @@ -111,7 +111,7 @@ async def runner(database, source, fixture, attempt=1): CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), connector.canonical_container_parents, page_source=connector, - files=MagicMock(), + files=files or MagicMock(), ) runtime.source, runtime.canonical_capture = connector, pipeline instance = orchestrator(ctx, pipeline, runtime, None, bus) @@ -249,3 +249,149 @@ async def request(url, **kwargs): assert missing.deleted_at is not None assert f"blocks/{BLOCK1}" in calls assert checkpoint["canonical_cycle"]["phase"] == "complete" + + +class PropertyFixture(NativeFixture): + def __init__(self): + super().__init__() + self.page["properties"] = { + name: {"id": name, "type": "title", "title": []} for name in ("a", "b") + } + self.property_calls = [] + self.interrupt_property = True + + async def request(self, url, *, headers, json=None, params=None): + path = url.removeprefix("https://api.notion.com/v1/") + if path == f"blocks/{PAGE}/children": + return httpx.Response(200, json=listing(), request=httpx.Request("GET", url)) + if "/properties/" in path: + prop = path.rsplit("/", 1)[-1] + cursor = (params or {}).get("start_cursor") + self.property_calls.append((prop, cursor)) + if prop == "b" and cursor and self.interrupt_property: + raise ConnectionError("synthetic interruption within one property") + more = prop == "b" and cursor is None + payload = { + **listing( + [ + { + "object": "property_item", + "id": prop, + "type": "title", + "title": {"plain_text": "retained"}, + } + ], + "following" if more else None, + ), + "type": "property_item", + "property_item": {"id": prop, "type": "title", "next_url": None}, + } + return httpx.Response(200, json=payload, request=httpx.Request("GET", url)) + return await super().request(url, headers=headers, json=json, params=params) + + +async def test_property_retry_keeps_committed_sibling_and_restarts_only_interrupted_property( + database, source, tmp_path, monkeypatch +): + import json + + from airweave.adapters.storage.filesystem import FilesystemBackend + from airweave.domains.storage.file_service import FileService + + 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 + ) + fixture = PropertyFixture() + first, _ = await runner(database, source, fixture, files=files) + with pytest.raises(ConnectionError, match="within one property"): + await run(first) + rows, checkpoint, scans = await saved(database) + properties = [row for row in rows if row.entity_definition_short_name == "page_property"] + assert [row.native_id for row in properties] == ["a"] + assert "canonical_checkpoint" not in checkpoint + scan = next(row for row in scans if row.record_type == "page_property") + assert scan.continuation["index"] == 1 + assert len(await files.storage.list_files()) == 1 + before = len(fixture.property_calls) + fixture.interrupt_property = False + second, _ = await runner(database, source, fixture, attempt=2, files=files) + await run(second) + assert fixture.property_calls[before:] == [("b", None), ("b", "following")] + rows, checkpoint, _ = await saved(database) + properties = [row for row in rows if row.entity_definition_short_name == "page_property"] + assert {row.native_id for row in properties} == {"a", "b"} + assert all(row.parent_record_type == "page" and row.container_id == PAGE for row in properties) + assert next(row for row in properties if row.native_id == "a").record_revision == 1 + second_property = next(row for row in properties if row.native_id == "b") + archive = json.loads(await files.storage.read_file(second_property.blob_references[0]["key"])) + assert len(archive["responses"]) == 2 + assert checkpoint["canonical_cycle"]["phase"] == "complete" + + # A genuinely later job gets a new cycle and reconciles a removed property. + from airweave.models.sync_job import SyncJob + + previous_cycle = checkpoint["canonical_cycle"]["version"]["cycle_id"] + next_job, next_attempt = uuid4(), uuid4() + async with database() as db: + old_job = await db.get(SyncJob, fence.job_id) + old_job.status = "completed" + db.add( + SyncJob( + id=next_job, + organization_id=fence.organization_id, + sync_id=fence.sync_id, + status="running", + ) + ) + await db.commit() + fixture.page["properties"].pop("a") + fixture.page["last_edited_time"] = "2026-09-30T02:00:00Z" + next_source = ( + service, + fence.model_copy(update={"job_id": next_job, "attempt_id": next_attempt}), + ) + third, _ = await runner(database, next_source, fixture, files=files) + await run(third) + rows, checkpoint, _ = await saved(database) + removed = next( + row + for row in rows + if row.entity_definition_short_name == "page_property" and row.native_id == "a" + ) + assert removed.deleted_at is not None and removed.removal_reason == "absent" + assert checkpoint["canonical_cycle"]["version"]["cycle_id"] != previous_cycle + assert checkpoint["canonical_cycle"]["phase"] == "complete" + + +async def test_removed_native_property_reconciles_only_after_current_inventory(database, source): + from datetime import datetime, timezone + + from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity + + fixture = NativeFixture() + fixture.interrupt = False + instance, connector = await runner(database, source, fixture) + service, fence = source + root = connector._record("page", fixture.page) + old_property = CaptureRecord( + identity=RecordIdentity( + record_type="page_property", native_id="removed", container_id=PAGE + ), + parent=root.identity, + payload={"native": "previously captured value"}, + observed_at=datetime.now(timezone.utc), + completeness="partial", + ) + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(root, old_property))) + await run(instance) + rows, checkpoint, _ = await saved(database) + removed = next(row for row in rows if row.native_id == "removed") + assert removed.deleted_at is not None and removed.removal_reason == "absent" + assert next(row for row in rows if row.native_id == PAGE).deleted_at is None + assert checkpoint["canonical_cycle"]["phase"] == "complete" diff --git a/backend/airweave/platform/sources/notion.py b/backend/airweave/platform/sources/notion.py index 1d447c9890..f9ec3ab78f 100644 --- a/backend/airweave/platform/sources/notion.py +++ b/backend/airweave/platform/sources/notion.py @@ -8,6 +8,7 @@ from __future__ import annotations import hashlib +import json from collections.abc import AsyncGenerator, Sequence from datetime import datetime, timezone from functools import partial @@ -40,6 +41,7 @@ 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 @@ -54,7 +56,7 @@ from airweave.schemas.source_connection import AuthenticationMethod, OAuthType API_VERSION = "2026-03-11" -FIELD_SET_VERSION = 1 +FIELD_SET_VERSION = 2 _ROOTS = ("page", "database", "data_source") _ENDPOINTS = { "page": "pages", @@ -116,6 +118,58 @@ class _List(_Native): 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 @@ -151,7 +205,7 @@ class _Progress(BaseModel): class NotionSource(BaseSource): """One source registration, native acquisition, and existing SQL-owned page progress.""" - canonical_record_types = (*_ROOTS, "block") + 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 = { @@ -159,6 +213,7 @@ class NotionSource(BaseSource): "database": (None, "data_source"), "data_source": (None, "database"), "block": ("page", "block"), + "page_property": "page", } @property @@ -177,7 +232,9 @@ async def create( ) -> NotionSource: """Build a native managed/direct-auth source without acquiring additional grants.""" instance = cls(auth=auth, logger=logger, http_client=http_client) - fingerprint = f"notion:{API_VERSION}:{FIELD_SET_VERSION}:partial-native-roots-blocks" + fingerprint = ( + f"notion:{API_VERSION}:{FIELD_SET_VERSION}:partial-native-roots-blocks-properties" + ) instance._cycle_configuration = CycleConfiguration.from_source( fingerprint=hashlib.sha256(fingerprint.encode()).hexdigest(), record_types=cls.canonical_record_types, @@ -306,7 +363,11 @@ def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: parent.identity.record_type ): raise ValueError("Unsupported Notion enumeration scope") - return CompletedScope(record_type=record_type, parent=parent.identity) + return CompletedScope( + record_type=record_type, + parent=parent.identity, + container_id=parent.identity.native_id if record_type == "page_property" else None, + ) async def capture_page( self, @@ -317,6 +378,16 @@ async def capture_page( 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 + ) progress = _Progress.model_validate(continuation.value) if scope.container_id is not None: raise ValueError("Notion native identities do not have a synthetic container") @@ -341,7 +412,7 @@ async def capture_page( @staticmethod def _page( records: Sequence[CaptureRecord], - progress: _Progress, + progress: _Progress | _PropertyProgress, final: bool, discovered: Sequence[CaptureRecord] = (), ) -> CapturePage: @@ -530,6 +601,163 @@ async def _blocks(self, parent: SourceRecord, progress: _Progress) -> CapturePag roots, ) + async def _property_inventory( + self, + 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" + ) + page = _PageProperties.model_validate(payload) + rows = sorted( + page.properties.items(), key=lambda item: _Property.model_validate(item[1]).id + ) + 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: + 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() + ) + 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", + ) + 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", + ) + 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( self, *, diff --git a/backend/tests/unit/platform/sources/test_notion_capture.py b/backend/tests/unit/platform/sources/test_notion_capture.py index 3d674ee578..84600cd936 100644 --- a/backend/tests/unit/platform/sources/test_notion_capture.py +++ b/backend/tests/unit/platform/sources/test_notion_capture.py @@ -374,3 +374,179 @@ async def test_query_window_resets_recent_cursor_history(): scope, window.continuation, files=MagicMock(), parent=owner ) assert following.continuation.value["cursor"] == "reusable" + + +@pytest.fixture +def property_files(tmp_path, monkeypatch): + from uuid import uuid4 + + from airweave.adapters.storage.filesystem import FilesystemBackend + from airweave.domains.storage.file_service import FileService + + monkeypatch.setattr( + "airweave.domains.storage.file_service.paths.temp_sync_dir", + lambda _: str(tmp_path / "temp"), + ) + return FileService(uuid4(), FilesystemBackend(tmp_path / "blobs"), sync_id=uuid4()) + + +def property_list(identity, kind, values, *, cursor=None, calculation=None): + metadata = {"id": identity, "type": kind, "next_url": None} + if calculation is not None: + metadata[kind] = calculation + return {**listing(values, cursor), "type": "property_item", "property_item": metadata} + + +@pytest.mark.asyncio +async def test_property_archive_preserves_all_native_pages_and_final_rollup(property_files): + import json + + prop = {"id": "roll%2Fup", "type": "rollup", "rollup": {"type": "incomplete"}} + page = native(properties={"Sum": prop}) + first = property_list( + prop["id"], + "rollup", + [], + cursor="next", + calculation={"type": "incomplete", "function": "sum", "incomplete": {}}, + ) + last = property_list( + prop["id"], "rollup", [], calculation={"type": "number", "function": "sum", "number": 13} + ) + connector, client = await source(gets=[page, first, last, page]) + owner = parent() + result = await connector.capture_page( + connector.child_scope(owner, "page_property"), + ScanContinuation(), + files=property_files, + parent=owner, + ) + record = result.records[0] + assert result.final and record.identity.native_id == prop["id"] + assert record.identity.container_id == PAGE and record.parent == owner.identity + assert record.payload["value_status"] == "available" and record.completeness == "partial" + assert record.source_updated_at is None + archive = json.loads(await property_files.storage.read_file(record.blobs[0].key)) + assert archive == { + "format_version": 1, + "page_id": PAGE, + "property_id": prop["id"], + "responses": [first, last], + "notion_version": "2026-03-11", + "page_last_edited_time": STAMP, + } + assert client.get.call_args_list[1].args[0].endswith("properties/roll%2Fup") + assert client.get.call_args_list[2].kwargs["params"]["start_cursor"] == "next" + assert record.payload["response_count"] == 2 + + +@pytest.mark.asyncio +async def test_unsupported_formula_is_explicit_retained_native_result(property_files): + prop = {"id": "calc", "type": "formula"} + page = native(properties={"Formula": prop}) + value = { + **prop, + "object": "property_item", + "formula": {"type": "unsupported", "unsupported": {}}, + } + connector, _ = await source(gets=[page, value, page]) + owner = parent() + result = await connector.capture_page( + connector.child_scope(owner, "page_property"), + ScanContinuation(), + files=property_files, + parent=owner, + ) + assert result.records[0].payload["value_status"] == "unsupported" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "case", + ["wrong_id", "unfinished_rollup", "shape_changed", "provider_error", "edited_page", "overflow"], +) +async def test_property_failure_never_publishes_partial_archive(case, property_files, monkeypatch): + from airweave.domains.entities.canonical.page_source import InvalidScanContinuation + + prop = {"id": "roll", "type": "rollup"} + page = native(properties={"Rollup": prop}) + value = property_list( + "wrong" if case == "wrong_id" else "roll", + "rollup", + [], + calculation={ + "type": "incomplete" if case == "unfinished_rollup" else "number", + "number": 3, + }, + ) + gets = [page, value, page] + expected = (ValueError, SourceError) + if case == "shape_changed": + gets = [ + page, + property_list("roll", "rollup", [], cursor="next", calculation={"type": "incomplete"}), + { + "object": "property_item", + "id": "roll", + "type": "rollup", + "rollup": {"type": "number", "number": 3}, + }, + ] + expected = InvalidScanContinuation + elif case == "provider_error": + gets = [page, response({"code": "restricted_resource"}, 403)] + expected = SourceEntityForbiddenError + elif case == "edited_page": + gets[-1] = {**page, "last_edited_time": "2026-09-30T01:00:00Z"} + expected = InvalidScanContinuation + elif case == "overflow": + monkeypatch.setattr("airweave.platform.sources.notion.MAX_FILE_SIZE_BYTES", 32) + expected = SourceError + connector, _ = await source(gets=gets) + owner = parent() + with pytest.raises(expected): + await connector.capture_page( + connector.child_scope(owner, "page_property"), + ScanContinuation(), + files=property_files, + parent=owner, + ) + assert await property_files.storage.list_files() == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", ["raw/slash", "bad%escape", "id?query=oops", "id#fragment"]) +async def test_unencoded_property_ids_never_enter_native_path(identity, property_files): + page = native(properties={"Field": {"id": identity, "type": "number"}}) + connector, client = await source(gets=[page]) + owner = parent() + with pytest.raises(ValidationError): + await connector.capture_page( + connector.child_scope(owner, "page_property"), + ScanContinuation(), + files=property_files, + parent=owner, + ) + assert client.get.call_count == 1 + + +@pytest.mark.asyncio +async def test_property_archive_roundtrips_unicode_and_empty_value(property_files): + import json + + prop = {"id": "p%25", "type": "rich_text", "rich_text": []} + page = native(properties={'Notes "日本語"': prop}) + value = {**property_list(prop["id"], "rich_text", []), "unknown": '日本語\n"escaped"'} + connector, _ = await source(gets=[page, value, page]) + owner = parent() + result = await connector.capture_page( + connector.child_scope(owner, "page_property"), + ScanContinuation(), + files=property_files, + parent=owner, + ) + record = result.records[0] + archive = json.loads(await property_files.storage.read_file(record.blobs[0].key)) + assert archive["responses"] == [value] + assert record.payload["name"] == 'Notes "日本語"' + assert archive["property_id"] == "p%25" From 3f6d496644a53541d4c13218ce3d14d5c627daf5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:39:02 -0700 Subject: [PATCH 084/266] Bound Gmail history page hydration and validate replay identity --- .../platform/sources/gmail_capture.py | 148 ++++++++++-- .../sources/tests/test_gmail_capture.py | 4 +- .../sources/tests/test_gmail_history_pages.py | 217 ++++++++++++++++++ 3 files changed, 349 insertions(+), 20 deletions(-) create mode 100644 backend/airweave/platform/sources/tests/test_gmail_history_pages.py diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py index 311234e395..e6af2f43d5 100644 --- a/backend/airweave/platform/sources/gmail_capture.py +++ b/backend/airweave/platform/sources/gmail_capture.py @@ -1,11 +1,15 @@ """Original Gmail message capture, independent of derived thread/search entities.""" import base64 +import hashlib +import json from collections.abc import AsyncGenerator, 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, @@ -14,6 +18,7 @@ StartedScope, ) from airweave.domains.entities.canonical.source import SourceObservation +from airweave.domains.sources.exceptions import SourceEntityNotFoundError from airweave.domains.storage.exceptions import FileSkippedException from airweave.domains.storage.file_service import FileService from airweave.domains.syncs.cursors.cursor import SyncCursor @@ -22,6 +27,90 @@ 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 HistoryBatch(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"): @@ -52,8 +141,8 @@ async def message(self, message_id: str) -> CaptureRecord: 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 as exc: - if exc.response.status_code != 404: + except (httpx.HTTPStatusError, SourceEntityNotFoundError) as exc: + if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code != 404: raise return CaptureRecord( identity=identity, @@ -124,27 +213,50 @@ async def body_bytes(self, message_id: str, body: dict) -> bytes: 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, + ) -> HistoryBatch: + """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") + end = min(offset + limit, len(page.message_ids)) + records = tuple([await self.message(mid) for mid in page.message_ids[offset:end]]) + return HistoryBatch(records=records, next_offset=end, complete=end == len(page.message_ids)) + async def history(self, boundary: str) -> tuple[list[str], str]: - """Drain history before emitting, so an expired boundary can restart cleanly.""" - params = {"startHistoryId": boundary, "maxResults": 500} + """Legacy generator bridge, removed when the durable page adapter is wired.""" affected: dict[str, None] = {} tokens: set[str] = set() + token = None while True: - page = await self.get(f"{BASE}/history", params=dict(params)) - for entry in page.get("history", []): - # `messages` is the full changed-ID list; typed lists can overlap it. - for message in entry.get("messages", []): - affected[message["id"]] = None - for field in ("messagesAdded", "messagesDeleted", "labelsAdded", "labelsRemoved"): - for change in entry.get(field, []): - affected[change["message"]["id"]] = None - token = page.get("nextPageToken") - if not token: - return list(affected), str(page["historyId"]) + page = await self.history_page(boundary, token) + affected.update(dict.fromkeys(page.message_ids)) + token = page.next_page_token + if token is None: + return list(affected), page.history_id if token in tokens: raise ValueError("Gmail history pagination repeated a token") tokens.add(token) - params["pageToken"] = token async def enumerate(self) -> AsyncGenerator[CaptureRecord, None]: """List every page in the provider's exact configured query scope.""" @@ -172,8 +284,8 @@ async def saved_history(self, cursor: SyncCursor | None) -> tuple[list[str], str return None try: return await self.history(boundary) - except httpx.HTTPStatusError as exc: - if exc.response.status_code != 404: + except (httpx.HTTPStatusError, SourceEntityNotFoundError) as exc: + if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code != 404: raise return None diff --git a/backend/airweave/platform/sources/tests/test_gmail_capture.py b/backend/airweave/platform/sources/tests/test_gmail_capture.py index 7573ca4911..967b7f14f9 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_capture.py +++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py @@ -181,8 +181,8 @@ async def test_legacy_entity_cursor_cannot_skip_canonical_bootstrap(): async def test_repeated_page_token_fails_without_advancing_cursor(): provider = Provider( [ - ("/history", {"nextPageToken": "same"}), - ("/history", {"nextPageToken": "same"}), + ("/history", {"nextPageToken": "same", "historyId": "21"}), + ("/history", {"nextPageToken": "same", "historyId": "21"}), ] ) state = cursor(history_id="20", canonical_query="") 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..15e9218144 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_gmail_history_pages.py @@ -0,0 +1,217 @@ +"""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, cursor, 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.saved_history(cursor(history_id="opaque", canonical_query="")) + + +@pytest.mark.asyncio +async def test_legacy_stale_history_handles_actual_translated_404(): + provider = Provider([("/history", translated_error(404))]) + assert ( + await GmailCapture(provider.get, None).saved_history( + cursor(history_id="opaque", canonical_query="") + ) + is None + ) + + +@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 From 45b9286d4c51bcf10e1941d2c08f124ba4e2d5ab Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:39:47 -0700 Subject: [PATCH 085/266] Separate full capture evidence from fenced provider changes checkpoints --- .../domains/entities/canonical/coverage.py | 39 +- .../entities/canonical/cycle_models.py | 76 +++- .../domains/entities/canonical/cycle_store.py | 111 ++++- .../domains/entities/canonical/requests.py | 8 + .../domains/entities/canonical/scan_models.py | 26 +- .../domains/entities/canonical/scan_store.py | 50 ++- .../canonical/tests/test_changes_cycles.py | 384 ++++++++++++++++++ 7 files changed, 665 insertions(+), 29 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_changes_cycles.py diff --git a/backend/airweave/domains/entities/canonical/coverage.py b/backend/airweave/domains/entities/canonical/coverage.py index efd1927cc6..3fabd5d2bb 100644 --- a/backend/airweave/domains/entities/canonical/coverage.py +++ b/backend/airweave/domains/entities/canonical/coverage.py @@ -3,7 +3,7 @@ from typing import Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -15,11 +15,23 @@ from airweave.models.sync_cursor import SyncCursor +class FullCaptureCoverage(BaseModel): + """Public coverage evidence excludes private provider continuation values.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + completed_at: AwareDatetime + discovery: Literal["incomplete", "scope_enumeration_complete"] + + class CaptureCoverage(BaseModel): """No cycle means unknown; completion only certifies the declared scope policies.""" model_config = ConfigDict(extra="forbid", frozen=True) phase: Literal["active", "complete"] + mode: Literal["full", "changes"] = "full" + last_full_capture: FullCaptureCoverage | None = None + provider_checkpoint_promoted_at: AwareDatetime | None = None policies: dict[str, CompletionPolicy] discovery: Literal["incomplete", "pending", "scope_enumeration_complete"] @@ -50,7 +62,30 @@ async def capture_coverage( if cycle.phase == "active" else "scope_enumeration_complete" ) + evidence = cycle.last_full_capture + if evidence is not None and evidence.configuration_digest != cycle.configuration.digest(): + continue + promoted = cycle.promoted_checkpoint + if cycle.mode == "changes": + if ( + evidence is None + or promoted is None + or (promoted.configuration_digest != cycle.configuration.digest()) + ): + continue + discovery = evidence.discovery result[row.sync_id] = CaptureCoverage( - phase=cycle.phase, policies=policies, discovery=discovery + phase=cycle.phase, + policies=policies, + discovery=discovery, + mode=cycle.mode, + last_full_capture=FullCaptureCoverage( + cycle_id=evidence.cycle_id, + completed_at=evidence.completed_at, + discovery=evidence.discovery, + ) + if evidence + else None, + provider_checkpoint_promoted_at=promoted.promoted_at if promoted else None, ) return result diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 27795a6b16..8740b89b91 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -1,16 +1,30 @@ """Typed engine-owned cycle state inside the existing SyncCursor.""" +import hashlib +import json from typing import Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import ( + AwareDatetime, + BaseModel, + ConfigDict, + Field, + JsonValue, + field_validator, + model_validator, +) from airweave.domains.entities.canonical.models import SourceRecord -from airweave.domains.entities.canonical.requests import RecordKind, WriterFence +from airweave.domains.entities.canonical.requests import RecordKind, ScanVersion, WriterFence CompletionPolicy = Literal["exhaustive", "discovery_only", "discovery_with_validation"] +CaptureMode = Literal["full", "changes"] +CompletedDiscovery = Literal["incomplete", "scope_enumeration_complete"] + CYCLE_KEY = "canonical_cycle" +TERMINAL_CHECKPOINT_KEY = "_canonical_terminal_checkpoint" class CycleConfiguration(BaseModel): @@ -21,6 +35,12 @@ class CycleConfiguration(BaseModel): parents: dict[RecordKind, tuple[RecordKind | None, ...]] completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) + def digest(self) -> str: + """Bind history to scope, topology and guarantees, including default policies.""" + value = self.model_dump(mode="json") + value["completion_policies"] = {kind: self.policy(kind) for kind in self.parents} + return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() + def policy(self, kind: str) -> CompletionPolicy: """Omitted declarations preserve historical exhaustive semantics.""" return self.completion_policies.get(kind, "exhaustive") @@ -110,6 +130,49 @@ class CycleVersion(BaseModel): revision: int = Field(ge=1) +class ProviderCheckpoint(BaseModel): + """Bounded connector state; distinct from the engine's change-sequence checkpoint.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + value: dict[str, JsonValue] + + @field_validator("value") + @classmethod + def bounded(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: + """Never store an unbounded provider response or work queue.""" + if not value or len(json.dumps(value, allow_nan=False).encode()) > 65536: + raise ValueError("Provider checkpoint must contain at most 64 KiB of state") + return value + + +class FullCaptureEvidence(BaseModel): + """A completed full pass certifies only its explicit discovery guarantee.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + completed_at: AwareDatetime + discovery: CompletedDiscovery + configuration_digest: str = Field(pattern=r"^[a-f0-9]{64}$") + + +class PromotedCheckpoint(BaseModel): + """Only successful cycle finalization advances this private provider boundary.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + checkpoint: ProviderCheckpoint + configuration_digest: str = Field(pattern=r"^[a-f0-9]{64}$") + promoted_at: AwareDatetime + + +class TerminalCheckpoint(BaseModel): + """Source-parsed boundary from an exact completed root scan, attested during publication.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + record_type: RecordKind + expected: ScanVersion + checkpoint: ProviderCheckpoint + + class CaptureCycle(BaseModel): """Completion stays durable until an explicit next-cycle CAS.""" @@ -119,6 +182,10 @@ class CaptureCycle(BaseModel): configuration: CycleConfiguration phase: Literal["active", "complete"] = "active" completed_job_id: UUID | None = None + mode: CaptureMode = "full" + starting_checkpoint: ProviderCheckpoint | None = None + promoted_checkpoint: PromotedCheckpoint | None = None + last_full_capture: FullCaptureEvidence | None = None @model_validator(mode="before") @classmethod @@ -138,6 +205,8 @@ class BeginCycle(BaseModel): fence: WriterFence configuration: CycleConfiguration expected: CycleVersion | None = None + mode: CaptureMode = "full" + starting_checkpoint: ProviderCheckpoint | None = None class CompleteCycle(BaseModel): @@ -146,6 +215,7 @@ class CompleteCycle(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) fence: WriterFence expected: CycleVersion + terminal_checkpoint: TerminalCheckpoint | None = None class RestartCycle(BaseModel): @@ -155,6 +225,8 @@ class RestartCycle(BaseModel): fence: WriterFence expected: CycleVersion configuration: CycleConfiguration + mode: CaptureMode = "full" + starting_checkpoint: ProviderCheckpoint | None = None class ScopeWork(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 763738504f..3109655e7b 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -10,10 +10,14 @@ from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint from airweave.domains.entities.canonical.cycle_models import ( CYCLE_KEY, + TERMINAL_CHECKPOINT_KEY, BeginCycle, CaptureCycle, CompleteCycle, CycleVersion, + FullCaptureEvidence, + PromotedCheckpoint, + ProviderCheckpoint, RestartCycle, ScopeWork, ) @@ -61,6 +65,40 @@ def persist(cursor: SyncCursor, state: CaptureCycle) -> None: cursor.last_updated = datetime.now(timezone.utc) +def next_cycle(previous: CaptureCycle | None, request: BeginCycle | RestartCycle) -> CaptureCycle: + """One immutable execution plan; incompatible history never becomes a delta baseline.""" + digest = request.configuration.digest() + evidence = previous.last_full_capture if previous else None + promoted = previous.promoted_checkpoint if previous else None + if evidence is not None and evidence.configuration_digest != digest: + evidence = None + if promoted is not None and promoted.configuration_digest != digest: + promoted = None + starting = request.starting_checkpoint + if request.mode == "changes": + if len(request.configuration.parents) != 1 or next( + iter(request.configuration.parents.values()) + ) != (None,): + raise CycleConflict("Changes currently require one independent root kind") + if evidence is None or promoted is None: + raise CycleConflict( + "Changes require a completed compatible full capture and checkpoint" + ) + if starting is not None and starting != promoted.checkpoint: + raise CycleConflict("Changes must start at the last promoted provider checkpoint") + starting = promoted.checkpoint + return CaptureCycle( + version=CycleVersion( + cycle_id=uuid4(), revision=1 if previous is None else previous.version.revision + 1 + ), + configuration=request.configuration, + mode=request.mode, + starting_checkpoint=starting, + promoted_checkpoint=promoted, + last_full_capture=evidence, + ) + + async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: """An explicit completed-version CAS starts the next cycle.""" cursor = await cursor_row(db, request.fence) @@ -73,6 +111,13 @@ async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: raise CycleConflict( "Active cycle configuration changed; explicit abandonment required" ) + if state.mode != request.mode or ( + request.starting_checkpoint is not None + and request.starting_checkpoint != state.starting_checkpoint + ): + raise CycleConflict( + "Active cycle execution plan changed; explicit restart required" + ) return state if request.expected is None: return state @@ -85,12 +130,7 @@ async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: cursor_data={}, ) db.add(cursor) - state = CaptureCycle( - version=CycleVersion( - cycle_id=uuid4(), revision=1 if state is None else state.version.revision + 1 - ), - configuration=request.configuration, - ) + state = next_cycle(state, request) persist(cursor, state) await db.flush() return state @@ -289,6 +329,43 @@ async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) return None +async def terminal_checkpoint( + db: AsyncSession, state: CaptureCycle, request: CompleteCycle, now: datetime +) -> PromotedCheckpoint | None: + """Attest the source's terminal scan before publishing its provider boundary.""" + terminal = request.terminal_checkpoint + if (state.starting_checkpoint is not None or state.mode == "changes") and terminal is None: + raise CycleConflict("Checkpoint-bearing capture requires its terminal scan boundary") + promoted = state.promoted_checkpoint if state.mode == "changes" else None + if terminal is not None: + if terminal.record_type not in state.configuration.root_record_types: + raise CycleConflict("Terminal checkpoint must belong to a declared root scope") + scan = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == request.fence.organization_id, + CaptureScan.sync_id == request.fence.sync_id, + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.record_type == terminal.record_type, + CaptureScan.container_id.is_(None), + CaptureScan.parent_record_id.is_(None), + CaptureScan.phase == "complete", + CaptureScan.sweep_id == terminal.expected.sweep_id, + CaptureScan.revision == terminal.expected.revision, + ) + ) + if scan is None: + raise CycleConflict("Terminal scan changed or is not complete") + committed = scan.continuation.get(TERMINAL_CHECKPOINT_KEY) + if committed is None or ProviderCheckpoint.model_validate(committed) != terminal.checkpoint: + raise CycleConflict("Provider checkpoint does not match committed terminal page") + promoted = PromotedCheckpoint( + checkpoint=terminal.checkpoint, + configuration_digest=state.configuration.digest(), + promoted_at=now, + ) + return promoted + + async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) -> CaptureCycle: """Publish the checkpoint only when SQL proves all currently visible scopes complete.""" cursor, state = await attest_cycle(db, request.fence, request.expected.cycle_id) @@ -315,9 +392,26 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - ) if missing is not None: raise CycleConflict("A visible parent still has an incomplete child scope") + now = datetime.now(timezone.utc) + promoted = await terminal_checkpoint(db, state, request, now) + evidence = state.last_full_capture + if state.mode == "full": + evidence = FullCaptureEvidence( + cycle_id=state.version.cycle_id, + completed_at=now, + configuration_digest=state.configuration.digest(), + discovery="scope_enumeration_complete" + if all( + state.configuration.policy(kind) == "exhaustive" + for kind in state.configuration.parents + ) + else "incomplete", + ) state = state.model_copy( update={ "phase": "complete", + "promoted_checkpoint": promoted, + "last_full_capture": evidence, "completed_job_id": request.fence.job_id, "version": CycleVersion( cycle_id=state.version.cycle_id, revision=state.version.revision + 1 @@ -341,10 +435,7 @@ async def restart_cycle(db: AsyncSession, request: RestartCycle) -> CaptureCycle cursor, previous = await attest_cycle(db, request.fence, request.expected.cycle_id) if previous.version != request.expected: raise CycleConflict("Cycle changed before explicit restart") - state = CaptureCycle( - version=CycleVersion(cycle_id=uuid4(), revision=previous.version.revision + 1), - configuration=request.configuration, - ) + state = next_cycle(previous, request) persist(cursor, state) await db.flush() return state diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index f4c3b8c79c..5984e403ed 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -11,6 +11,14 @@ RecordKind = Annotated[str, Field(min_length=1, max_length=200)] +class ScanVersion(BaseModel): + """Exact durable scope version observed before provider I/O.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + sweep_id: UUID + revision: int = Field(ge=1) + + class RecordIdentity(BaseModel): """Native identity inside one sync, including provider-required container.""" diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index 3848e49e96..838204ccc5 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -6,17 +6,18 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, field_validator -from airweave.domains.entities.canonical.cycle_models import CompletionPolicy +from airweave.domains.entities.canonical.cycle_models import ( + TERMINAL_CHECKPOINT_KEY, + CompletionPolicy, + ProviderCheckpoint, +) from airweave.domains.entities.canonical.models import CaptureResult, SourceRecord -from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, WriterFence - - -class ScanVersion(BaseModel): - """The exact durable state a caller observed before external I/O.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - sweep_id: UUID - revision: int = Field(ge=1) +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + ScanVersion, + WriterFence, +) class ScanContinuation(BaseModel): @@ -29,6 +30,8 @@ class ScanContinuation(BaseModel): @classmethod def bounded(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: """Bound serialized state before it enters a transaction.""" + if TERMINAL_CHECKPOINT_KEY in value: + raise ValueError("Source continuation uses a reserved engine key") if len(json.dumps(value, allow_nan=False, ensure_ascii=False).encode()) > 65536: raise ValueError("Scan continuation exceeds 64 KiB") return value @@ -62,6 +65,8 @@ class ScanState(BaseModel): parent_visibility_epoch: int | None = None membership_attempt_id: UUID | None = None completion_policy: CompletionPolicy = "exhaustive" + mode: Literal["full", "changes"] = "full" + provider_checkpoint: ProviderCheckpoint | None = None class CommitScanPage(BaseModel): @@ -80,6 +85,7 @@ class CommitScanPage(BaseModel): ) continuation: ScanContinuation final: bool = False + provider_checkpoint: ProviderCheckpoint | None = None class ReconcileScan(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index c83c69447f..b515ec5f42 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -4,10 +4,16 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 +from pydantic import JsonValue from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.cycle_models import CycleVersion +from airweave.domains.entities.canonical.cycle_models import ( + TERMINAL_CHECKPOINT_KEY, + CaptureMode, + CycleVersion, + ProviderCheckpoint, +) from airweave.domains.entities.canonical.cycle_store import ( CycleConflict, attest_cycle, @@ -72,7 +78,18 @@ def scan_state(row: CaptureScan, parent: Entity | None = None) -> ScanState: version=ScanVersion(sweep_id=row.sweep_id, revision=row.revision), phase=row.phase, fingerprint=row.fingerprint, - continuation=ScanContinuation(value=row.continuation), + continuation=ScanContinuation( + value={ + key: value + for key, value in row.continuation.items() + if key != TERMINAL_CHECKPOINT_KEY + } + ), + provider_checkpoint=ProviderCheckpoint.model_validate( + row.continuation[TERMINAL_CHECKPOINT_KEY] + ) + if TERMINAL_CHECKPOINT_KEY in row.continuation + else None, started_at=row.started_at, completed_at=row.completed_at, parent_visibility_epoch=row.parent_visibility_epoch, @@ -80,6 +97,20 @@ def scan_state(row: CaptureScan, parent: Entity | None = None) -> ScanState: ) +def page_continuation(request: CommitScanPage, mode: CaptureMode) -> dict[str, JsonValue]: + """The terminal boundary commits with the records it certifies, never in a later call.""" + if request.provider_checkpoint is not None and not request.final: + raise ScanConflict("Only a terminal page may certify a provider checkpoint") + if mode == "changes" and request.final and request.provider_checkpoint is None: + raise ScanConflict("Changes terminal page requires a provider checkpoint") + if mode == "changes" and any(record.removal_reason == "absent" for record in request.records): + raise ScanConflict("Changes require explicit deletion evidence, never absence") + continuation = dict(request.continuation.value) + if request.provider_checkpoint is not None: + continuation[TERMINAL_CHECKPOINT_KEY] = request.provider_checkpoint.model_dump(mode="json") + return continuation + + class CanonicalScanStore: """Reuse record capture, lock ordering and journal semantics for recoverable scans.""" @@ -129,7 +160,8 @@ async def _state(self, db: AsyncSession, row: CaptureScan) -> ScanState: update={ "completion_policy": cycle.configuration.policy(row.record_type) if cycle - else "exhaustive" + else "exhaustive", + "mode": cycle.mode if cycle else "full", } ) @@ -230,6 +262,7 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: """Capture and advance the page as one transaction, never a partial acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) + continuation = page_continuation(request, cycle.mode) await attest_scope(db, request.fence, cycle, request.scope) row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id @@ -287,10 +320,11 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: sequence=discovered.sequence, unchanged=captured.unchanged + discovered.unchanged, ) - row.continuation = request.continuation.value + row.continuation = continuation row.revision += 1 if request.final: - row.phase = "reconciling" + row.phase = "complete" if cycle.mode == "changes" else "reconciling" + row.completed_at = datetime.now(timezone.utc) if cycle.mode == "changes" else None await db.flush() return ScanResult(state=await self._state(db, row), capture=captured) @@ -346,6 +380,8 @@ async def missing( """Bounded exact inventory omissions, only after current-attempt collection.""" await self.records._fenced_sync(db, fence) _, cycle = await attest_cycle(db, fence, state.cycle_id) + if cycle.mode == "changes": + raise ScanConflict("Changes scans cannot reconcile or validate enumeration absence") await attest_scope(db, fence, cycle, state.scope) row = self._expect(await self._row(db, fence, state.scope), state.version, state.cycle_id) parent = await scope_owner(db, fence, cycle, state.scope) @@ -387,6 +423,8 @@ async def omission(self, db: AsyncSession, request: CommitOmission) -> ScanResul """Commit exact known-object validation with its sweep acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.state.cycle_id) + if cycle.mode == "changes": + raise ScanConflict("Changes scans cannot reconcile or validate enumeration absence") if ( cycle.configuration.policy(request.state.scope.record_type) != "discovery_with_validation" @@ -425,6 +463,8 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul """A final page is necessary; absence completion is durable and bounded.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) + if cycle.mode == "changes": + raise ScanConflict("Changes scans cannot reconcile or validate enumeration absence") await attest_scope(db, request.fence, cycle, request.scope) row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id diff --git a/backend/airweave/domains/entities/canonical/tests/test_changes_cycles.py b/backend/airweave/domains/entities/canonical/tests/test_changes_cycles.py new file mode 100644 index 0000000000..06f5f261cb --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_changes_cycles.py @@ -0,0 +1,384 @@ +"""Real SQL: provider progress never substitutes for committed capture or inventory evidence.""" + +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.coverage import capture_coverage +from airweave.domains.entities.canonical.cycle_models import ( + BeginCycle, + CompleteCycle, + CycleConfiguration, + ProviderCheckpoint, + RestartCycle, + TerminalCheckpoint, +) +from airweave.domains.entities.canonical.cycle_store import CycleConflict +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, +) +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.store import StaleWriter +from airweave.models.entity import Entity +from airweave.models.sync_cursor import SyncCursor + +CONFIG = CycleConfiguration(fingerprint="a" * 64, parents={"message": (None,)}) +SCOPE = CompletedScope(record_type="message") + + +def checkpoint(value): + return ProviderCheckpoint(value={"history_id": value}) + + +def record(native_id, **kwargs): + return CaptureRecord( + identity=RecordIdentity(record_type="message", native_id=native_id), + payload={"id": native_id}, + observed_at=datetime.now(timezone.utc), + **kwargs, + ) + + +async def begin(database, service, fence, **kwargs): + async with database() as db: + return await service.begin_cycle( + db, BeginCycle(fence=fence, configuration=CONFIG, **kwargs) + ) + + +async def start(database, service, fence, cycle): + async with database() as db: + previous = await service.read_scan(db, fence, SCOPE) + return await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=SCOPE, + cycle_id=cycle.version.cycle_id, + fingerprint=CONFIG.fingerprint, + expected=previous.version if previous else None, + ), + ) + + +async def page(database, service, fence, scan, *records, final=True, value=None): + async with database() as db: + result = await service.commit_scan_page( + db, + CommitScanPage( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + records=records, + continuation=ScanContinuation(value={"terminal": "200"}), + final=final, + provider_checkpoint=checkpoint( + value or ("300" if scan.mode == "changes" else "200") + ) + if final + else None, + ), + ) + return result.state + + +async def complete(database, service, fence, cycle, scan, value="200"): + async with database() as db: + cycle = await service.read_cycle(db, fence) + return await service.complete_cycle( + db, + CompleteCycle( + fence=fence, + expected=cycle.version, + terminal_checkpoint=TerminalCheckpoint( + record_type="message", expected=scan.version, checkpoint=checkpoint(value) + ), + ), + ) + + +async def baseline(database, service, fence): + cycle = await begin(database, service, fence, starting_checkpoint=checkpoint("100")) + scan = await start(database, service, fence, cycle) + scan = await page(database, service, fence, scan, record("A"), record("B")) + assert scan.phase == "reconciling" + async with database() as db: + result = await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + observed_at=datetime.now(timezone.utc), + ), + ) + return await complete(database, service, fence, cycle, result.state) + + +async def test_changes_commit_without_absence_and_promote_only_after_terminal(database, source): + service, fence = source + full = await baseline(database, service, fence) + delta = await begin(database, service, fence, mode="changes", expected=full.version) + assert delta.starting_checkpoint == checkpoint("200") + scan = await start(database, service, fence, delta) + scan = await page(database, service, fence, scan, record("A"), final=False) + async with database() as db: + state = await service.read_cycle(db, fence) + assert state.promoted_checkpoint == full.promoted_checkpoint + with pytest.raises(CycleConflict, match="incomplete"): + await complete(database, service, fence, delta, scan, "300") + with pytest.raises(ScanConflict, match="Changes"): + async with database() as db: + await service.scan_missing(db, fence, scan) + with pytest.raises(ScanConflict, match="Changes"): + async with database() as db: + await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + observed_at=datetime.now(timezone.utc), + ), + ) + scan = await page(database, service, fence, scan) + assert scan.phase == "complete" and scan.completed_at is not None + # A fresh session recovers after a crash between the final page and cycle publication. + async with database() as db: + scan = await service.read_scan(db, fence, SCOPE) + finished = await complete(database, service, fence, delta, scan, "300") + assert finished.last_full_capture == full.last_full_capture + assert finished.promoted_checkpoint.checkpoint == checkpoint("300") + assert await begin(database, service, fence, mode="changes") == finished # lost ACK + async with database() as db: + rows = (await db.scalars(select(Entity))).all() + assert {r.native_id for r in rows if r.deleted_at is None} == {"A", "B"} + cursor = await db.scalar(select(SyncCursor)) + assert cursor.cursor_data["canonical_checkpoint"]["observed_change_sequence"] >= 2 + coverage = (await capture_coverage(db, fence.organization_id, (fence.sync_id,)))[ + fence.sync_id + ] + assert coverage.mode == "changes" and coverage.discovery == "scope_enumeration_complete" + assert "history_id" not in coverage.model_dump_json() + assert "configuration_digest" not in coverage.model_dump_json() + + +async def test_changes_rejects_no_baseline_wrong_start_and_changed_guarantee(database, source): + service, fence = source + with pytest.raises(CycleConflict, match="compatible"): + await begin(database, service, fence, mode="changes") + full = await baseline(database, service, fence) + with pytest.raises(CycleConflict, match="last promoted"): + await begin( + database, + service, + fence, + mode="changes", + expected=full.version, + starting_checkpoint=checkpoint("999"), + ) + changed = CONFIG.model_copy(update={"completion_policies": {"message": "discovery_only"}}) + with pytest.raises(CycleConflict, match="compatible"): + async with database() as db: + await service.begin_cycle( + db, + BeginCycle( + fence=fence, configuration=changed, expected=full.version, mode="changes" + ), + ) + + +async def test_terminal_version_and_lost_writer_cannot_publish(database, source): + service, fence = source + full = await baseline(database, service, fence) + delta = await begin(database, service, fence, mode="changes", expected=full.version) + first = await start(database, service, fence, delta) + final = await page(database, service, fence, first, record("A")) + with pytest.raises(CycleConflict, match="Terminal scan"): + await complete(database, service, fence, delta, first, "300") + async with database() as db: + newer = await service.activate_writer( + db, + organization_id=fence.organization_id, + sync_id=fence.sync_id, + job_id=fence.job_id, + attempt_id=uuid4(), + attempt_number=2, + ) + with pytest.raises(StaleWriter): + await complete(database, service, fence, delta, final, "300") + async with database() as db: + state = await service.read_cycle(db, newer) + assert state.promoted_checkpoint == full.promoted_checkpoint + finished = await complete(database, service, newer, delta, final, "300") + assert finished.completed_job_id == newer.job_id + + +async def test_restart_preserves_successful_history_not_partial_delta(database, source): + service, fence = source + full = await baseline(database, service, fence) + delta = await begin(database, service, fence, mode="changes", expected=full.version) + scan = await start(database, service, fence, delta) + with pytest.raises(ScanConflict, match="never absence"): + await page( + database, service, fence, scan, record("A", kind="delete", removal_reason="absent") + ) + async with database() as db: + delta = await service.read_cycle(db, fence) + fresh = await service.restart_cycle( + db, + RestartCycle( + fence=fence, + expected=delta.version, + configuration=CONFIG, + starting_checkpoint=checkpoint("500"), + ), + ) + assert fresh.mode == "full" and fresh.starting_checkpoint == checkpoint("500") + assert fresh.promoted_checkpoint == full.promoted_checkpoint + assert fresh.last_full_capture == full.last_full_capture + + +async def test_exact_terminal_version_cannot_publish_different_checkpoint(database, source): + service, fence = source + full = await baseline(database, service, fence) + delta = await begin(database, service, fence, mode="changes", expected=full.version) + scan = await page(database, service, fence, await start(database, service, fence, delta)) + with pytest.raises(CycleConflict, match="does not match committed"): + await complete(database, service, fence, delta, scan, "forged") + async with database() as db: + persisted = await service.read_scan(db, fence, SCOPE) + state = await service.read_cycle(db, fence) + assert persisted.provider_checkpoint == checkpoint("300") + assert persisted.continuation.value == {"terminal": "200"} + assert state.promoted_checkpoint == full.promoted_checkpoint + + +async def test_full_refresh_without_checkpoint_clears_old_delta_authority(database, source): + service, fence = source + old = await baseline(database, service, fence) + fresh = await begin(database, service, fence, expected=old.version) + scan = await start(database, service, fence, fresh) + async with database() as db: + result = await service.commit_scan_page( + db, + CommitScanPage( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + records=(record("B"),), + continuation=ScanContinuation(), + final=True, + ), + ) + scan = result.state + async with database() as db: + await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + observed_at=datetime.now(timezone.utc), + ), + ) + current = await service.read_cycle(db, fence) + done = await service.complete_cycle( + db, CompleteCycle(fence=fence, expected=current.version) + ) + assert done.promoted_checkpoint is None + assert done.last_full_capture.cycle_id == fresh.version.cycle_id + with pytest.raises(CycleConflict, match="compatible"): + await begin(database, service, fence, mode="changes", expected=done.version) + async with database() as db: + assert ( + await db.scalar(select(Entity).where(Entity.native_id == "A")) + ).deleted_at is not None + + +async def test_incomplete_full_evidence_stays_incomplete_during_changes(database, source): + service, fence = source + config = CONFIG.model_copy(update={"completion_policies": {"message": "discovery_only"}}) + async with database() as db: + cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + scan = await start(database, service, fence, cycle) + scan = await page(database, service, fence, scan, record("A")) + async with database() as db: + result = await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + observed_at=datetime.now(timezone.utc), + ), + ) + done = await complete(database, service, fence, cycle, result.state) + async with database() as db: + await service.begin_cycle( + db, BeginCycle(fence=fence, configuration=config, expected=done.version, mode="changes") + ) + coverage = (await capture_coverage(db, fence.organization_id, (fence.sync_id,)))[ + fence.sync_id + ] + assert coverage.phase == "active" and coverage.mode == "changes" + assert coverage.discovery == coverage.last_full_capture.discovery == "incomplete" + assert coverage.provider_checkpoint_promoted_at is not None + + +async def test_changes_forest_is_explicitly_unsupported(database, source): + service, fence = source + config = CycleConfiguration( + fingerprint="b" * 64, parents={"channel": (None,), "message": ("channel",)} + ) + with pytest.raises(CycleConflict, match="one independent root"): + async with database() as db: + await service.begin_cycle( + db, BeginCycle(fence=fence, configuration=config, mode="changes") + ) + + +@pytest.mark.parametrize("final, boundary", [(False, "early"), (True, None)]) +async def test_invalid_checkpoint_page_rolls_back_records_and_progress( + database, source, final, boundary +): + service, fence = source + full = await baseline(database, service, fence) + delta = await begin(database, service, fence, mode="changes", expected=full.version) + scan = await start(database, service, fence, delta) + with pytest.raises(ScanConflict, match="[Tt]erminal"): + async with database() as db: + await service.commit_scan_page( + db, + CommitScanPage( + fence=fence, + scope=SCOPE, + cycle_id=scan.cycle_id, + expected=scan.version, + records=(record("uncommitted"),), + continuation=ScanContinuation(), + final=final, + provider_checkpoint=checkpoint(boundary) if boundary else None, + ), + ) + async with database() as db: + assert await service.read_scan(db, fence, SCOPE) == scan + assert await db.scalar(select(Entity.id).where(Entity.native_id == "uncommitted")) is None + state = await service.read_cycle(db, fence) + assert state.promoted_checkpoint == full.promoted_checkpoint From aa6d3587b2cd82443765b600f124ec747fe9fbce Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:47:29 -0700 Subject: [PATCH 086/266] Wire checkpointed page sources through existing capture lifecycle --- .../entities/canonical/cycle_models.py | 13 +- .../domains/entities/canonical/page_source.py | 35 ++- .../domains/entities/canonical/scan_store.py | 15 +- .../canonical/tests/test_checkpoint_driver.py | 253 ++++++++++++++++++ .../sync_pipeline/canonical_capture.py | 26 +- .../domains/sync_pipeline/canonical_scan.py | 68 ++++- 6 files changed, 398 insertions(+), 12 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 8740b89b91..713f0d1d57 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -34,10 +34,13 @@ class CycleConfiguration(BaseModel): fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") parents: dict[RecordKind, tuple[RecordKind | None, ...]] completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) + known_object_validation: tuple[RecordKind, ...] = () def digest(self) -> str: """Bind history to scope, topology and guarantees, including default policies.""" value = self.model_dump(mode="json") + if not self.known_object_validation: + value.pop("known_object_validation") # Preserve existing default configuration digests. value["completion_policies"] = {kind: self.policy(kind) for kind in self.parents} return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() @@ -53,6 +56,7 @@ def from_source( record_types: tuple[str, ...], container_parents: dict[str, str | tuple[str | None, ...]], completion_policies: dict[RecordKind, CompletionPolicy] | None = None, + known_object_validation: tuple[RecordKind, ...] = (), ) -> "CycleConfiguration": """Snapshot the existing source declaration; no second topology registry.""" if len(set(record_types)) != len(record_types): @@ -64,7 +68,10 @@ def from_source( declared = container_parents.get(kind, (None,)) parents[kind] = (declared,) if isinstance(declared, str) else declared return cls( - fingerprint=fingerprint, parents=parents, completion_policies=completion_policies or {} + fingerprint=fingerprint, + parents=parents, + completion_policies=completion_policies or {}, + known_object_validation=known_object_validation, ) @model_validator(mode="before") @@ -86,6 +93,10 @@ def upgrade_flat(cls, value): @model_validator(mode="after") def reachable_types(self) -> "CycleConfiguration": """Type cycles are valid; every declared kind must be reachable from a root.""" + if len(set(self.known_object_validation)) != len(self.known_object_validation) or ( + set(self.known_object_validation) - set(self.parents) + ): + raise ValueError("Known-object validation requires unique declared kinds") if set(self.completion_policies) - set(self.parents): raise ValueError("Completion policy has an undeclared record kind") if not self.parents or any(not kind or not values for kind, values in self.parents.items()): diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index a62526c60e..c6425c1f48 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -4,7 +4,12 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter -from airweave.domains.entities.canonical.cycle_models import CycleConfiguration +from airweave.domains.entities.canonical.cycle_models import ( + CaptureCycle, + CaptureMode, + CycleConfiguration, + ProviderCheckpoint, +) from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.requests import ( CaptureRecord, @@ -27,6 +32,32 @@ class CapturePage(BaseModel): ) 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 + + +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): @@ -77,7 +108,7 @@ async def confirm_absent(self, record: SourceRecord) -> None: @runtime_checkable class KnownObjectSource(Protocol): - """Only discovery-with-validation sources implement exact current reads.""" + """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.""" diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index b515ec5f42..18a44c79ba 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -410,7 +410,10 @@ async def missing( Entity.record_revision > 0, Entity.last_seen_run_id.is_distinct_from(row.sweep_id), ] - if cycle.configuration.policy(row.record_type) != "discovery_with_validation": + if ( + cycle.configuration.policy(row.record_type) != "discovery_with_validation" + and row.record_type not in cycle.configuration.known_object_validation + ): predicates.append(Entity.deleted_at.is_(None)) if after is not None: predicates.append(Entity.id > after) @@ -428,6 +431,7 @@ async def omission(self, db: AsyncSession, request: CommitOmission) -> ScanResul if ( cycle.configuration.policy(request.state.scope.record_type) != "discovery_with_validation" + and request.state.scope.record_type not in cycle.configuration.known_object_validation ): raise ScanConflict("This scope does not accept known-object validation") candidates = await self.missing(db, request.fence, request.state) @@ -480,11 +484,12 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul if row.phase != "reconciling": raise ScanConflict("Only fully collected scans can reconcile absence") policy = cycle.configuration.policy(row.record_type) + if ( + policy == "discovery_with_validation" + or row.record_type in cycle.configuration.known_object_validation + ) and await self.missing(db, request.fence, await self._state(db, row)): + raise ScanConflict("Known objects still require exact validation") if policy != "exhaustive": - if policy == "discovery_with_validation" and await self.missing( - db, request.fence, await self._state(db, row) - ): - raise ScanConflict("Known objects still require exact validation") row.phase = "complete" row.completed_at = datetime.now(timezone.utc) row.revision += 1 diff --git a/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py b/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py new file mode 100644 index 0000000000..992b1806a9 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py @@ -0,0 +1,253 @@ +"""Actual driver/pipeline and PostgreSQL, with an explicitly synthetic native-page source.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import CycleConfiguration, ProviderCheckpoint +from airweave.domains.entities.canonical.page_source import ( + CapturePage, + CapturePlan, + InvalidCaptureCheckpoint, +) +from airweave.domains.entities.canonical.requests import CaptureBatch +from airweave.domains.entities.canonical.scan_models import ScanContinuation +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.sync_pipeline.canonical_capture import CanonicalCapturePipeline +from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt +from airweave.models.entity import Entity + + +class NativeFixture: + canonical_record_types = ("message",) + canonical_container_parents = {} + capture_cycle_configuration = CycleConfiguration( + fingerprint="f" * 64, parents={"message": (None,)}, known_object_validation=("message",) + ) + + def __init__(self, *, stop=False, expired=0): + self.stop, self.expired = stop, expired + self.plans = [] + self.pages = [] + self.refreshes = [] + + async def prepare_cycle(self, previous): + self.plans.append(previous) + return CapturePlan( + starting_checkpoint=ProviderCheckpoint(value={"start": str(len(self.plans))}) + ) + + def initial_continuation(self, cycle): + return ScanContinuation( + value={"start": next(iter(cycle.starting_checkpoint.value.values())), "page": 0} + ) + + async def capture_page(self, scope, continuation, *, files, parent=None): + self.pages.append(continuation.value) + if self.expired: + self.expired -= 1 + raise InvalidCaptureCheckpoint("synthetic expired history") + if continuation.value["page"] == 0: + return CapturePage( + records=(record("listed"),), + continuation=ScanContinuation(value={**continuation.value, "page": 1}), + ) + if self.stop: + raise ConnectionError("synthetic interruption") + return CapturePage( + records=(), + continuation=ScanContinuation(value={**continuation.value, "page": 2}), + final=True, + provider_checkpoint=ProviderCheckpoint(value={"terminal": "done"}), + ) + + async def refresh_known(self, old, *, files): + self.refreshes.append(old.identity.native_id) + return record(old.identity.native_id) + + def child_scope(self, parent, record_type): + raise AssertionError("No child scope") + + async def confirm_absent(self, record): + raise AssertionError("Exact hydration must replace speculative absence") + + +def pipeline(database, source, connector, *, attempt=1): + service, fence = source + ctx, _, runtime, bus = components(database, source) + capture = 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.canonical_capture = capture + return ctx, runtime, capture + + +async def test_actual_pipeline_resumes_plan_and_hydrates_known_omissions(database, source): + service, fence = source + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(record("omitted"),))) + connector = NativeFixture(stop=True) + ctx, runtime, capture = pipeline(database, source, connector) + await capture.start(ctx) + with pytest.raises(ConnectionError, match="interruption"): + await capture.run_scans(ctx, runtime, AsyncMock()) + connector.stop = False + ctx, runtime, capture = pipeline(database, source, connector, attempt=2) + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + assert len(connector.plans) == 1 + assert [p["page"] for p in connector.pages] == [0, 1, 1] + assert connector.refreshes == ["omitted"] + async with database() as db: + rows = (await db.scalars(select(Entity))).all() + assert {r.native_id for r in rows if r.deleted_at is None} == {"listed", "omitted"} + cycle = await service.read_cycle(db, capture._writer()) + assert cycle.phase == "complete" + assert cycle.promoted_checkpoint.checkpoint.value == {"terminal": "done"} + assert cycle.last_full_capture.discovery == "scope_enumeration_complete" + + +@pytest.mark.parametrize("expirations", [1, 2]) +async def test_expired_checkpoint_restarts_whole_cycle_once(database, source, expirations): + connector = NativeFixture(expired=expirations) + ctx, runtime, capture = pipeline(database, source, connector) + await capture.start(ctx) + if expirations == 2: + with pytest.raises(InvalidCaptureCheckpoint): + await capture.run_scans(ctx, runtime, AsyncMock()) + else: + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + assert len(connector.plans) == 2 + assert [p["start"] for p in connector.pages[:2]] == ["1", "2"] + async with database() as db: + current = await source[0].read_cycle(db, capture._writer()) + assert current.starting_checkpoint.value == {"start": "2"} + assert current.phase == ("active" if expirations == 2 else "complete") + if expirations == 2: + assert current.promoted_checkpoint is None + + +async def test_driver_runs_changes_without_enumeration_or_omission_cleanup(database, source): + from airweave.domains.entities.canonical.cycle_models import BeginCycle + + connector = NativeFixture() + ctx, runtime, capture = pipeline(database, source, connector) + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + service, fence = source + async with database() as db: + await service.capture(db, CaptureBatch(fence=fence, records=(record("not-in-delta"),))) + previous = await service.read_cycle(db, fence) + await service.begin_cycle( + db, + BeginCycle( + fence=fence, + configuration=connector.capture_cycle_configuration, + expected=previous.version, + mode="changes", + ), + ) + connector.pages.clear() + ctx, runtime, capture = pipeline(database, source, connector, attempt=2) + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + assert connector.refreshes == [] + assert len(connector.plans) == 1 + async with database() as db: + retained = await db.scalar(select(Entity).where(Entity.native_id == "not-in-delta")) + current = await service.read_cycle(db, capture._writer()) + assert retained.deleted_at is None + assert current.mode == "changes" and current.phase == "complete" + assert current.last_full_capture == previous.last_full_capture + + +async def test_exhaustive_known_validation_is_enforced_by_sql(database, source): + from datetime import datetime, timezone + + from airweave.domains.entities.canonical.cycle_models import BeginCycle + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ( + BeginScan, + CommitOmission, + CommitScanPage, + ReconcileScan, + ) + from airweave.domains.entities.canonical.scan_store import ScanConflict + + service, fence = source + config = NativeFixture.capture_cycle_configuration + scope = CompletedScope(record_type="message") + async with database() as db: + old = ( + (await service.capture(db, CaptureBatch(fence=fence, records=(record("omitted"),)))) + .changes[0] + .record + ) + cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=config)) + state = await service.begin_scan( + db, + BeginScan( + fence=fence, + scope=scope, + cycle_id=cycle.version.cycle_id, + fingerprint=config.fingerprint, + ), + ) + state = ( + await service.commit_scan_page( + db, + CommitScanPage( + fence=fence, + scope=scope, + cycle_id=state.cycle_id, + expected=state.version, + records=(), + continuation=ScanContinuation(), + final=True, + ), + ) + ).state + with pytest.raises(ScanConflict, match="still require"): + async with database() as db: + await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=scope, + cycle_id=state.cycle_id, + expected=state.version, + observed_at=datetime.now(timezone.utc), + ), + ) + async with database() as db: + refreshed = await service.commit_omission( + db, + CommitOmission( + fence=fence, state=state, expected_record=old, observation=record("omitted") + ), + ) + result = await service.reconcile_scan( + db, + ReconcileScan( + fence=fence, + scope=scope, + cycle_id=state.cycle_id, + expected=refreshed.state.version, + observed_at=datetime.now(timezone.utc), + ), + ) + assert result.state.phase == "complete" and result.capture.changes == () diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index a87ddfd7ab..7d66ce702d 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -12,6 +12,7 @@ CaptureCycle, CompleteCycle, CycleConfiguration, + TerminalCheckpoint, ) from airweave.domains.entities.canonical.models import CaptureResult from airweave.domains.entities.canonical.page_source import CanonicalPageSource @@ -68,6 +69,7 @@ def __init__( record_types=record_types, container_parents=self._container_parents, completion_policies=config.completion_policies, + known_object_validation=config.known_object_validation, ) if config != declared: raise ValueError("Page source cycle must match its declared container topology") @@ -313,9 +315,13 @@ async def save_checkpoint(self, sync_context: SyncContext, runtime: SyncRuntime) 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) + db, + CompleteCycle( + fence=self._writer(), expected=current.version, terminal_checkpoint=terminal + ), ) return if runtime.cursor is None or not runtime.cursor.cursor_data: @@ -325,6 +331,24 @@ async def save_checkpoint(self, sync_context: SyncContext, runtime: SyncRuntime) 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 index 02ca0aa103..5dcf1c4047 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -7,11 +7,14 @@ from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle +from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle, RestartCycle 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, KnownObjectSource, ScopeAccessLost, @@ -27,6 +30,7 @@ CommitOmission, CommitScanPage, ReconcileScan, + ScanContinuation, ScanState, ) from airweave.domains.entities.canonical.service import CanonicalCaptureService @@ -62,10 +66,55 @@ async def run(self) -> CaptureCycle: expected = current.version else: expected = None + plan = ( + CapturePlan(mode=current.mode, starting_checkpoint=current.starting_checkpoint) + if current is not None and current.phase == "active" + else ( + await self.source.prepare_cycle(current) + if isinstance(self.source, CheckpointedPageSource) + else CapturePlan() + ) + ) async with self.sessions() as db: cycle = await self.service.begin_cycle( - db, BeginCycle(fence=self.fence, configuration=configuration, expected=expected) + db, + BeginCycle( + fence=self.fence, + configuration=configuration, + expected=expected, + mode=plan.mode, + starting_checkpoint=plan.starting_checkpoint, + ), ) + for restart in range(2): + try: + return await self.run_cycle(cycle) + except InvalidCaptureCheckpoint: + if restart or not isinstance(self.source, CheckpointedPageSource): + raise + await self.check_limits() + plan = await self.source.prepare_cycle(None) + if plan.mode != "full": + raise CycleConflict("An invalid checkpoint 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 != cycle.version.cycle_id: + raise CycleConflict("Cycle changed during checkpoint recovery") + cycle = await self.service.restart_cycle( + db, + RestartCycle( + fence=self.fence, + expected=current.version, + configuration=configuration, + mode=plan.mode, + starting_checkpoint=plan.starting_checkpoint, + ), + ) + raise AssertionError("Unreachable checkpoint recovery") + + async def run_cycle(self, cycle: CaptureCycle) -> CaptureCycle: + """Process the current frontier; every acknowledgement remains in the existing SQL store.""" + configuration = cycle.configuration while True: async with self.sessions() as db: work = await self.service.next_scope_work(db, self.fence, cycle.version.cycle_id) @@ -132,6 +181,7 @@ async def scan( fingerprint=cycle.configuration.fingerprint, expected=previous.version if previous else None, restart=restart, + continuation=self.initial_continuation(cycle), ), ) restarts = 0 @@ -155,6 +205,7 @@ async def scan( fingerprint=state.fingerprint, expected=state.version, restart=True, + continuation=self.initial_continuation(cycle), ), ) continue @@ -171,17 +222,28 @@ async def scan( discovered_records=page.discovered_records, continuation=page.continuation, final=page.final, + provider_checkpoint=page.provider_checkpoint, ), ) state = result.state await self.progress(result.capture, (*page.records, *page.discovered_records)) if state.phase == "reconciling": - if state.completion_policy == "discovery_with_validation": + 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) + 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.""" From 75e318b8f95e13b1b532faa8938600a572812f65 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:51:12 -0700 Subject: [PATCH 087/266] Verify full capture metadata in indexed search response --- .../domains/entities/canonical/tests/test_owned_search.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 44db15469e..817e288fd8 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -478,6 +478,9 @@ async def test_capture_discovery_coverage_is_separate_from_retrieval( body = response.json() assert body["sources"][0]["capture"] == { "phase": "active", + "mode": "full", + "last_full_capture": None, + "provider_checkpoint_promoted_at": None, "policies": {"event": "discovery_only"}, "discovery": "incomplete", } From d557252b4f16db02d21d51d7dffedea3a5591c7b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 09:58:19 -0700 Subject: [PATCH 088/266] Distinguish verified inference CI from pending staging readiness --- deploy/staging/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deploy/staging/README.md b/deploy/staging/README.md index 32a6d8ff55..981b0f225b 100644 --- a/deploy/staging/README.md +++ b/deploy/staging/README.md @@ -138,5 +138,7 @@ returned404, so model retrieval assertions did not run. Startup validation, Compose, staging readiness and the integration probe now use `/.well-known/ready`, matching the [upstream inference API](https://docs.weaviate.io/weaviate/modules/custom-modules). The integration test requires204 and still verifies actual384-dimensional vectors -and retrieval afterward. This fixes the contract; a successful rerun is required -before claiming the pinned deployment works. +and retrieval afterward. Later CI36744440692 at2bb84cf passed the real-model +retrieval lane and image build. That verifies the disposable CI topology; it +does not establish Kubernetes admission, deployed readiness, capacity or restore +for these staging templates. From 7d42c8a634a57499f80a7dcc4b696577b1a84457 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:03:13 -0700 Subject: [PATCH 089/266] Prepare bounded fresh-process Gmail recovery probe --- .../tests/test_gmail_resume_probe.py | 171 ++++++++++++++++++ backend/tests/live/README.md | 39 +++- backend/tests/live/gmail_resume.py | 105 +++++++++++ backend/tests/live/provider_lifecycle.py | 170 +++++++++++++++-- backend/tests/live/provider_sample.py | 14 +- 5 files changed, 480 insertions(+), 19 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py create mode 100644 backend/tests/live/gmail_resume.py 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..5d04eca2e7 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py @@ -0,0 +1,171 @@ +"""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.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( + 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" diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index c4d92246e1..f281c39616 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -104,7 +104,44 @@ 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. -## Complete configured Gmail reconciliation lifecycle +## 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. diff --git a/backend/tests/live/gmail_resume.py b/backend/tests/live/gmail_resume.py new file mode 100644 index 0000000000..d643f9f576 --- /dev/null +++ b/backend/tests/live/gmail_resume.py @@ -0,0 +1,105 @@ +"""Test-only process loss after an acknowledged Gmail page; never copies provider state.""" + +import json +import os + +from capture_resume import ResumeTarget +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import CaptureCycle +from airweave.models.capture_scan import CaptureScan +from airweave.models.sync import Sync +from airweave.models.sync_cursor import SyncCursor +from airweave.models.sync_job import SyncJob + + +class GmailResumeProbe: + """Interrupt at the next page boundary; verify one further commit after a fresh process.""" + + def __init__( + self, sessions, organization_id, sync_id, job_id, attempt_id, target_path, mode, counters + ): + self.sessions = sessions + self.organization_id, self.sync_id = organization_id, sync_id + self.job_id, self.attempt_id = job_id, attempt_id + self.target_path, self.mode, self.counters = target_path, mode, counters + self.target = ( + ResumeTarget.model_validate_json(target_path.read_text()) if mode == "resume" else None + ) + self.resumed_saved_page = False + self.saved_cursor_expired = False + + async def before_page(self, scope, continuation): + if scope.record_type != "message" or scope.parent is not None: + raise AssertionError("Gmail recovery requires the independent message scope") + async with self.sessions() as db: + row = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == self.organization_id, + CaptureScan.sync_id == self.sync_id, + CaptureScan.record_type == "message", + CaptureScan.container_id.is_(None), + CaptureScan.parent_record_id.is_(None), + ) + ) + cursor = await db.scalar( + select(SyncCursor).where( + SyncCursor.organization_id == self.organization_id, + SyncCursor.sync_id == self.sync_id, + ) + ) + sync, job = await db.get(Sync, self.sync_id), await db.get(SyncJob, self.job_id) + assert row is not None and cursor is not None and sync is not None and job is not None + cycle = CaptureCycle.model_validate(cursor.cursor_data["canonical_cycle"]) + assert sync.organization_id == job.organization_id == self.organization_id + assert job.sync_id == self.sync_id and job.status == "running" + assert sync.writer_job_id == self.job_id and sync.writer_attempt_id == self.attempt_id + assert cycle.phase == "active" and cycle.mode == "full" + assert cycle.starting_checkpoint is not None + assert row.cycle_id == cycle.version.cycle_id and row.phase == "collecting" + assert row.continuation == continuation.value + target = ResumeTarget( + scan_id=row.id, cycle_id=row.cycle_id, sweep_id=row.sweep_id, revision=row.revision + ) + if self.mode == "interrupt": + if not self.counters["records_observed"]: + return + self.target_path.write_text(target.model_dump_json()) + os.chmod(self.target_path, 0o600) + self._stop(False) + assert self.counters["capture_profile_requests"] == 0, ( + "Resume fetched another starting boundary" + ) + if target == self.target: + self.resumed_saved_page = True + elif self.resumed_saved_page and not self.saved_cursor_expired: + assert (target.scan_id, target.cycle_id, target.sweep_id) == ( + self.target.scan_id, + self.target.cycle_id, + self.target.sweep_id, + ) + assert target.revision > self.target.revision + self._stop(True) + + def cursor_expired(self): + self.saved_cursor_expired = True + + def _stop(self, verified): + print( + json.dumps( + { + "failed": False, + "counters_complete": True, + "intentional_interruption": not verified, + "gmail_recovery_verified": verified, + "resumed_saved_page": self.resumed_saved_page, + "saved_cursor_expired": self.saved_cursor_expired, + "full_scope_completed": False, + "checkpoint_saved": False, + "job_status": "running", + **self.counters, + } + ), + flush=True, + ) + os._exit(76 if verified else 75) diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index eb0386a95d..43f15a6b4b 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -29,7 +29,10 @@ from airweave.adapters.event_bus.fake import FakeEventBus from airweave.domains.entities.canonical.page_source import ( CanonicalPageSource, + CheckpointedPageSource, + InvalidCaptureCheckpoint, InvalidScanContinuation, + KnownObjectSource, ) from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope from airweave.domains.entities.canonical.service import CanonicalCaptureService @@ -64,16 +67,19 @@ class BudgetExceeded(RuntimeError): class BoundedStorage(harness.FilesystemBackend): """Cap all writes, even repeated MIME reads, before retaining another blob.""" - def __init__(self, root, byte_limit): + def __init__(self, root, byte_limit, counters=None): super().__init__(root) self.written_bytes = 0 self.byte_limit = byte_limit + self.counters = counters async def write_file(self, path, content): if self.written_bytes + len(content) > self.byte_limit: raise BudgetExceeded("blob_bytes") await super().write_file(path, content) self.written_bytes += len(content) + if self.counters is not None: + self.counters["blob_bytes_written"] = self.written_bytes class BoundedPageSource: @@ -91,11 +97,11 @@ async def capture_page(self, scope, continuation, *, files: FileService, parent= await self.resume_probe.before_page(scope, continuation) try: page = await self.source.capture_page(scope, continuation, files=files, parent=parent) - except InvalidScanContinuation: + except (InvalidScanContinuation, InvalidCaptureCheckpoint): if self.resume_probe is not None: self.resume_probe.cursor_expired() raise - self.counters["records_observed"] += len(page.records) + self.counters["records_observed"] += len(page.records) + len(page.discovered_records) if self.counters["records_observed"] > self.record_limit: raise BudgetExceeded("records") return page @@ -107,6 +113,39 @@ async def confirm_absent(self, record): await self.source.confirm_absent(record) +class BoundedKnownSource(BoundedPageSource): + async def refresh_known(self, record, *, files): + observation = await self.source.refresh_known(record, files=files) + self.counters["records_observed"] += 1 + if self.counters["records_observed"] > self.record_limit: + raise BudgetExceeded("records") + return observation + + +class BoundedCheckpointSource(BoundedPageSource): + async def prepare_cycle(self, previous): + return await self.source.prepare_cycle(previous) + + def initial_continuation(self, cycle): + return self.source.initial_continuation(cycle) + + +class BoundedCheckpointKnownSource(BoundedCheckpointSource, BoundedKnownSource): + pass + + +def bounded_page_source(source, counters, record_limit, probe=None): + """Preserve only protocols the actual source implements.""" + checkpoint = isinstance(source, CheckpointedPageSource) + known = isinstance(source, KnownObjectSource) + wrapper = ( + (BoundedCheckpointKnownSource if known else BoundedCheckpointSource) + if checkpoint + else (BoundedKnownSource if known else BoundedPageSource) + ) + return wrapper(source, counters, record_limit, probe) + + async def bounded_observations(source, cursor, files, counters, record_limit): """A breached test budget fails enumeration before scope completion.""" async with aclosing(source.generate_observations(cursor=cursor, files=files)) as stream: @@ -186,8 +225,16 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert saved["canonical_cycle"]["phase"] == "complete" assert saved["canonical_cycle"]["completed_job_id"] else: - assert saved["canonical_query"] == manifest["query"] - assert saved["history_id"] == "" and counters["started"] == counters["completed"] == 1 + cycle = saved["canonical_cycle"] + assert cycle["phase"] == "complete" and cycle["completed_job_id"] + if manifest.get("gmail_unfiltered"): + assert cycle["promoted_checkpoint"] and cycle["last_full_capture"] + if previous.get("canonical_cycle", {}).get("phase") == "complete": + assert cycle["mode"] == "changes" + assert counters["resumed_changes_requests"] > 0 + assert counters["capture_profile_requests"] == 0 + else: + assert cycle["mode"] == "full" and cycle["promoted_checkpoint"] is None assert saved["canonical_checkpoint"] == { "writer_attempt_id": str(attempt.id), "observed_change_sequence": sequence, @@ -218,9 +265,14 @@ def prepare_resume_probe(manifest, sessions, organization_id, sync_id, job_id, a if manifest.get("resume_stage") not in {"interrupt", "resume"}: return None from capture_resume import PageResumeProbe + from gmail_resume import GmailResumeProbe from wispr_resume import WisprResumeProbe - probe_type = WisprResumeProbe if manifest.get("wispr_resume") else PageResumeProbe + probe_type = ( + GmailResumeProbe + if manifest.get("gmail_resume") + else (WisprResumeProbe if manifest.get("wispr_resume") else PageResumeProbe) + ) return probe_type( sessions, organization_id, @@ -233,6 +285,23 @@ def prepare_resume_probe(manifest, sessions, organization_id, sync_id, job_id, a ) +def count_gmail_request(name, request, previous, counters, identity_verified): + """Count required identity verification separately from production capture planning.""" + if name != "gmail": + return + if request.url.path.endswith("/profile"): + key = "capture_profile_requests" if identity_verified else "identity_profile_requests" + counters[key] += 1 + if request.url.path.endswith("/history"): + boundary = previous.get("canonical_cycle", {}).get("promoted_checkpoint") + if ( + boundary + and request.url.params.get("startHistoryId") + == boundary["checkpoint"]["value"]["history_id"] + ): + counters["resumed_changes_requests"] += 1 + + async def child(manifest): engine = create_async_engine( harness.test_database_url(), @@ -249,6 +318,10 @@ async def child(manifest): is_calendar = name == "google_calendar" counters["sync_token_requests"] = 0 counters["resumed_changes_requests"] = 0 + counters["identity_profile_requests"] = 0 + counters["capture_profile_requests"] = 0 + counters["blob_bytes_written"] = 0 + identity_verified = False from functools import partial from wispr_diagnostics import WisprDiagnostics, observe_request @@ -285,6 +358,7 @@ async def request_hook(request): if counters["provider_requests"] >= manifest["request_limit"]: raise BudgetExceeded("provider_requests") counters["provider_requests"] += 1 + count_gmail_request(name, request, previous, counters, identity_verified) if request.url.params.get("syncToken"): counters["sync_token_requests"] += 1 if ( @@ -337,7 +411,9 @@ async def db_context(): }.get(name) cursor = SyncCursor(sync_id, cursor_schema, previous or None) if cursor_schema else None loaded = cursor.loaded_from_db if cursor else False - storage = BoundedStorage(Path(manifest["root"]) / "blobs", manifest["blob_byte_limit"]) + storage = BoundedStorage( + Path(manifest["root"]) / "blobs", manifest["blob_byte_limit"], counters + ) files = FileService(job_id, storage, sync_id=sync_id) files.MAX_FILE_SIZE_BYTES = manifest["file_byte_limit"] fence = SimpleNamespace(organization_id=organization_id, sync_id=sync_id, job_id=job_id) @@ -366,6 +442,7 @@ async def db_context(): key, fence, gmail_query=manifest.get("query", ""), + gmail_unfiltered=manifest.get("gmail_unfiltered", False), calendar_config=GoogleCalendarConfig.model_validate(manifest["calendar_config"]) if is_calendar else None, @@ -378,6 +455,7 @@ async def db_context(): asyncio.timeout(manifest["timeout"]), connection as (source, identity_verification), ): + identity_verified = True assert source.cursor_class == cursor_schema bus = FakeEventBus() attempt = CaptureAttempt(id=uuid4(), number=attempt_number) @@ -385,7 +463,7 @@ async def db_context(): manifest, sessions, organization_id, sync_id, job_id, attempt, counters ) page_source = ( - BoundedPageSource(source, counters, manifest["record_limit"], resume_probe) + bounded_page_source(source, counters, manifest["record_limit"], resume_probe) if isinstance(source, CanonicalPageSource) else None ) @@ -480,7 +558,8 @@ async def db_context(): ) result = { "failed": False, - "full_scope_completed": name != "wispr", + "full_scope_completed": name != "wispr" + and not (name == "gmail" and saved["canonical_cycle"]["mode"] == "changes"), "counters_complete": True, "resumed_saved_page": resume_probe.resumed_saved_page if resume_probe else False, "saved_cursor_expired": resume_probe.saved_cursor_expired if resume_probe else False, @@ -498,7 +577,9 @@ async def db_context(): "job_status": status, "payload_revision_digest": digest, "mode": { - "gmail": "filtered_full_reconciliation", + "gmail": ("unfiltered_" + saved["canonical_cycle"]["mode"]) + if manifest.get("gmail_unfiltered") + else "filtered_full_reconciliation", "google_calendar": "calendar_incremental" if loaded else "calendar_initial", "google_drive": "drive_incremental" if loaded else "drive_initial", "slack": "accessible_history_full_reconciliation", @@ -578,19 +659,43 @@ async def execute_trial(path, timeout): return process.returncode, lines[-1] +def remaining_blob_budget(manifest, results, total): + """Repeated writes consume the same parent-owned trial allowance.""" + if total is None: + return True + manifest["blob_byte_limit"] = total - sum(item.get("blob_bytes_written", 0) for item in results) + if manifest["blob_byte_limit"] > 0: + return True + results.append( + {"failed": True, "safe_reason": "aggregate_blob_budget", "full_scope_completed": False} + ) + return False + + +def mark_early_completion(gmail_resume, stage, code, result): + if gmail_resume and stage == "interrupt" and code == 0: + result["recovery_not_exercised"] = True + + async def run_trials(manifest, evidence_reader=None): """One parent-owned aggregate budget across interruption, retry, and new-job proof.""" path = Path(manifest["root"]) / "manifest.json" results = [] name = manifest["provider"] wispr_resume = manifest.get("wispr_resume", False) - resume_trial = name == "slack" or wispr_resume + gmail_resume = manifest.get("gmail_resume", False) + resume_trial = name == "slack" or wispr_resume or gmail_resume stages = ( - (("interrupt", "resume") if wispr_resume else ("interrupt", "resume", "new_cycle")) + ( + ("interrupt", "resume") + if wispr_resume or gmail_resume + else ("interrupt", "resume", "new_cycle") + ) if resume_trial else ("first", "second") ) total_requests, total_records = manifest["request_limit"], manifest["record_limit"] + total_blob_bytes = manifest.get("blob_byte_limit") deadline = time.monotonic() + manifest["timeout"] if resume_trial else None interrupted_job_id = str(uuid4()) for stage in stages: @@ -610,6 +715,8 @@ async def run_trials(manifest, evidence_reader=None): } ) break + if not remaining_blob_budget(manifest, results, total_blob_bytes): + break manifest.update( resume_stage=stage, job_id=interrupted_job_id if stage != "new_cycle" else str(uuid4()), @@ -628,16 +735,17 @@ async def run_trials(manifest, evidence_reader=None): result["resumed_partial"] = bool( result.get("resumed_page_committed") and not result.get("full_scope_completed") ) + mark_early_completion(gmail_resume, stage, code, result) results.append(result) print(json.dumps({"run": len(results), "stage": stage, **result}), flush=True) if stage == "interrupt": if code != 75 or not result.get("intentional_interruption"): break elif ( - wispr_resume + (wispr_resume or gmail_resume) and stage == "resume" and code == 76 - and result.get("wispr_recovery_verified") + and (result.get("wispr_recovery_verified") or result.get("gmail_recovery_verified")) ): break elif code: @@ -670,11 +778,22 @@ def wispr_request_limit(value: str) -> int: return limit +def gmail_trial_options(name): + resume = os.environ.get("LIVE_GMAIL_RESUME") == "1" + unfiltered = os.environ.get("LIVE_GMAIL_UNFILTERED") == "1" + if (resume or unfiltered) and name != "gmail": + raise ValueError("Gmail trial options require Gmail provider") + if resume and not unfiltered: + raise ValueError("Gmail recovery proof requires explicit unfiltered scope") + return resume, unfiltered + + async def main(): name = os.environ.get("LIVE_LIFECYCLE_PROVIDER", "gmail") if name not in {"gmail", "google_calendar", "google_drive", "slack", "wispr"}: raise ValueError("Unsupported lifecycle provider") wispr_resume = os.environ.get("LIVE_WISPR_RESUME") == "1" + gmail_resume, gmail_unfiltered = gmail_trial_options(name) if wispr_resume and name != "wispr": raise ValueError("Wispr resume requires Wispr provider") request_limit = ( @@ -755,6 +874,8 @@ async def main(): "sync_id": str(sync_id), "query": f"after:{end - 7 * 86400} before:{end}", } + if name == "gmail": + manifest.update(gmail_resume=gmail_resume, gmail_unfiltered=gmail_unfiltered) if wispr_resume: manifest.update( wispr_resume=True, request_limit=request_limit, record_limit=600, timeout=180 @@ -780,7 +901,9 @@ async def resume_evidence(): root / "resume-target.json", ) - results = await run_trials(manifest, resume_evidence if name == "slack" else None) + results = await run_trials( + manifest, resume_evidence if name == "slack" or gmail_resume else None + ) finally: os.umask(old_umask) @@ -798,6 +921,15 @@ async def resume_evidence(): and results[1].get("full_scope_completed") and results[2].get("full_scope_completed") ) + gmail_verified = bool( + gmail_resume + and len(results) == 2 + and results[0].get("intentional_interruption") + and results[1].get("resumed_page_committed") + and results[1].get("resumed_saved_page") + and results[1].get("capture_profile_requests") == 0 + and not results[1].get("saved_cursor_expired") + ) print( json.dumps( { @@ -806,6 +938,7 @@ async def resume_evidence(): "schema_removed": True, "blob_directory_removed": True, "resume_verified": bool(resume_verified), + "gmail_recovery_verified": gmail_verified, "wispr_recovery_verified": bool( wispr_resume and len(results) == 2 @@ -816,7 +949,8 @@ async def resume_evidence(): ), "resumed_partial": any(item.get("resumed_partial", False) for item in results), "saved_cursor_reused": bool( - resume_verified and not results[1].get("saved_cursor_expired") + (resume_verified or gmail_verified) + and not results[1].get("saved_cursor_expired") ), **trial_counter_summary(results), "identical_second_read": len(completed) == 2 @@ -829,7 +963,9 @@ async def resume_evidence(): wispr_resume and len(results) == 2 and results[1].get("wispr_recovery_verified", False) ) success = ( - wispr_verified + gmail_verified + if gmail_resume + else wispr_verified if wispr_resume else (resume_verified if name == "slack" else len(results) == 2) ) diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index 222ff69cb5..e10db47fdf 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -67,6 +67,7 @@ async def rest_source( gmail_query="newer_than:7d smaller:5M", request_hook=None, response_hook=None, + gmail_unfiltered=False, calendar_config=None, max_file_bytes=10 * 1024 * 1024, ): @@ -82,7 +83,18 @@ async def rest_source( client, fence.organization_id, name, feature_flag_enabled=False ) source_type, config = { - "gmail": (GmailSource, GmailConfig(gmail_query=gmail_query)), + "gmail": ( + GmailSource, + GmailConfig( + gmail_query=None, + included_labels=[], + excluded_labels=[], + excluded_categories=[], + after_date=None, + ) + if gmail_unfiltered + else GmailConfig(gmail_query=gmail_query), + ), "google_drive": (GoogleDriveSource, GoogleDriveConfig()), "google_calendar": (GoogleCalendarSource, calendar_config or GoogleCalendarConfig()), "slack": (SlackSource, SlackConfig()), From 0ac9abe0427d0c229348bcec33d5a42241b3f4be Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:10:16 -0700 Subject: [PATCH 090/266] Restart Slack inventory when a queued thread disappears --- .../canonical/tests/test_slack_recovery.py | 35 ++++++++++++++- backend/airweave/platform/sources/slack.py | 4 ++ .../platform/sources/test_slack_capture.py | 43 +++++++++++++++++++ 3 files changed, 80 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index da78ba6ec4..1e07836f46 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -131,11 +131,14 @@ async def test_expired_reply_cursor_restarts_whole_scope_and_never_partial_recon assert next(scan for scan in scans if scan.record_type == "message").sweep_id != old_sweep -async def test_repeated_invalid_cursor_fails_without_checkpoint_or_absence(database, source): +@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, _, _ = runner( database, source, - [ROOT, HISTORY, SlackApiError("invalid_cursor"), HISTORY, SlackApiError("invalid_cursor")], + [ROOT, HISTORY, SlackApiError(error_code), HISTORY, SlackApiError(error_code)], ) from airweave.domains.entities.canonical.page_source import InvalidScanContinuation @@ -502,3 +505,31 @@ async def change_during_network(scope, continuation, *, files, parent=None): 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, _ = 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) diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index a9238f8376..774442ff91 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -334,6 +334,10 @@ async def _capture_page_response( 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 items = TypeAdapter(list[dict[str, JsonValue]]).validate_python(payload.get(item_key)) metadata = SlackPageMetadata.model_validate(payload.get("response_metadata") or {}) diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 87c8c4a4b0..bb47c9feba 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -131,3 +131,46 @@ async def test_proxy_rate_limit_stops_after_five_attempts(): assert connector.http_client.get.await_count == 5 assert waits.await_count == 4 assert all(float(call.args[0]) == 240 for call in waits.call_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code, expected", + [ + ("thread_not_found", "restart"), + ("channel_not_found", "access"), + ("not_in_channel", "access"), + ("missing_scope", "error"), + ("internal_error", "error"), + ], +) +async def test_queued_thread_failure_preserves_error_meaning(code, expected): + from airweave.domains.entities.canonical.page_source import ( + InvalidScanContinuation, + ScopeAccessLost, + ) + + connector = source() + connector._get = AsyncMock(side_effect=SlackApiError(code)) + error = {"restart": InvalidScanContinuation, "access": ScopeAccessLost, "error": SlackApiError}[ + expected + ] + with pytest.raises(error): + await connector.capture_page( + CompletedScope(record_type="message", container_id="C1"), + ScanContinuation(value={"pending_threads": ["1"]}), + files=MagicMock(), + ) + assert connector._get.await_count == 1 + + +@pytest.mark.asyncio +async def test_thread_not_found_outside_reply_fetch_does_not_invalidate_inventory(): + connector = source() + connector._get = AsyncMock(side_effect=SlackApiError("thread_not_found")) + with pytest.raises(SlackApiError): + await connector.capture_page( + CompletedScope(record_type="message", container_id="C1"), + ScanContinuation(), + files=MagicMock(), + ) From 75d85b53e653add5673f82e10e617864773422c3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:11:25 -0700 Subject: [PATCH 091/266] Capture Gmail through resumable baseline and history pages --- .../canonical/tests/test_capture_pipeline.py | 1 + .../canonical/tests/test_checkpoint_driver.py | 43 ++ .../canonical/tests/test_gmail_recovery.py | 241 ++++++++++ .../sync_pipeline/canonical_capture.py | 10 +- .../domains/sync_pipeline/canonical_scan.py | 90 ++-- backend/airweave/platform/sources/gmail.py | 212 ++++++-- .../platform/sources/gmail_capture.py | 124 ++--- .../airweave/platform/sources/gmail_pages.py | 268 +++++++++++ .../sources/tests/test_gmail_capture.py | 455 ++++++++++++------ .../sources/tests/test_gmail_history_pages.py | 15 +- 10 files changed, 1115 insertions(+), 344 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py create mode 100644 backend/airweave/platform/sources/gmail_pages.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py b/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py index a58672771d..e4e87c41a3 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py +++ b/backend/airweave/domains/entities/canonical/tests/test_capture_pipeline.py @@ -45,6 +45,7 @@ def components(database, source): source_connection_id=uuid4(), source_short_name="test", execution_config=config, + force_full_sync=False, logger=logger, batch_size=1, max_batch_latency_ms=0, diff --git a/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py b/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py index 992b1806a9..2a8a0badd7 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py +++ b/backend/airweave/domains/entities/canonical/tests/test_checkpoint_driver.py @@ -251,3 +251,46 @@ async def test_exhaustive_known_validation_is_enforced_by_sql(database, source): ), ) assert result.state.phase == "complete" and result.capture.changes == () + + +@pytest.mark.parametrize("force_flag", ["force_full_sync", "skip_load"]) +async def test_explicit_full_sync_restarts_delta_but_completed_job_retry_is_idempotent( + database, source, force_flag +): + from airweave.domains.entities.canonical.cycle_models import BeginCycle + + connector = NativeFixture() + ctx, runtime, capture = pipeline(database, source, connector) + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + service, fence = source + async with database() as db: + before = await service.read_cycle(db, fence) + delta = await service.begin_cycle( + db, + BeginCycle( + fence=fence, + configuration=connector.capture_cycle_configuration, + expected=before.version, + mode="changes", + ), + ) + ctx, runtime, capture = pipeline(database, source, connector, attempt=2) + ctx.force_full_sync = force_flag == "force_full_sync" + ctx.execution_config.cursor.skip_load = force_flag == "skip_load" + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + async with database() as db: + completed = await service.read_cycle(db, capture._writer()) + assert completed.mode == "full" and completed.version.cycle_id != delta.version.cycle_id + assert completed.starting_checkpoint.value == {"start": "2"} + pages = len(connector.pages) + ctx, runtime, capture = pipeline(database, source, connector, attempt=3) + ctx.force_full_sync = force_flag == "force_full_sync" + ctx.execution_config.cursor.skip_load = force_flag == "skip_load" + await capture.start(ctx) + await capture.run_scans(ctx, runtime, AsyncMock()) + await capture.save_checkpoint(ctx, runtime) + assert len(connector.pages) == pages and len(connector.plans) == 2 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..f25091f91f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py @@ -0,0 +1,241 @@ +"""Actual Gmail transport/page source, lifecycle and PostgreSQL; native HTTP is synthetic.""" + +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): + 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( + 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 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 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 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 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 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 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 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 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" diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index 7d66ce702d..e052758303 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -15,7 +15,10 @@ TerminalCheckpoint, ) from airweave.domains.entities.canonical.models import CaptureResult -from airweave.domains.entities.canonical.page_source import CanonicalPageSource +from airweave.domains.entities.canonical.page_source import ( + CanonicalPageSource, + CheckpointedPageSource, +) from airweave.domains.entities.canonical.requests import ( CaptureBatch, CaptureRecord, @@ -132,6 +135,11 @@ async def progress(result: CaptureResult, records: tuple[CaptureRecord, ...]) -> 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) + else False, ).run() def _writer(self) -> WriterFence: diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 5dcf1c4047..9445873584 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -49,10 +49,13 @@ def __init__( 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 async def run(self) -> CaptureCycle: @@ -66,26 +69,35 @@ async def run(self) -> CaptureCycle: expected = current.version else: expected = None - plan = ( - CapturePlan(mode=current.mode, starting_checkpoint=current.starting_checkpoint) - if current is not None and current.phase == "active" - else ( - await self.source.prepare_cycle(current) - if isinstance(self.source, CheckpointedPageSource) - 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, - ), + if ( + self.force_full + and isinstance(self.source, CheckpointedPageSource) + and current is not None + and current.phase == "active" + and (current.mode == "changes" or current.configuration != configuration) + ): + cycle = await self.fresh_full(current) + else: + plan = ( + CapturePlan(mode=current.mode, starting_checkpoint=current.starting_checkpoint) + 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) + 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, + ), + ) for restart in range(2): try: return await self.run_cycle(cycle) @@ -93,25 +105,31 @@ async def run(self) -> CaptureCycle: if restart or not isinstance(self.source, CheckpointedPageSource): raise await self.check_limits() - plan = await self.source.prepare_cycle(None) - if plan.mode != "full": - raise CycleConflict("An invalid checkpoint 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 != cycle.version.cycle_id: - raise CycleConflict("Cycle changed during checkpoint recovery") - cycle = await self.service.restart_cycle( - db, - RestartCycle( - fence=self.fence, - expected=current.version, - configuration=configuration, - mode=plan.mode, - starting_checkpoint=plan.starting_checkpoint, - ), - ) + 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): + raise CycleConflict("Source has no checkpoint recovery contract") + plan = await self.source.prepare_cycle(None) + if plan.mode != "full": + 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, + ), + ) + async def run_cycle(self, cycle: CaptureCycle) -> CaptureCycle: """Process the current frontier; every acknowledgement remains in the existing SQL store.""" configuration = cycle.configuration diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index aec5ec7e48..426d6812f6 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -1,16 +1,14 @@ -"""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 contextlib import aclosing from datetime import datetime @@ -22,8 +20,21 @@ 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.source import SourceObservation -from airweave.domains.sources.exceptions import SourceAuthError +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, @@ -47,6 +58,7 @@ ) from airweave.platform.sources._base import BaseSource from airweave.platform.sources.gmail_capture import GmailCapture +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 @@ -60,7 +72,10 @@ def _should_retry_gmail_request(exception: Exception) -> bool: if exception.response.status_code == 429: return True return True - if isinstance(exception, (httpx.ConnectTimeout, httpx.ReadTimeout)): + if isinstance( + exception, + (httpx.ConnectTimeout, httpx.ReadTimeout, SourceRateLimitError, SourceServerError), + ): return True return False @@ -92,21 +107,77 @@ class GmailSource(BaseSource): canonical_record_types = ("message",) - async def generate_observations( + canonical_container_parents = {} + capture_config: GmailConfig + + @property + def capture_cycle_configuration(self) -> CycleConfiguration: + """Bind the actual query and typed filters; preserve every existing filtered scope.""" + query = self._build_gmail_query() + fingerprint = hashlib.sha256( + json.dumps( + { + "capture_version": 1, + "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: + 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.""" + return GmailPages.initial(cycle) + + async def capture_page( self, + scope: CompletedScope, + continuation: ScanContinuation, *, - cursor: SyncCursor | None = None, - files: FileService | None = None, - node_selections: list[NodeSelectionData] | None = None, - ) -> AsyncGenerator[SourceObservation, None]: - """Capture native messages once; thread/search projections are downstream.""" - if node_selections: - raise ValueError("Gmail capture does not support selected node scopes") - capture = GmailCapture( - self._get, self._build_gmail_query(), files=files, attachment_get=self._get_mime_json - ) - async for observation in capture.generate(cursor): - yield observation + 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 @@ -123,6 +194,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)) @@ -282,9 +354,7 @@ 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( - "Gmail rate limit reached" - ) + self.logger.warning("Gmail rate limit reached") raise_for_status( response, @@ -296,26 +366,80 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: self.logger.debug(f"Response data keys: {list(data.keys())}") return data - async def _get_mime_json(self, url: str, *, max_bytes: int) -> dict: + @retry( + stop=stop_after_attempt(5), + retry=_should_retry_gmail_request, + wait=wait_rate_limit_with_backoff, + 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() - async with self.http_client.stream("GET", url, headers=headers) as response: - raise_for_status( - response, source_short_name=self.short_name, - token_provider_kind=self.auth.provider_kind, - ) - body = bytearray() - async with aclosing(response.aiter_bytes()) as chunks: - async for chunk in chunks: - if len(body) + len(chunk) > max_bytes: - raise FileSkippedException( - "Encoded MIME body exceeds size limit", "MIME body" - ) - body.extend(chunk) - value = json.loads(body) - if not isinstance(value, dict): - raise ValueError("Gmail MIME response must be an object") - return value + 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_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.""" + if response.is_stream_consumed: + if len(response.content) > maximum: + raise FileSkippedException( + "Gmail JSON response exceeds size limit", "Gmail response" + ) + 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( + "Gmail JSON response exceeds size limit", "Gmail response" + ) + body.extend(chunk) + return bytes(body) # ----------------------- # Cursor helper diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py index e6af2f43d5..322d70ac89 100644 --- a/backend/airweave/platform/sources/gmail_capture.py +++ b/backend/airweave/platform/sources/gmail_capture.py @@ -3,7 +3,7 @@ import base64 import hashlib import json -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator +from collections.abc import Awaitable, Callable, Iterator from datetime import datetime, timezone import httpx @@ -13,15 +13,11 @@ from airweave.domains.entities.canonical.requests import ( BlobReference, CaptureRecord, - CompletedScope, RecordIdentity, - StartedScope, ) -from airweave.domains.entities.canonical.source import SourceObservation -from airweave.domains.sources.exceptions import SourceEntityNotFoundError +from airweave.domains.sources.exceptions import SourceEntityNotFoundError, SourceError from airweave.domains.storage.exceptions import FileSkippedException from airweave.domains.storage.file_service import FileService -from airweave.domains.syncs.cursors.cursor import SyncCursor BASE = "https://gmail.googleapis.com/gmail/v1/users/me" Get = Callable[..., Awaitable[dict]] @@ -63,7 +59,7 @@ class HistoryPage(BaseModel): fingerprint: str -class HistoryBatch(BaseModel): +class HydratedMessages(BaseModel): """Bounded current observations; callers atomically commit records and offset.""" model_config = ConfigDict(extra="forbid", frozen=True) @@ -231,7 +227,7 @@ async def hydrate_history_page( offset: int = 0, limit: int = 500, expected_fingerprint: str | None = None, - ) -> HistoryBatch: + ) -> 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") @@ -239,90 +235,28 @@ async def hydrate_history_page( 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") - end = min(offset + limit, len(page.message_ids)) - records = tuple([await self.message(mid) for mid in page.message_ids[offset:end]]) - return HistoryBatch(records=records, next_offset=end, complete=end == len(page.message_ids)) - - async def history(self, boundary: str) -> tuple[list[str], str]: - """Legacy generator bridge, removed when the durable page adapter is wired.""" - affected: dict[str, None] = {} - tokens: set[str] = set() - token = None - while True: - page = await self.history_page(boundary, token) - affected.update(dict.fromkeys(page.message_ids)) - token = page.next_page_token - if token is None: - return list(affected), page.history_id - if token in tokens: - raise ValueError("Gmail history pagination repeated a token") - tokens.add(token) - - async def enumerate(self) -> AsyncGenerator[CaptureRecord, None]: - """List every page in the provider's exact configured query scope.""" - params = {"maxResults": 500, "includeSpamTrash": "true"} - if self.query: - params["q"] = self.query - tokens: set[str] = set() - while True: - page = await self.get(f"{BASE}/messages", params=dict(params)) - for item in page.get("messages", []): - yield await self.message(item["id"]) - token = page.get("nextPageToken") - if not token: - return - if token in tokens: - raise ValueError("Gmail messages pagination repeated a token") - tokens.add(token) - params["pageToken"] = token - - async def saved_history(self, cursor: SyncCursor | None) -> tuple[list[str], str] | None: - """Only canonical unfiltered checkpoints can resume this record collection.""" - saved = cursor.data if cursor else {} - boundary = saved.get("history_id") if saved.get("canonical_query") == "" else None - if not boundary: - return None - try: - return await self.history(boundary) - except (httpx.HTTPStatusError, SourceEntityNotFoundError) as exc: - if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code != 404: - raise - return None - - async def generate(self, cursor: SyncCursor | None) -> AsyncGenerator[SourceObservation, None]: - """Filtered queries reconcile each run; unfiltered mailboxes use Gmail history. - - Gmail queries cannot be faithfully evaluated against history JSON. Full query - enumeration is a best-effort changing mailbox view, not a snapshot. Never - advance an incremental cursor for such a view. External MIME bytes require - successful immutable storage; missing storage or oversize bodies remain partial. - """ - if self.query: - yield StartedScope(record_type="message") - async for record in self.enumerate(): - yield record - yield CompletedScope(record_type="message") - if cursor: - cursor.update(history_id="", canonical_query=self.query) - return - - history = await self.saved_history(cursor) - if history is None: - profile = await self.get(f"{BASE}/profile") - boundary = str(profile["historyId"]) - yield StartedScope(record_type="message") - async for record in self.enumerate(): - yield record - # If this boundary expires mid-crawl, fail the run. Starting another full - # crawl in this attempt would incorrectly retain earlier seen records. - history = await self.history(boundary) - full = True - else: - full = False - message_ids, latest = history - for message_id in message_ids: - yield await self.message(message_id) - if full: - yield CompletedScope(record_type="message") - if cursor: - cursor.update(history_id=latest, canonical_query="") + 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_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/tests/test_gmail_capture.py b/backend/airweave/platform/sources/tests/test_gmail_capture.py index 967b7f14f9..62ada414ef 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_capture.py +++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py @@ -1,14 +1,21 @@ -"""Gmail capture races and failures against scripted provider responses.""" +"""Gmail native pages and MIME retention; SQL recovery is tested separately.""" from uuid import uuid4 import httpx import pytest -from airweave.domains.entities.canonical.requests import CompletedScope, StartedScope -from airweave.domains.syncs.cursors.cursor import SyncCursor -from airweave.platform.cursors.gmail import GmailCursor +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): @@ -34,195 +41,126 @@ async def get(self, url, params=None): return response -def cursor(**data): - return SyncCursor(uuid4(), cursor_schema=GmailCursor, cursor_data=data) +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_bootstrap_boundary_precedes_pagination_and_replays_concurrent_edits(): +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", labelIds=["INBOX"])), - ("/messages", {"messages": [{"id": "b"}]}), - ("/messages/b", message("b")), + ("/messages/a", message("a")), ( "/history", - { - "history": [{"labelsRemoved": [{"message": {"id": "a"}}]}], - "nextPageToken": "h2", - "historyId": "11", - }, + {"historyId": "11", "history": [{"labelsRemoved": [{"message": {"id": "a"}}]}]}, ), + ("/messages/a", message("a", labelIds=[])), + ("/messages", {"messages": [{"id": "b"}]}), + ("/messages/b", message("b")), ( "/history", - {"history": [{"messagesAdded": [{"message": {"id": "c"}}]}], "historyId": "12"}, + {"historyId": "12", "history": [{"messagesAdded": [{"message": {"id": "c"}}]}]}, ), - ("/messages/a", message("a", labelIds=[])), ("/messages/c", message("c")), ] ) - state = cursor() - results = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert results[0] == StartedScope(record_type="message") - assert [r.payload["id"] for r in results[1:-1]] == ["a", "b", "a", "c"] - assert results[3].payload["labelIds"] == [] - assert isinstance(results[-1], CompletedScope) - assert state.data["history_id"] == "12" - assert provider.calls[3][1]["pageToken"] == "p2" - assert provider.calls[6][1] == {"startHistoryId": "10", "maxResults": 500, "pageToken": "h2"} - assert not provider.steps + 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_incremental_labels_and_deleted_messages_fetch_current_state(): - provider = Provider( - [ - ( - "/history", - { - "historyId": "21", - "history": [ - { - "messagesDeleted": [{"message": {"id": "gone"}}], - "labelsAdded": [{"message": {"id": "a"}}], - "messages": [{"id": "a"}], - } - ], - }, - ), - ("/messages/a", message("a", labelIds=["STARRED"])), - ("/messages/gone", missing()), - ] - ) - state = cursor(history_id="20", canonical_query="") - results = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert len(results) == 2 - assert results[1].kind == "delete" - assert results[1].removal_reason == "provider_deleted" - assert state.data["history_id"] == "21" +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_incremental_boundary_bootstraps_and_reconciles(): +async def test_expired_history_requires_whole_cycle_recovery(): provider = Provider( - [ - ("/history", missing()), - ("/profile", {"historyId": "30"}), - ("/messages", {}), - ("/history", {"historyId": "31"}), - ] + [("/profile", {"historyId": "10"}), ("/messages", {}), ("/history", missing())] ) - state = cursor(history_id="20", canonical_query="") - results = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert results == [StartedScope(record_type="message"), CompletedScope(record_type="message")] - assert state.data["history_id"] == "31" + 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 -@pytest.mark.parametrize("error", [RuntimeError("network"), missing()]) -async def test_failure_after_full_crawl_does_not_finish_scope_or_advance_checkpoint(error): - provider = Provider( - [ - ("/profile", {"historyId": "30"}), - ("/messages", {"messages": [{"id": "a"}]}), - ("/messages/a", message("a")), - ("/history", error), - ] - ) - state = cursor() - before = state.data - observed = [] - with pytest.raises(type(error)): - async for record in GmailCapture(provider.get, None).generate(state): - observed.append(record) - assert len(observed) == 2 - assert observed[0] == StartedScope(record_type="message") - assert state.data == before +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_filtered_query_always_reconciles_and_external_mime_bytes_are_partial(): - raw = message("a", payload={"parts": [{"body": {"attachmentId": "external", "size": 3}}]}) +async def test_history_pagination_keeps_start_boundary_until_terminal(): provider = Provider( [ - ("/messages", {"messages": [{"id": "a"}]}), - ("/messages/a", raw), - ] - ) - state = cursor(history_id="20", canonical_query="") - results = [ - r async for r in GmailCapture(provider.get, "from:person@example.com").generate(state) - ] - assert results[0] == StartedScope(record_type="message") - assert results[1].payload == raw - assert results[1].completeness == "partial" - assert results[1].blobs == () - assert results[-1] == CompletedScope(record_type="message") - assert provider.calls[0][1]["q"] == "from:person@example.com" - assert state.data["history_id"] == "" - - -@pytest.mark.asyncio -async def test_legacy_entity_cursor_cannot_skip_canonical_bootstrap(): - provider = Provider( - [ - ("/profile", {"historyId": "30"}), + ("/profile", {"historyId": "10"}), ("/messages", {}), - ("/history", {"historyId": "31"}), + ("/history", {"historyId": "999", "nextPageToken": "next"}), + ("/history", {"historyId": "1000"}), ] ) - results = [r async for r in GmailCapture(provider.get, None).generate(cursor(history_id="20"))] - assert results == [StartedScope(record_type="message"), CompletedScope(record_type="message")] + 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_page_token_fails_without_advancing_cursor(): +async def test_repeated_listing_cursor_fails_without_a_completed_page(): provider = Provider( - [ - ("/history", {"nextPageToken": "same", "historyId": "21"}), - ("/history", {"nextPageToken": "same", "historyId": "21"}), - ] + [("/messages", {"nextPageToken": "same"}), ("/messages", {"nextPageToken": "same"})] ) - state = cursor(history_id="20", canonical_query="") + pages = GmailPages(GmailCapture(provider.get, "in:inbox")) + first = await pages.page(await initial(pages)) with pytest.raises(ValueError, match="repeated"): - _ = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert state.data["history_id"] == "20" - - -@pytest.mark.asyncio -async def test_incremental_detail_failure_retains_old_boundary(): - provider = Provider( - [ - ("/history", {"historyId": "21", "history": [{"messages": [{"id": "a"}]}]}), - ("/messages/a", RuntimeError("upstream unavailable")), - ] - ) - state = cursor(history_id="20", canonical_query="") - with pytest.raises(RuntimeError): - _ = [r async for r in GmailCapture(provider.get, None).generate(state)] - assert state.data["history_id"] == "20" - - -@pytest.mark.asyncio -async def test_filtered_page_failure_never_claims_completed_scope(): - provider = Provider( - [ - ("/messages", {"messages": [{"id": "a"}], "nextPageToken": "p2"}), - ("/messages/a", message("a")), - ("/messages", RuntimeError("upstream unavailable")), - ] - ) - state = cursor(history_id="20", canonical_query="") - results = [] - with pytest.raises(RuntimeError): - async for record in GmailCapture(provider.get, "in:inbox").generate(state): - results.append(record) - assert len(results) == 2 - assert results[0] == StartedScope(record_type="message") - assert not any(type(r) is CompletedScope for r in results) - assert state.data["history_id"] == "20" + await pages.page(first.continuation) + assert not first.final @pytest.mark.asyncio @@ -274,12 +212,11 @@ async def test_blob_storage_failure_propagates_and_prevents_record_checkpoint(): ("/messages/a", message("a", payload={"body": {"attachmentId": "body-id", "size": 3}})), ] ) - state = cursor(history_id="20", canonical_query="") 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): - _ = [r async for r in capture.generate(state)] - assert state.data["history_id"] == "20" + await capture.hydrate_history_page(page) @pytest.mark.asyncio @@ -373,3 +310,211 @@ async def test_mime_json_stream_size_limit_closes_response(): 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_history_pages.py b/backend/airweave/platform/sources/tests/test_gmail_history_pages.py index 15e9218144..64619734cb 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_history_pages.py +++ b/backend/airweave/platform/sources/tests/test_gmail_history_pages.py @@ -7,7 +7,7 @@ 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, cursor, message +from airweave.platform.sources.tests.test_gmail_capture import Provider, message def history(ids, **extra): @@ -117,18 +117,7 @@ async def test_non_missing_errors_propagate_from_message_and_history(status): with pytest.raises(type(error)): await capture.hydrate_history_page(parse_history_page(history(["a"]))) with pytest.raises(type(error)): - await capture.saved_history(cursor(history_id="opaque", canonical_query="")) - - -@pytest.mark.asyncio -async def test_legacy_stale_history_handles_actual_translated_404(): - provider = Provider([("/history", translated_error(404))]) - assert ( - await GmailCapture(provider.get, None).saved_history( - cursor(history_id="opaque", canonical_query="") - ) - is None - ) + await capture.history_page("opaque") @pytest.mark.asyncio From 6c45db8605f98a86232f5b01edfa87b48e9ddade Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:13:07 -0700 Subject: [PATCH 092/266] Supply explicit sync mode in live lifecycle context --- backend/tests/live/provider_lifecycle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 43f15a6b4b..89e5a23438 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -395,6 +395,7 @@ async def db_context(): source_short_name=name, connection=SimpleNamespace(id=uuid4(), short_name=name), execution_config=config, + force_full_sync=False, batch_size=10, max_batch_latency_ms=20, should_batch=True, From 965f6e6fb0a8f4e3b6cfe0a96f4fbbafffbc88f8 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:14:50 -0700 Subject: [PATCH 093/266] Exercise actual Gmail lifecycle child with synthetic provider --- .../tests/test_gmail_resume_probe.py | 86 +++++++++++++++++++ 1 file changed, 86 insertions(+) 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 index 5d04eca2e7..1f368f041f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py @@ -169,3 +169,89 @@ async def execute(path, timeout): 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 provider_sample import verify_rest_identity +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', {'emailAddress': expected_email, 'historyId': 'identity-only'}), + ('/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( + included_labels=[], excluded_labels=[], excluded_categories=[])) + await verify_rest_identity(name, connector, expected_email) + yield connector, 'provider_email' + +lifecycle.rest_source = synthetic_source +# 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 From cd423b87e69be0c7eb283855795d28e336c95d14 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:16:04 -0700 Subject: [PATCH 094/266] Record bounded live Gmail recovery and cleanup evidence --- .../evidence/gmail-recovery-20260930.json | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 backend/tests/live/evidence/gmail-recovery-20260930.json diff --git a/backend/tests/live/evidence/gmail-recovery-20260930.json b/backend/tests/live/evidence/gmail-recovery-20260930.json new file mode 100644 index 0000000000..87ddcd31d3 --- /dev/null +++ b/backend/tests/live/evidence/gmail-recovery-20260930.json @@ -0,0 +1,68 @@ +{ + "verified_at": "2026-09-30", + "source_commit": "6c45db8", + "provider": "gmail", + "mode": "bounded_unfiltered_same_job_recovery", + "provider_writes": false, + "limits": { + "requests": 600, + "observations": 250, + "seconds": 600, + "blob_bytes": 268435456 + }, + "initial_attempt": { + "source_commit": "75d85b5", + "harness_exit_code": 1, + "error_type": "AttributeError", + "identified_harness_gap": "missing force_full_sync context field", + "provider_requests": 1, + "records_observed": 0, + "schema_removed": true, + "blob_directory_removed": true + }, + "stages": [ + { + "stage": "interrupt", + "intentional_exit": 75, + "provider_requests": 68, + "records_observed": 50, + "identity_profile_requests": 1, + "capture_profile_requests": 1, + "blob_bytes_written": 3039870, + "full_scope_completed": false, + "checkpoint_saved": false + }, + { + "stage": "resume", + "intentional_exit": 76, + "provider_requests": 2, + "records_observed": 0, + "identity_profile_requests": 1, + "capture_profile_requests": 0, + "blob_bytes_written": 0, + "resumed_saved_page": true, + "saved_cursor_expired": false, + "resumed_page_committed": true, + "full_scope_completed": false, + "checkpoint_saved": false + } + ], + "successful_trial": { + "harness_exit_code": 0, + "provider_requests": 70, + "records_observed": 50, + "gmail_recovery_verified": true, + "schema_removed": true, + "blob_directory_removed": true + }, + "independent_cleanup": { + "remaining_private_schemas": 0, + "remaining_private_live_directories": 0, + "immutable_archive_removed": true + }, + "limitations": [ + "Resumed commit was an empty history-check page; no additional messages were captured by the resumed process.", + "No completed full-mailbox baseline, subsequent change cycle, deletion or hosted owner-binding proof.", + "This trial uses local disposable PostgreSQL and fixture collaborators for non-capture services; it is not a staging deployment." + ] +} From 912ee37a714d2c8d2ea2a495c9e042c20f42228b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:24:59 -0700 Subject: [PATCH 095/266] Record failed bounded Gmail query trial without completion claim --- .../live/evidence/gmail-query-20260930.json | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 backend/tests/live/evidence/gmail-query-20260930.json diff --git a/backend/tests/live/evidence/gmail-query-20260930.json b/backend/tests/live/evidence/gmail-query-20260930.json new file mode 100644 index 0000000000..c8d101d9ac --- /dev/null +++ b/backend/tests/live/evidence/gmail-query-20260930.json @@ -0,0 +1,26 @@ +{ + "verified_at": "2026-09-30", + "source_commit": "cd423b8", + "provider": "gmail", + "mode": "frozen_seven_day_query_two_cycle_attempt", + "provider_writes": false, + "completed_runs": 0, + "attempted_runs": 1, + "harness_exit_code": 1, + "error_type": "SourceServerError", + "last_operation": "metadata", + "provider_requests": 184, + "records_observed": 100, + "blob_bytes_written": 3482736, + "full_scope_completed": false, + "schema_removed": true, + "blob_directory_removed": true, + "independent_remaining_private_schemas": 0, + "independent_remaining_private_live_directories": 0, + "immutable_archive_removed": true, + "limitations": [ + "Second cycle did not run.", + "No complete query scope or subsequent cycle proof.", + "Error type identifies an upstream server failure; originating provider versus authentication proxy was not established by sanitized output." + ] +} From d9f9629e785e304885a569ca35f96182fa4feacd Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:32:12 -0700 Subject: [PATCH 096/266] Restore Gmail request retries and verify failed-job continuation --- .../canonical/tests/test_gmail_recovery.py | 72 ++++++++++++++++++ backend/airweave/platform/sources/gmail.py | 6 +- .../sources/tests/test_gmail_retry.py | 75 +++++++++++++++++++ 3 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 backend/airweave/platform/sources/tests/test_gmail_retry.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py index f25091f91f..c78695840b 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py @@ -239,3 +239,75 @@ async def test_history_offset_resume_and_lost_final_ack_do_not_rehydrate_committ 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 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 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"} diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index 426d6812f6..830d623ca0 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -15,7 +15,7 @@ from typing import Any, AsyncGenerator, Dict, List, Optional, Set import httpx -from tenacity import retry, stop_after_attempt +from tenacity import retry, retry_if_exception, stop_after_attempt from airweave.core.logging import ContextualLogger from airweave.core.shared_models import RateLimitLevel @@ -335,7 +335,7 @@ async def _refresh_and_get_headers(self) -> Dict[str, str]: @retry( stop=stop_after_attempt(5), - retry=_should_retry_gmail_request, + retry=retry_if_exception(_should_retry_gmail_request), wait=wait_rate_limit_with_backoff, reraise=True, ) @@ -368,7 +368,7 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: @retry( stop=stop_after_attempt(5), - retry=_should_retry_gmail_request, + retry=retry_if_exception(_should_retry_gmail_request), wait=wait_rate_limit_with_backoff, reraise=True, ) 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..5ea0b45927 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_gmail_retry.py @@ -0,0 +1,75 @@ +"""Exercise the decorated request boundary, not just its exception predicate.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from tenacity import wait_none + +from airweave.domains.sources.exceptions import SourceAuthError, 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"]) +async def test_transient_server_error_retries_actual_request(method): + # For _get, return a real response; canonical JSON returns an already parsed dict. + import httpx + + 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, [SourceServerError(status_code=503), 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"]) +async def test_server_error_exhausts_exactly_five_attempts(method): + error = SourceServerError(status_code=503) + source, upstream, call = request(method, error) + with pytest.raises(SourceServerError) 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, + ), + 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 From 1b28bf75481751e078409adee0f6934110328cab Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:34:18 -0700 Subject: [PATCH 097/266] Retry bounded Gmail proxy transients without early rate-limit retries --- backend/airweave/platform/sources/gmail.py | 24 +++++-- .../sources/tests/test_gmail_retry.py | 68 ++++++++++++++++--- 2 files changed, 79 insertions(+), 13 deletions(-) diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index 830d623ca0..c9777a6da9 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -19,6 +19,10 @@ 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.entities.canonical.cycle_models import CaptureCycle, CycleConfiguration from airweave.domains.entities.canonical.models import SourceRecord @@ -66,15 +70,25 @@ def _should_retry_gmail_request(exception: Exception) -> bool: """Custom retry condition that excludes 404 errors but includes 429 and timeouts.""" + 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 + 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, SourceRateLimitError, SourceServerError), + ( + httpx.ConnectTimeout, + httpx.ReadTimeout, + SourceServerError, + AuthProviderServerError, + ), ): return True return False diff --git a/backend/airweave/platform/sources/tests/test_gmail_retry.py b/backend/airweave/platform/sources/tests/test_gmail_retry.py index 5ea0b45927..895ad84f05 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_retry.py +++ b/backend/airweave/platform/sources/tests/test_gmail_retry.py @@ -4,10 +4,20 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, Mock +import httpx import pytest from tenacity import wait_none -from airweave.domains.sources.exceptions import SourceAuthError, SourceServerError +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 @@ -26,16 +36,22 @@ def request(method, effects): @pytest.mark.asyncio @pytest.mark.parametrize("method", ["_get", "_get_capture_json"]) -async def test_transient_server_error_retries_actual_request(method): +@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. - import httpx - 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, [SourceServerError(status_code=503), success]) + 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} @@ -44,10 +60,17 @@ async def test_transient_server_error_retries_actual_request(method): @pytest.mark.asyncio @pytest.mark.parametrize("method", ["_get", "_get_capture_json"]) -async def test_server_error_exhausts_exactly_five_attempts(method): - error = SourceServerError(status_code=503) +@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(SourceServerError) as caught: + with pytest.raises(type(error)) as caught: await call(source, "https://example.test") assert caught.value is error assert upstream.await_count == 5 @@ -64,6 +87,19 @@ async def test_server_error_exhausts_exactly_five_attempts(method): 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(), ], @@ -73,3 +109,19 @@ async def test_auth_budget_and_cancellation_do_not_retry(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 From 06d9ec6646fe71d1edc03c00d93174b06a0f26e4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:35:20 -0700 Subject: [PATCH 098/266] Capture Drive files with durable full and changes pages --- .../canonical/tests/test_drive_recovery.py | 256 +++++++++++++++++ .../airweave/platform/sources/google_drive.py | 103 +++++-- .../platform/sources/records/google_drive.py | 138 +-------- .../sources/records/google_drive_pages.py | 261 ++++++++++++++++++ .../sources/records/test_google_drive.py | 175 ++++++------ 5 files changed, 703 insertions(+), 230 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py create mode 100644 backend/airweave/platform/sources/records/google_drive_pages.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py new file mode 100644 index 0000000000..29dedd19d6 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py @@ -0,0 +1,256 @@ +"""Real canonical SQL/pipeline and Drive source; native HTTP is synthetic.""" + +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +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.models.entity import Entity +from airweave.models.sync_job import SyncJob +from airweave.platform.configs.config import GoogleDriveConfig +from airweave.platform.sources.google_drive import GoogleDriveSource +from airweave.platform.sources.records.google_drive import BASE + + +class NativeHTTP: + def __init__(self, replies): + self.replies = list(replies) + self.calls = [] + + async def handle(self, request): + self.calls.append((request.url.path, dict(request.url.params))) + path, value = self.replies.pop(0) + assert str(request.url).split("?")[0] == BASE + path + if isinstance(value, Exception): + raise value + status, payload = value if isinstance(value, tuple) else (200, value) + return httpx.Response(status, json=payload) + + +def folder(native_id): + return {"id": native_id, "mimeType": "application/vnd.google-apps.folder"} + + +async def setup(database, source, native, attempt=1): + service, fence = source + client = httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) + connector = await GoogleDriveSource.create( + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=client, + config=GoogleDriveConfig(), + ) + 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_resume_pending_file_and_changes_then_lost_final_ack(database, source): + native = NativeHTTP( + [ + ("/changes/startPageToken", {"startPageToken": "before"}), + ("/files", {"files": [{"id": "a"}, {"id": "b"}]}), + ("/files/a", folder("a")), + ("/files/b", ConnectionError("interrupted")), + ("/files/b", folder("b")), + ( + "/changes", + { + "changes": [{"changeType": "file", "fileId": "a", "removed": True}], + "nextPageToken": "next", + }, + ), + ("/files/a", (404, {"error": {"message": "gone"}})), + ("/changes", {"changes": [], "newStartPageToken": "after"}), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with client: + with pytest.raises(ConnectionError): + await run(pipeline, ctx, runtime) + async with database() as db: + cycle = await source[0].read_cycle(db, pipeline._writer()) + assert cycle.promoted_checkpoint is None + assert len((await db.scalars(select(Entity))).all()) == 1 + pipeline, ctx, runtime, client = await setup(database, source, native, 2) + async with client: + await pipeline.run_scans(ctx, runtime, no_limits) + # Crash after final durable page, before publication. No provider call on retry. + pipeline, ctx, runtime, client = await setup(database, source, native, 3) + async with client: + await run(pipeline, ctx, runtime) + assert not native.replies + assert sum(path.endswith("/files") for path, _ in native.calls) == 1 + assert sum(path.endswith("/startPageToken") for path, _ in native.calls) == 1 + async with database() as db: + cycle = await source[0].read_cycle(db, pipeline._writer()) + assert cycle.promoted_checkpoint.checkpoint.value == {"page_token": "after"} + records = {r.native_id: r for r in (await db.scalars(select(Entity))).all()} + assert records["a"].removal_reason == "scope_removed" + assert records["b"].deleted_at is None + + full_evidence = cycle.last_full_capture + service, fence = source + async with database() as db: + 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( + [ + ( + "/changes", + { + "changes": [{"changeType": "file", "fileId": "a", "removed": True}], + "newStartPageToken": "later", + }, + ), + ("/files/a", folder("a")), + ] + ) + pipeline, ctx, runtime, client = await setup(database, (service, next_fence), delta) + async with client: + await run(pipeline, ctx, runtime) + assert delta.calls[0][1]["pageToken"] == "after" + assert not delta.replies + async with database() as db: + cycle = await service.read_cycle(db, pipeline._writer()) + assert cycle.mode == "changes" and cycle.last_full_capture == full_evidence + assert cycle.promoted_checkpoint.checkpoint.value == {"page_token": "later"} + assert all(row.deleted_at is None for row in (await db.scalars(select(Entity))).all()) + + +async def test_drive_membership_restart_preserves_records_until_exact_validation(database, source): + native = NativeHTTP( + [ + ("/changes/startPageToken", {"startPageToken": "old"}), + ("/files", {"files": [{"id": "a"}]}), + ("/files/a", folder("a")), + ( + "/changes", + {"changes": [{"changeType": "drive", "driveId": "lost", "removed": True}]}, + ), + ("/changes/startPageToken", {"startPageToken": "fresh"}), + ("/files", {"files": []}), + ("/changes", {"changes": [], "newStartPageToken": "after"}), + ("/files/a", folder("a")), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with client: + await run(pipeline, ctx, runtime) + assert not native.replies + async with database() as db: + row = await db.scalar(select(Entity)) + assert row.deleted_at is None + cycle = await source[0].read_cycle(db, pipeline._writer()) + assert cycle.starting_checkpoint.value == {"page_token": "fresh"} + + +async def test_rejected_inventory_token_restarts_sweep_not_boundary(database, source): + native = NativeHTTP( + [ + ("/changes/startPageToken", {"startPageToken": "before"}), + ("/files", {"files": [{"id": "a"}], "nextPageToken": "invalid"}), + ("/files/a", folder("a")), + ( + "/files", + ( + 400, + { + "error": { + "errors": [ + { + "location": "pageToken", + "locationType": "parameter", + "reason": "invalid", + } + ] + } + }, + ), + ), + ("/files", {"files": [{"id": "b"}]}), + ("/files/b", folder("b")), + ("/changes", {"changes": [], "newStartPageToken": "after"}), + ("/files/a", (404, {"error": {"message": "gone"}})), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with client: + await run(pipeline, ctx, runtime) + assert not native.replies + assert sum(path.endswith("/startPageToken") for path, _ in native.calls) == 1 + async with database() as db: + assert ( + await db.scalar(select(Entity).where(Entity.native_id == "a")) + ).removal_reason == "scope_removed" + + +@pytest.mark.parametrize("failure", ["incomplete", "forbidden"]) +async def test_provider_failure_cannot_promote_or_remove_prior_capture(database, source, failure): + from airweave.domains.sources.exceptions import SourceEntityForbiddenError + + ending = ( + ("/files", {"incompleteSearch": True}) + if failure == "incomplete" + else ("/files/b", (403, {"error": {"message": "forbidden"}})) + ) + listing = ( + {"files": [{"id": "a"}], "nextPageToken": "second"} + if failure == "incomplete" + else {"files": [{"id": "a"}, {"id": "b"}]} + ) + native = NativeHTTP( + [ + ("/changes/startPageToken", {"startPageToken": "before"}), + ("/files", listing), + ("/files/a", folder("a")), + ending, + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + async with client: + with pytest.raises(ValueError if failure == "incomplete" else SourceEntityForbiddenError): + await run(pipeline, ctx, runtime) + async with database() as db: + cycle = await source[0].read_cycle(db, pipeline._writer()) + assert cycle.promoted_checkpoint is None and cycle.phase == "active" + assert (await db.scalar(select(Entity))).deleted_at is None diff --git a/backend/airweave/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py index 54d3b6bd76..9394fdec08 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -16,6 +16,7 @@ from __future__ import annotations +import hashlib from datetime import datetime from typing import Any, AsyncGenerator, Dict, List, Optional @@ -25,8 +26,15 @@ 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.requests import CaptureRecord -from airweave.domains.entities.canonical.source import SourceObservation +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 ( SourceAuthProvider, @@ -52,8 +60,8 @@ ) from airweave.platform.sources._base import BaseSource from airweave.platform.sources.http_helpers import raise_for_status -from airweave.platform.sources.records.google_drive import generate_drive_observations 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 @@ -86,27 +94,69 @@ class GoogleDriveSource(BaseSource): canonical_record_types = ("file",) - async def generate_observations( + canonical_container_parents = {} + + @property + def capture_cycle_configuration(self) -> CycleConfiguration: + """One unfiltered accessible corpus; legacy path selections remain unsupported.""" + if self.include_patterns: + raise ValueError("Drive path selection is not yet supported by canonical capture") + return CycleConfiguration( + fingerprint=hashlib.sha256(b"drive-all-accessible-v1").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.""" + 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, + ) + + async def capture_page( self, + scope: CompletedScope, + continuation: ScanContinuation, *, - cursor: SyncCursor | None = None, - files: FileService | None = None, - node_selections: list[NodeSelectionData] | None = None, - ) -> AsyncGenerator[SourceObservation, None]: - """Capture original file metadata; missing body bytes remain explicitly metadata-only.""" - if self.include_patterns or node_selections: - raise ValueError("Drive path selection is not yet supported by canonical capture") - async for observation in generate_drive_observations(self._get, cursor): - if isinstance(observation, CaptureRecord) and files is not None: - observation = await capture_file_content( - observation, - files=files, - get=self._get, - client=self.http_client, - auth=self.auth, - logger=self.logger, - ) - yield observation + files: FileService, + parent: SourceRecord | None = None, + ) -> CapturePage: + """Commit one version-checked file without retaining a native response queue.""" + 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.""" + 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( @@ -161,6 +211,15 @@ async def _get(self, url: str, params: Optional[Dict] = None) -> Dict: 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/records/google_drive.py b/backend/airweave/platform/sources/records/google_drive.py index fac15cd806..36bb28b45b 100644 --- a/backend/airweave/platform/sources/records/google_drive.py +++ b/backend/airweave/platform/sources/records/google_drive.py @@ -1,45 +1,17 @@ -"""Drive metadata capture with boundary-before-scan and complete changes replay.""" +"""Native Drive identity and original record mapping shared with body capture.""" -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import Awaitable, Callable from datetime import datetime, timezone -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import JsonValue -from airweave.domains.entities.canonical.requests import ( - CaptureRecord, - CompletedScope, - RecordIdentity, - StartedScope, -) -from airweave.domains.entities.canonical.source import SourceObservation -from airweave.domains.sources.exceptions import SourceGoneError -from airweave.domains.syncs.cursors.cursor import SyncCursor +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]] -class DrivePage(BaseModel): - """Original item dictionaries survive validation intact.""" - - model_config = ConfigDict(extra="ignore") - files: list[dict[str, JsonValue]] = Field(default_factory=list) - changes: list[dict[str, JsonValue]] = Field(default_factory=list) - nextPageToken: str | None = None - newStartPageToken: str | None = None - incompleteSearch: bool = False - - -class DriveScopeChanged(Exception): - """A shared-drive membership change requires fresh corpus enumeration.""" - - -class DriveCheckpoint(BaseModel): - """In-memory result of draining one change stream, never persisted prematurely.""" - - token: str | None = None - - 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") @@ -56,102 +28,6 @@ def file_record(payload: dict[str, JsonValue], *, removed: bool = False) -> Capt 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")), ) - - -async def enumerate_files(get: GetJSON) -> AsyncGenerator[CaptureRecord, None]: - """Fail incomplete all-drive listings rather than deleting unenumerated records.""" - params: dict[str, str | int] = { - "pageSize": 1000, - "fields": "*", - "spaces": "drive", - "corpora": "allDrives", - "includeItemsFromAllDrives": "true", - "supportsAllDrives": "true", - } - seen: set[str] = set() - while True: - page = DrivePage.model_validate(await get(f"{BASE}/files", params=params)) - if page.incompleteSearch: - raise ValueError("Drive reported incomplete enumeration; checkpoint was not advanced") - for item in page.files: - yield file_record(item) - if not page.nextPageToken: - return - if page.nextPageToken in seen: - raise ValueError("Drive returned a repeated listing page token") - seen.add(page.nextPageToken) - params["pageToken"] = page.nextPageToken - - -async def replay_changes( - get: GetJSON, token: str, checkpoint: DriveCheckpoint -) -> AsyncGenerator[CaptureRecord, None]: - """Read every change page; metadata edits and trash state are real updates.""" - params: dict[str, str | int] = { - "pageToken": token, - "pageSize": 1000, - "fields": "*", - "spaces": "drive", - "includeRemoved": "true", - "includeCorpusRemovals": "true", - "includeItemsFromAllDrives": "true", - "supportsAllDrives": "true", - } - seen: set[str] = {token} - while True: - page = DrivePage.model_validate(await get(f"{BASE}/changes", params=params)) - for change in page.changes: - if change.get("changeType") == "drive": - # Drive membership changes require a new exhaustive file scan. - raise DriveScopeChanged("Shared-drive membership changed") - if change.get("removed") is True: - yield file_record({"id": change.get("fileId"), "change": change}, removed=True) - else: - item = change.get("file") - if not isinstance(item, dict): - raise ValueError("Drive file change omitted its current file resource") - yield file_record(item) - if not page.nextPageToken: - if not page.newStartPageToken: - raise ValueError("Drive change enumeration ended without newStartPageToken") - checkpoint.token = page.newStartPageToken - return - if page.nextPageToken in seen: - raise ValueError("Drive returned a repeated changes page token") - seen.add(page.nextPageToken) - params["pageToken"] = page.nextPageToken - - -async def generate_drive_observations( - get: GetJSON, cursor: SyncCursor | None -) -> AsyncGenerator[SourceObservation, None]: - """Only canonical checkpoints can skip bootstrap; legacy index cursors cannot.""" - state = cursor.data if cursor else {} - token = state.get("canonical_page_token") - for reset_attempt in range(2): - full = not token - if full: - boundary = await get( - f"{BASE}/changes/startPageToken", params={"supportsAllDrives": "true"} - ) - token = boundary.get("startPageToken") - if not isinstance(token, str) or not token: - raise ValueError("Drive omitted its starting change boundary") - yield StartedScope(record_type="file") - async for item in enumerate_files(get): - yield item - checkpoint = DriveCheckpoint() - try: - async for item in replay_changes(get, token, checkpoint): - yield item - except (DriveScopeChanged, SourceGoneError): - if reset_attempt: - raise - token = None - continue - if full: - yield CompletedScope(record_type="file") - if cursor is not None: - cursor.update(canonical_page_token=checkpoint.token) - return 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/tests/unit/platform/sources/records/test_google_drive.py b/backend/tests/unit/platform/sources/records/test_google_drive.py index 9dd8c63d9b..8e905dc13f 100644 --- a/backend/tests/unit/platform/sources/records/test_google_drive.py +++ b/backend/tests/unit/platform/sources/records/test_google_drive.py @@ -1,14 +1,15 @@ -"""Drive bootstrap/delta fidelity, scope completeness and checkpoint failure behavior.""" - -from copy import deepcopy -from uuid import uuid4 +"""Drive bounded acquisition; lifecycle checkpoint proof lives in canonical SQL tests.""" import pytest +from pydantic import ValidationError -from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope -from airweave.domains.syncs.cursors.cursor import SyncCursor -from airweave.platform.cursors.google_drive import GoogleDriveCursor -from airweave.platform.sources.records.google_drive import generate_drive_observations +from airweave.domains.entities.canonical.page_source import InvalidCaptureCheckpoint +from airweave.domains.sources.exceptions import ( + SourceEntityForbiddenError, + SourceEntityNotFoundError, +) +from airweave.platform.sources.records.google_drive import file_record +from airweave.platform.sources.records.google_drive_pages import DrivePages, _Progress class ScriptedGet: @@ -17,92 +18,112 @@ def __init__(self, replies): self.calls = [] async def __call__(self, url, params): - self.calls.append((url, deepcopy(params))) + self.calls.append((url, params.copy())) value = next(self.replies) if isinstance(value, Exception): raise value return value -def state(token=""): - return SyncCursor(uuid4(), GoogleDriveCursor, {"canonical_page_token": token}) +async def identity(record): + return record -async def test_initial_boundary_precedes_scan_and_replay_updates_metadata(): +async def test_pending_native_ids_resume_without_relisting_and_terminal_only_checkpoint(): get = ScriptedGet( [ - {"startPageToken": "before"}, - { - "files": [{"id": "one", "name": "old", "parents": ["folder-a"]}], - "nextPageToken": "page2", - }, - {"files": []}, + {"files": [{"id": "one"}, {"id": "two"}]}, + {"id": "one"}, + {"id": "two"}, { - "changes": [ - {"file": {"id": "one", "name": "new", "parents": ["folder-b"], "trashed": True}} - ], + "changes": [{"changeType": "file", "fileId": "one", "removed": True}], "newStartPageToken": "after", }, + {"id": "one", "trashed": True}, ] ) - cursor = state() - results = [item async for item in generate_drive_observations(get, cursor)] - records = [item for item in results if isinstance(item, CaptureRecord)] - assert get.calls[0][0].endswith("startPageToken") - assert get.calls[-1][1]["pageToken"] == "before" - assert [item.payload["name"] for item in records] == ["old", "new"] - assert records[-1].kind == "upsert" and records[-1].payload["trashed"] is True - assert all(item.completeness == "metadata_only" for item in records) - assert cursor.data["canonical_page_token"] == "after" - assert type(results[-1]) is CompletedScope - - -async def test_incremental_removal_pagination_and_no_full_scan(): - get = ScriptedGet( - [ - {"changes": [{"fileId": "gone", "removed": True}], "nextPageToken": "next"}, - {"changes": [], "newStartPageToken": "after"}, - ] + first = await DrivePages(get).page( + _Progress(phase="list", boundary="before").continuation(), identity ) - cursor = state("before") - results = [item async for item in generate_drive_observations(get, cursor)] - assert len(results) == 1 and results[0].kind == "delete" - assert results[0].removal_reason == "scope_removed" - assert cursor.data["canonical_page_token"] == "after" - assert all(url.endswith("/changes") for url, _ in get.calls) - - -async def test_incomplete_enumeration_cannot_mark_scope_complete_or_save_cursor(): - get = ScriptedGet([{"startPageToken": "before"}, {"files": [], "incompleteSearch": True}]) - cursor = state() - observed = [] - with pytest.raises(ValueError, match="incomplete enumeration"): - async for item in generate_drive_observations(get, cursor): - observed.append(item) - assert cursor.data["canonical_page_token"] == "" - assert [type(item) for item in observed] == [StartedScope] - - -async def test_shared_drive_membership_change_restarts_with_new_boundary(): - get = ScriptedGet( - [ - {"changes": [{"changeType": "drive", "removed": True, "driveId": "lost"}]}, - {"startPageToken": "new-before"}, - {"files": []}, - {"changes": [], "newStartPageToken": "new-after"}, - ] + assert first.provider_checkpoint is None and first.continuation.value["pending"] == ["two"] + second = await DrivePages(get).page(first.continuation, identity) + assert second.provider_checkpoint is None + final = await DrivePages(get).page(second.continuation, identity) + assert final.final and final.provider_checkpoint.value == {"page_token": "after"} + assert final.records[0].kind == "upsert" and final.records[0].payload["trashed"] + assert get.calls[3][1]["pageToken"] == "before" + assert len([url for url, _ in get.calls if url.endswith("/files")]) == 1 + + +@pytest.mark.parametrize( + "raw,error", + [ + ({"incompleteSearch": True}, ValueError), + ({"files": [{"id": str(n)} for n in range(101)]}, ValidationError), + ], +) +async def test_incomplete_or_oversized_listing_cannot_return_progress(raw, error): + with pytest.raises(error): + await DrivePages(ScriptedGet([raw])).page( + _Progress(phase="list", boundary="before").continuation(), identity + ) + + +async def test_shared_drive_event_requires_whole_cycle_restart(): + with pytest.raises(InvalidCaptureCheckpoint): + await DrivePages(ScriptedGet([{"changes": [{"changeType": "drive"}]}])).page( + _Progress(phase="changes", boundary="before", token="before").continuation(), identity + ) + + +async def test_exact_unavailability_is_not_a_general_provider_error(): + record = await DrivePages(ScriptedGet([SourceEntityNotFoundError("gone")])).current("one") + assert record.kind == "delete" and record.removal_reason == "scope_removed" + with pytest.raises(SourceEntityForbiddenError): + await DrivePages(ScriptedGet([SourceEntityForbiddenError("forbidden")])).current("one") + with pytest.raises(ValueError, match="different file"): + await DrivePages(ScriptedGet([{"id": "other"}])).current("one") + + +def test_continuation_bound_and_native_dates(): + state = _Progress( + phase="changes", + boundary="b" * 8192, + token="t" * 8192, + pending=tuple(f"{n:03d}" + ("x" * 253) for n in range(100)), + recent=tuple("a" * 64 for _ in range(128)), + loaded=True, ) - cursor = state("old") - results = [item async for item in generate_drive_observations(get, cursor)] - assert [type(item) for item in results] == [StartedScope, CompletedScope] - assert cursor.data["canonical_page_token"] == "new-after" + state.continuation() -async def test_failed_change_page_retains_saved_checkpoint(): - get = ScriptedGet( - [{"changes": [], "nextPageToken": "next"}, ConnectionError("provider unavailable")] +def test_native_dates_preserved_without_inventing_folder_edit_times(): + record = file_record( + {"id": "a", "createdTime": "2026-01-01T00:00:00Z", "modifiedTime": "2026-02-01T00:00:00Z"} + ) + assert record.source_created_at.month == 1 and record.source_updated_at.month == 2 + assert file_record({"id": "b"}).source_updated_at is None + + +async def test_actual_drive_page_source_preserves_factory_selected_node_rejection(): + from types import SimpleNamespace + from unittest.mock import MagicMock + + from airweave.domains.entities.canonical.page_source import CanonicalPageSource + from airweave.domains.sources.token_providers.static import StaticTokenProvider + from airweave.domains.sync_pipeline.factory import SourceBuildResult, SyncFactory + from airweave.platform.configs.config import GoogleDriveConfig + from airweave.platform.sources.google_drive import GoogleDriveSource + + source = await GoogleDriveSource.create( + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=MagicMock(), + config=GoogleDriveConfig(), + ) + assert isinstance(source, CanonicalPageSource) + result = SourceBuildResult( + source=source, cursor=None, files=None, node_selections=[MagicMock()] ) - cursor = state("before") - with pytest.raises(ConnectionError): - _ = [item async for item in generate_drive_observations(get, cursor)] - assert cursor.data["canonical_page_token"] == "before" + with pytest.raises(ValueError, match="does not support selected nodes"): + SyncFactory._build_stream(None, SimpleNamespace(source=source), result, MagicMock()) From 3357dd48a5089e10fef97ac4cc99eed0944a94ee Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:35:46 -0700 Subject: [PATCH 099/266] Use disposable fixture database in offline lifecycle child test --- .../entities/canonical/tests/test_gmail_resume_probe.py | 3 +++ 1 file changed, 3 insertions(+) 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 index 1f368f041f..48058cd6e2 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py @@ -208,6 +208,9 @@ async def synthetic_source(name, account, expected_email, key, fence, **options) 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' From de7ed36b6ee89fbae44bd7b834fc57f573bb00e8 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:39:13 -0700 Subject: [PATCH 100/266] Discard Calendar delta token when occurrence access is lost --- .../tests/test_calendar_exclusions.py | 53 +++++++++++++++++++ .../sources/records/google_calendar.py | 5 ++ .../sources/records/test_google_calendar.py | 34 ++++++++++++ 3 files changed, 92 insertions(+) diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index 302d46cde4..342dd4228e 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -185,3 +185,56 @@ async def test_operational_occurrences_publish_zero_documents_without_retry(data for w in await projections.pending(db, fence.organization_id, fence.sync_id) ) assert (await db.get(Entity, row.id)).indexed_chunk_count == 0 + + +async def test_occurrence_access_loss_then_full_reacquisition_restores_original(database, source): + from airweave.domains.entities.canonical.tests.test_capture_pipeline import components + from airweave.domains.sources.exceptions import SourceEntityNotFoundError + from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline + from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt + from airweave.domains.syncs.cursors.cursor import SyncCursor + from airweave.platform.configs.config import CalendarOccurrenceWindow + from airweave.platform.sources.records.google_calendar import generate_calendar_observations + + capture, fence = source + state = SyncCursor(fence.sync_id) + window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") + native = {"id": "unchanged", "status": "confirmed", "summary": "Original"} + + async def run(restored, attempt): + async def get(url, params): + if url.endswith("calendarList"): + return {"items": [{"id": "cal", "timeZone": "UTC"}]} + if params["singleEvents"] == "true": + if not restored: + raise SourceEntityNotFoundError("synthetic access loss") + return {"items": []} + assert "syncToken" not in params + return {"items": [native], "nextSyncToken": "restored" if restored else "discard"} + + ctx, _, runtime, bus = components(database, source) + pipeline = CanonicalCapturePipeline( + capture, + database, + bus, + ("calendar", "event", "event_occurrence"), + CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), + {"event": "calendar", "event_occurrence": "calendar"}, + ) + await pipeline.start(ctx) + async for observation in generate_calendar_observations(get, state, window): + await pipeline.process([observation], ctx, runtime) + await pipeline.cleanup_orphaned_entities(ctx, runtime) + + await run(False, 1) + async with database() as db: + entity = await db.scalar(select(Entity).where(Entity.native_id == "unchanged")) + record_id = entity.id + lost = await capture.store.read(db, fence.organization_id, fence.sync_id, record_id) + assert lost.content_access == "unavailable" and lost.payload == {} + assert state.data["calendar_tokens"] == {} + await run(True, 2) + async with database() as db: + restored = await capture.store.read(db, fence.organization_id, fence.sync_id, record_id) + assert restored.content_access == "available" + assert restored.deleted_at is None and restored.payload == native diff --git a/backend/airweave/platform/sources/records/google_calendar.py b/backend/airweave/platform/sources/records/google_calendar.py index a3e76bb37d..61aca399db 100644 --- a/backend/airweave/platform/sources/records/google_calendar.py +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -223,6 +223,11 @@ async def _capture_member( scan_id=uuid4(), ) except SourceEntityNotFoundError: + # Master enumeration can finish before expanded instances lose access. + # Its token cannot survive withdrawing those originals: restoration must + # reacquire unchanged masters, which an incremental response may omit. + tokens.pop(calendar_id, None) + coverage.pop(calendar_id, None) yield CaptureRecord( identity=RecordIdentity(record_type="calendar", native_id=calendar_id), payload={"id": calendar_id}, diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py index fa8d60a433..8dd26ed4d7 100644 --- a/backend/tests/unit/platform/sources/records/test_google_calendar.py +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -315,3 +315,37 @@ async def test_selected_calendar_loses_access_between_membership_and_events(): revoked = [item for item in result if isinstance(item, RemovedScope)] assert len(revoked) == 2 and all(item.removal_reason == "access_revoked" for item in revoked) assert state.data["calendar_tokens"] == {} + + +async def test_occurrence_access_loss_discards_master_token_before_restoration(): + from airweave.domains.sources.exceptions import SourceEntityNotFoundError + from airweave.platform.configs.config import CalendarOccurrenceWindow + + window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") + state = cursor({"cal": "old"}) + get = ScriptedGet( + [ + {"items": [{"id": "cal", "timeZone": "UTC"}]}, + {"items": [{"id": "unchanged"}], "nextSyncToken": "must-discard"}, + SourceEntityNotFoundError("synthetic access loss"), + ] + ) + lost = [item async for item in generate_calendar_observations(get, state, window)] + assert state.data["calendar_tokens"] == {} + assert state.data["occurrence_coverage"] == {} + assert {item.record_type for item in lost if isinstance(item, RemovedScope)} == { + "event", + "event_occurrence", + } + restored_get = ScriptedGet( + [ + {"items": [{"id": "cal", "timeZone": "UTC"}]}, + {"items": [{"id": "unchanged"}], "nextSyncToken": "restored"}, + {"items": []}, + ] + ) + restored = [item async for item in generate_calendar_observations(restored_get, state, window)] + assert "syncToken" not in restored_get.calls[1][1] + assert any(isinstance(item, StartedScope) and item.record_type == "event" for item in restored) + assert state.data["calendar_tokens"] == {"cal": "restored"} + assert "cal" in state.data["occurrence_coverage"] From 35cb5fbf0a0c57c6a79cb5f9432d5f08c8222175 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:42:39 -0700 Subject: [PATCH 101/266] Record two completed Gmail query cycles after retry correction --- .../gmail-query-retry-fixed-20260930.json | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 backend/tests/live/evidence/gmail-query-retry-fixed-20260930.json diff --git a/backend/tests/live/evidence/gmail-query-retry-fixed-20260930.json b/backend/tests/live/evidence/gmail-query-retry-fixed-20260930.json new file mode 100644 index 0000000000..70a6eeecea --- /dev/null +++ b/backend/tests/live/evidence/gmail-query-retry-fixed-20260930.json @@ -0,0 +1,46 @@ +{ + "verified_at": "2026-09-30", + "source_commit": "1b28bf7", + "provider": "gmail", + "mode": "frozen_seven_day_filtered_full_reconciliation", + "provider_writes": false, + "harness_exit_code": 0, + "runs": [ + { + "run": 1, + "provider_requests": 302, + "records_observed": 242, + "stored_records": 242, + "partial_records": 0, + "blob_bytes_written": 5660964, + "full_scope_completed": true, + "loaded_durable_checkpoint": false, + "change_sequence": 242 + }, + { + "run": 2, + "provider_requests": 302, + "records_observed": 242, + "stored_records": 242, + "partial_records": 0, + "blob_bytes_written": 5660964, + "full_scope_completed": true, + "loaded_durable_checkpoint": true, + "change_sequence": 277 + } + ], + "aggregate_requests": 604, + "aggregate_observations": 484, + "identical_second_read": false, + "schema_removed": true, + "blob_directory_removed": true, + "independent_private_directories": 0, + "independent_private_schemas": 0, + "immutable_archive_removed": true, + "limitations": [ + "Frozen seven-day query, not whole-mailbox baseline or history-delta proof.", + "Fresh second job loads durable cycle state; this trial does not interrupt a page.", + "Second raw payload/revision digest changed; no claim of identical provider responses.", + "Local disposable database and capture harness, not hosted owner-binding or staging proof." + ] +} From 46172672bcc721fd8bd8b0761d3d28dd4c7d97e1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 10:58:33 -0700 Subject: [PATCH 102/266] Verify Drive lifecycle against durable page checkpoints --- .../tests/test_drive_lifecycle_probe.py | 90 +++++++++++++++++++ .../drive-page-harness-gap-20260930.json | 20 +++++ backend/tests/live/provider_lifecycle.py | 28 ++++-- 3 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py create mode 100644 backend/tests/live/evidence/drive-page-harness-gap-20260930.json diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py new file mode 100644 index 0000000000..72ff306dd2 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py @@ -0,0 +1,90 @@ +"""Actual lifecycle children verify the Drive page checkpoint contract offline.""" + +import json + +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 unittest.mock import MagicMock +sys.path.insert(0, 'tests/live') +import conftest +import httpx +import provider_lifecycle as lifecycle +from provider_sample import verify_rest_identity +from airweave.domains.entities.canonical.tests.test_drive_recovery import NativeHTTP, folder +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.platform.configs.config import GoogleDriveConfig +from airweave.platform.sources.google_drive import GoogleDriveSource + +@asynccontextmanager +async def synthetic_source(name, account, expected_email, key, fence, **options): + assert name == 'google_drive' + responses = [('/about', {'user': {'emailAddress': expected_email}})] + if os.environ['TEST_STAGE'] == 'first': + responses += [('/changes/startPageToken', {'startPageToken': 'before'}), + ('/files', {'files': [{'id': 'folder-a'}]}), + ('/files/folder-a', folder('folder-a')), + ('/changes', {'changes': [], 'newStartPageToken': 'after'})] + else: + responses += [('/changes', {'changes': [], 'newStartPageToken': 'later'})] + native = NativeHTTP(responses) + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle), + event_hooks={'request': [options['request_hook']]}) as client: + connector = await GoogleDriveSource.create(auth=StaticTokenProvider('synthetic'), + logger=MagicMock(), http_client=client, config=GoogleDriveConfig()) + await verify_rest_identity(name, connector, expected_email) + yield connector, 'provider_email' + assert not native.replies + +lifecycle.rest_source = synthetic_source +# All provider I/O is synthetic; reuse the disposable fixture (TCP in CI). +# The real account harness retains its private Unix-socket requirement. +lifecycle.harness.test_database_url = lambda: os.environ['CANONICAL_TEST_DATABASE_URL'] +os.environ['COMPOSIO_API_KEY'] = 'synthetic-test-only' +os.environ['LIVE_DRIVE_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'])))) +""" + + +async def test_drive_lifecycle_initial_then_fresh_process_changes(database, source, tmp_path): + fence = source[1] + async with database() as db: + schema = await db.scalar(text("select current_schema()")) + manifest = { + "provider": "google_drive", + "schema": schema, + "root": str(tmp_path), + "organization_id": str(fence.organization_id), + "sync_id": str(fence.sync_id), + "record_limit": 10, + "request_limit": 20, + "timeout": 30, + "file_byte_limit": 1024, + "blob_byte_limit": 1024, + } + results = [] + for stage in ("first", "second"): + current = dict(manifest) + if stage == "first": + current.update(job_id=str(fence.job_id), attempt_number=2) + code, output, error = await subprocess_script( + SCRIPT, {"TEST_MANIFEST": json.dumps(current), "TEST_STAGE": stage} + ) + if code != 0: + raise AssertionError(output + error) + result = json.loads(output.splitlines()[-1]) + assert not result["failed"] and result["job_status"] == "completed" + assert result["stored_records"] == 1 and result["partial_records"] == 0 + results.append(result) + first, second = results + assert first["provider_requests"] == 5 and second["provider_requests"] == 2 + assert first["records_observed"] == 1 and second["records_observed"] == 0 + assert not first["loaded_durable_checkpoint"] and second["loaded_durable_checkpoint"] + assert first["resumed_changes_requests"] == 0 and second["resumed_changes_requests"] == 1 + assert first["observed_change_sequence"] == second["observed_change_sequence"] + assert first["payload_revision_digest"] == second["payload_revision_digest"] diff --git a/backend/tests/live/evidence/drive-page-harness-gap-20260930.json b/backend/tests/live/evidence/drive-page-harness-gap-20260930.json new file mode 100644 index 0000000000..8b1004e5b1 --- /dev/null +++ b/backend/tests/live/evidence/drive-page-harness-gap-20260930.json @@ -0,0 +1,20 @@ +{ + "source_commit": "35cb5fb", + "verified_at": "2026-09-30", + "provider": "google_drive", + "harness_exit_code": 1, + "job_status": "completed", + "error_type": "KeyError", + "identified_harness_gap": "validate_checkpoint expects removed legacy canonical_page_token; new adapter publishes canonical_cycle checkpoint", + "provider_requests": 140, + "records_observed": 55, + "blob_bytes_written": 207614203, + "second_cycle_ran": false, + "complete_trial_verified": false, + "schema_removed": true, + "blob_directory_removed": true, + "independent_private_schemas": 0, + "independent_private_directories": 0, + "immutable_archive_removed": true, + "provider_writes": false +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 89e5a23438..d3daa83cd4 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -218,8 +218,12 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert counters["started"] == counters["completed"] == expected_scopes assert bool(counters["sync_token_requests"]) == loaded elif name == "google_drive": - assert saved["canonical_page_token"] - assert counters["started"] == counters["completed"] == (0 if loaded else 1) + cycle = saved["canonical_cycle"] + assert cycle["phase"] == "complete" and cycle["completed_job_id"] + assert cycle["last_full_capture"] and cycle["promoted_checkpoint"] + assert cycle["promoted_checkpoint"]["checkpoint"]["value"]["page_token"] + assert cycle["mode"] == ("changes" if loaded else "full") + assert counters["started"] == counters["completed"] == 0 assert bool(counters["resumed_changes_requests"]) == loaded elif name in {"slack", "wispr"}: assert saved["canonical_cycle"]["phase"] == "complete" @@ -302,6 +306,19 @@ def count_gmail_request(name, request, previous, counters, identity_verified): counters["resumed_changes_requests"] += 1 +def count_drive_request(name, request, previous, counters): + """Count only an actual request using the previously published Drive boundary.""" + if name != "google_drive" or not request.url.path.endswith("/changes"): + return + promoted = previous.get("canonical_cycle", {}).get("promoted_checkpoint") + if ( + promoted + and request.url.params.get("pageToken") + == promoted["checkpoint"]["value"]["page_token"] + ): + counters["resumed_changes_requests"] += 1 + + async def child(manifest): engine = create_async_engine( harness.test_database_url(), @@ -361,12 +378,7 @@ async def request_hook(request): count_gmail_request(name, request, previous, counters, identity_verified) if request.url.params.get("syncToken"): counters["sync_token_requests"] += 1 - if ( - request.url.path.endswith("/changes") - and previous.get("canonical_page_token") - and request.url.params.get("pageToken") == previous["canonical_page_token"] - ): - counters["resumed_changes_requests"] += 1 + count_drive_request(name, request, previous, counters) @asynccontextmanager async def db_context(): From b61f73bc5d4b3c8ac758013b5aa58412ff6b8ef2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:03:53 -0700 Subject: [PATCH 103/266] Persist fenced per-scope full and changes capture evidence --- .../domains/entities/canonical/coverage.py | 124 ++++- .../entities/canonical/cycle_models.py | 30 +- .../domains/entities/canonical/cycle_store.py | 39 +- .../domains/entities/canonical/page_source.py | 36 +- .../domains/entities/canonical/scan_models.py | 5 + .../domains/entities/canonical/scan_store.py | 124 ++++- .../entities/canonical/scope_execution.py | 55 ++ .../entities/canonical/tests/conftest.py | 1 + .../canonical/tests/test_mixed_scopes.py | 497 ++++++++++++++++++ .../canonical/tests/test_owned_search.py | 1 + .../sync_pipeline/canonical_capture.py | 4 +- .../domains/sync_pipeline/canonical_scan.py | 188 +++++-- backend/airweave/models/capture_scan.py | 1 + .../alembic/versions/0008_scope_execution.py | 21 + 14 files changed, 1060 insertions(+), 66 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/scope_execution.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py create mode 100644 backend/alembic/versions/0008_scope_execution.py diff --git a/backend/airweave/domains/entities/canonical/coverage.py b/backend/airweave/domains/entities/canonical/coverage.py index 3fabd5d2bb..fe237eacb0 100644 --- a/backend/airweave/domains/entities/canonical/coverage.py +++ b/backend/airweave/domains/entities/canonical/coverage.py @@ -3,8 +3,8 @@ from typing import Literal from uuid import UUID -from pydantic import AwareDatetime, BaseModel, ConfigDict, ValidationError -from sqlalchemy import select +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError, model_validator +from sqlalchemy import and_, case, func, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.cycle_models import ( @@ -12,6 +12,12 @@ CaptureCycle, CompletionPolicy, ) +from airweave.domains.entities.canonical.cycle_store import root_ready +from airweave.domains.entities.canonical.requests import WriterFence +from airweave.domains.entities.canonical.store import content_is_available +from airweave.models.capture_scan import CaptureScan +from airweave.models.entity import Entity +from airweave.models.sync import Sync from airweave.models.sync_cursor import SyncCursor @@ -24,18 +30,118 @@ class FullCaptureCoverage(BaseModel): discovery: Literal["incomplete", "scope_enumeration_complete"] +class CaptureScopeSummary(BaseModel): + """Counts describe eligible exact scopes, never original/indexed record totals.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + eligible: int = Field(ge=0, strict=True) + completed_full: int = Field(ge=0, strict=True) + completed_changes: int = Field(ge=0, strict=True) + unfinished: int = Field(ge=0, strict=True) + + @model_validator(mode="after") + def partition(self): + """Every eligible scope belongs to exactly one progress category.""" + if self.eligible != self.completed_full + self.completed_changes + self.unfinished: + raise ValueError("Scope progress must partition eligible scopes") + return self + + class CaptureCoverage(BaseModel): """No cycle means unknown; completion only certifies the declared scope policies.""" model_config = ConfigDict(extra="forbid", frozen=True) phase: Literal["active", "complete"] - mode: Literal["full", "changes"] = "full" + mode: Literal["full", "changes", "mixed"] = "full" + scope_summary: CaptureScopeSummary | None = None last_full_capture: FullCaptureCoverage | None = None provider_checkpoint_promoted_at: AwareDatetime | None = None policies: dict[str, CompletionPolicy] discovery: Literal["incomplete", "pending", "scope_enumeration_complete"] +async def mixed_scope_summary( + db: AsyncSession, organization_id: UUID, sync_id: UUID, cycle: CaptureCycle +) -> CaptureScopeSummary | None: + """Aggregate the current visible forest, withholding counts until roots are refreshed.""" + sync = await db.scalar( + select(Sync).where(Sync.id == sync_id, Sync.organization_id == organization_id) + ) + if sync is None or sync.writer_job_id is None or sync.writer_attempt_id is None: + return None + fence = WriterFence( + organization_id=organization_id, + sync_id=sync_id, + job_id=sync.writer_job_id, + epoch=sync.writer_epoch, + attempt_id=sync.writer_attempt_id, + attempt_number=sync.writer_attempt_number, + ) + if not await root_ready(db, fence, cycle): + return None + eligible = full = changes = 0 + common = ( + CaptureScan.organization_id == organization_id, + CaptureScan.sync_id == sync_id, + CaptureScan.cycle_id == cycle.version.cycle_id, + CaptureScan.phase == "complete", + ) + for kind in cycle.configuration.root_record_types: + eligible += 1 + mode = await db.scalar( + select(CaptureScan.execution_state["plan"]["mode"].astext).where( + *common, + CaptureScan.record_type == kind, + CaptureScan.parent_record_id.is_(None), + CaptureScan.container_id.is_(None), + ) + ) + full += int(mode == "full") + changes += int(mode == "changes") + for kind, parents in cycle.configuration.parents.items(): + parent_types = tuple(parent for parent in parents if parent is not None) + if not parent_types: + continue + completed = and_( + *common, + CaptureScan.record_type == kind, + CaptureScan.parent_record_id == Entity.id, + CaptureScan.parent_visibility_epoch == Entity.visibility_epoch, + ) + if cycle.configuration.children_of(kind): + completed = and_(completed, CaptureScan.membership_attempt_id == fence.attempt_id) + counts = ( + await db.execute( + select( + func.count(Entity.id), + func.count( + case((CaptureScan.execution_state["plan"]["mode"].astext == "full", 1)) + ), + func.count( + case((CaptureScan.execution_state["plan"]["mode"].astext == "changes", 1)) + ), + ) + .select_from(Entity) + .outerjoin(CaptureScan, completed) + .where( + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Entity.entity_definition_short_name.in_(parent_types), + content_is_available(), + ) + ) + ).one() + eligible += counts[0] + full += counts[1] + changes += counts[2] + return CaptureScopeSummary( + eligible=eligible, + completed_full=full, + completed_changes=changes, + unfinished=eligible - full - changes, + ) + + async def capture_coverage( db: AsyncSession, organization_id: UUID, sync_ids: tuple[UUID, ...] ) -> dict[UUID, CaptureCoverage]: @@ -74,11 +180,23 @@ async def capture_coverage( ): continue discovery = evidence.discovery + summary = ( + await mixed_scope_summary(db, organization_id, row.sync_id, cycle) + if cycle.mode == "mixed" + else None + ) + if cycle.mode == "mixed" and ( + evidence is not None + or promoted is not None + or (cycle.phase == "complete" and (summary is None or summary.unfinished != 0)) + ): + continue # Withhold unknown/incompatible coverage instead of emitting a false promise. result[row.sync_id] = CaptureCoverage( phase=cycle.phase, policies=policies, discovery=discovery, mode=cycle.mode, + scope_summary=summary, last_full_capture=FullCaptureCoverage( cycle_id=evidence.cycle_id, completed_at=evidence.completed_at, diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 713f0d1d57..1a38e0b2df 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -2,10 +2,11 @@ import hashlib import json -from typing import Literal +from typing import Annotated, Literal from uuid import UUID from pydantic import ( + AfterValidator, AwareDatetime, BaseModel, ConfigDict, @@ -20,13 +21,23 @@ CompletionPolicy = Literal["exhaustive", "discovery_only", "discovery_with_validation"] -CaptureMode = Literal["full", "changes"] +CaptureMode = Literal["full", "changes", "mixed"] CompletedDiscovery = Literal["incomplete", "scope_enumeration_complete"] CYCLE_KEY = "canonical_cycle" TERMINAL_CHECKPOINT_KEY = "_canonical_terminal_checkpoint" +def bounded_source_plan(value: dict[str, JsonValue]) -> dict[str, JsonValue]: + """Immutable cycle parameters must not become a response archive.""" + if len(json.dumps(value, allow_nan=False).encode()) > 65536: + raise ValueError("Cycle source plan exceeds64KiB") + return value + + +SourcePlan = Annotated[dict[str, JsonValue], AfterValidator(bounded_source_plan)] + + class CycleConfiguration(BaseModel): """Immutable copy of the source's allowed record-parent relationships.""" @@ -35,12 +46,15 @@ class CycleConfiguration(BaseModel): parents: dict[RecordKind, tuple[RecordKind | None, ...]] completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) known_object_validation: tuple[RecordKind, ...] = () + scope_changes: tuple[RecordKind, ...] = () def digest(self) -> str: """Bind history to scope, topology and guarantees, including default policies.""" value = self.model_dump(mode="json") if not self.known_object_validation: value.pop("known_object_validation") # Preserve existing default configuration digests. + if not self.scope_changes: + value.pop("scope_changes") value["completion_policies"] = {kind: self.policy(kind) for kind in self.parents} return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest() @@ -57,6 +71,7 @@ def from_source( container_parents: dict[str, str | tuple[str | None, ...]], completion_policies: dict[RecordKind, CompletionPolicy] | None = None, known_object_validation: tuple[RecordKind, ...] = (), + scope_changes: tuple[RecordKind, ...] = (), ) -> "CycleConfiguration": """Snapshot the existing source declaration; no second topology registry.""" if len(set(record_types)) != len(record_types): @@ -72,6 +87,7 @@ def from_source( parents=parents, completion_policies=completion_policies or {}, known_object_validation=known_object_validation, + scope_changes=scope_changes, ) @model_validator(mode="before") @@ -97,6 +113,10 @@ def reachable_types(self) -> "CycleConfiguration": set(self.known_object_validation) - set(self.parents) ): raise ValueError("Known-object validation requires unique declared kinds") + if len(set(self.scope_changes)) != len(self.scope_changes) or any( + kind not in self.parents or self.children_of(kind) for kind in self.scope_changes + ): + raise ValueError("Scope changes require unique declared leaf kinds") if set(self.completion_policies) - set(self.parents): raise ValueError("Completion policy has an undeclared record kind") if not self.parents or any(not kind or not values for kind, values in self.parents.items()): @@ -195,6 +215,8 @@ class CaptureCycle(BaseModel): completed_job_id: UUID | None = None mode: CaptureMode = "full" starting_checkpoint: ProviderCheckpoint | None = None + source_plan: SourcePlan = Field(default_factory=dict) + force_full_scopes: bool = False promoted_checkpoint: PromotedCheckpoint | None = None last_full_capture: FullCaptureEvidence | None = None @@ -218,6 +240,8 @@ class BeginCycle(BaseModel): expected: CycleVersion | None = None mode: CaptureMode = "full" starting_checkpoint: ProviderCheckpoint | None = None + source_plan: SourcePlan = Field(default_factory=dict) + force_full_scopes: bool = False class CompleteCycle(BaseModel): @@ -238,6 +262,8 @@ class RestartCycle(BaseModel): configuration: CycleConfiguration mode: CaptureMode = "full" starting_checkpoint: ProviderCheckpoint | None = None + source_plan: SourcePlan = Field(default_factory=dict) + force_full_scopes: bool = False class ScopeWork(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 3109655e7b..260444c504 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -74,6 +74,12 @@ def next_cycle(previous: CaptureCycle | None, request: BeginCycle | RestartCycle evidence = None if promoted is not None and promoted.configuration_digest != digest: promoted = None + if request.mode == "mixed": + if not request.configuration.scope_changes or request.starting_checkpoint is not None: + raise CycleConflict("Mixed capture requires declared scope changes and no global token") + evidence = promoted = None + elif request.configuration.scope_changes: + raise CycleConflict("Scoped changes require explicit mixed cycle execution") starting = request.starting_checkpoint if request.mode == "changes": if len(request.configuration.parents) != 1 or next( @@ -94,11 +100,28 @@ def next_cycle(previous: CaptureCycle | None, request: BeginCycle | RestartCycle configuration=request.configuration, mode=request.mode, starting_checkpoint=starting, + source_plan=request.source_plan, + force_full_scopes=request.force_full_scopes, promoted_checkpoint=promoted, last_full_capture=evidence, ) +def attest_active_plan(state: CaptureCycle, request: BeginCycle) -> None: + """Resume never changes the already persisted execution intent.""" + if state.configuration != request.configuration: + raise CycleConflict("Active cycle configuration changed; explicit abandonment required") + if state.force_full_scopes != request.force_full_scopes: + raise CycleConflict("Active cycle force-full intent changed") + if state.source_plan != request.source_plan: + raise CycleConflict("Active cycle source plan changed; explicit restart required") + if state.mode != request.mode or ( + request.starting_checkpoint is not None + and request.starting_checkpoint != state.starting_checkpoint + ): + raise CycleConflict("Active cycle execution plan changed; explicit restart required") + + async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: """An explicit completed-version CAS starts the next cycle.""" cursor = await cursor_row(db, request.fence) @@ -107,17 +130,7 @@ async def begin_cycle(db: AsyncSession, request: BeginCycle) -> CaptureCycle: if request.expected is not None and request.expected != state.version: raise CycleConflict("Cycle changed; reload durable progress") if state.phase == "active": - if state.configuration != request.configuration: - raise CycleConflict( - "Active cycle configuration changed; explicit abandonment required" - ) - if state.mode != request.mode or ( - request.starting_checkpoint is not None - and request.starting_checkpoint != state.starting_checkpoint - ): - raise CycleConflict( - "Active cycle execution plan changed; explicit restart required" - ) + attest_active_plan(state, request) return state if request.expected is None: return state @@ -334,6 +347,10 @@ async def terminal_checkpoint( ) -> PromotedCheckpoint | None: """Attest the source's terminal scan before publishing its provider boundary.""" terminal = request.terminal_checkpoint + if state.mode == "mixed": + if terminal is not None: + raise CycleConflict("Mixed capture has no global provider checkpoint") + return None if (state.starting_checkpoint is not None or state.mode == "changes") and terminal is None: raise CycleConflict("Checkpoint-bearing capture requires its terminal scan boundary") promoted = state.promoted_checkpoint if state.mode == "changes" else None diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index c6425c1f48..023cdebeb7 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -9,6 +9,7 @@ CaptureMode, CycleConfiguration, ProviderCheckpoint, + SourcePlan, ) from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.requests import ( @@ -16,7 +17,8 @@ CompletedScope, ScopeRemovalReason, ) -from airweave.domains.entities.canonical.scan_models import ScanContinuation +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 @@ -41,6 +43,7 @@ class CapturePlan(BaseModel): 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): @@ -113,3 +116,34 @@ class KnownObjectSource(Protocol): 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/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index 838204ccc5..cb3090f0eb 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -18,6 +18,7 @@ ScanVersion, WriterFence, ) +from airweave.domains.entities.canonical.scope_execution import ScopeExecution, ScopePlan class ScanContinuation(BaseModel): @@ -47,6 +48,9 @@ class BeginScan(BaseModel): fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") expected: ScanVersion | None = None restart: bool = False + plan: ScopePlan | None = None + expected_parent_epoch: int | None = Field(default=None, ge=1) + expected_parent_revision: int | None = Field(default=None, ge=1) continuation: ScanContinuation = Field(default_factory=ScanContinuation) @@ -67,6 +71,7 @@ class ScanState(BaseModel): completion_policy: CompletionPolicy = "exhaustive" mode: Literal["full", "changes"] = "full" provider_checkpoint: ProviderCheckpoint | None = None + execution: ScopeExecution | None = None class CommitScanPage(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 18a44c79ba..596d308859 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -10,6 +10,7 @@ from airweave.domains.entities.canonical.cycle_models import ( TERMINAL_CHECKPOINT_KEY, + CaptureCycle, CaptureMode, CycleVersion, ProviderCheckpoint, @@ -42,6 +43,11 @@ ScanState, ScanVersion, ) +from airweave.domains.entities.canonical.scope_execution import ( + ScopeEvidence, + ScopeExecution, + ScopePlan, +) from airweave.domains.entities.canonical.store import ( CanonicalRecordStore, CanonicalStoreError, @@ -94,6 +100,9 @@ def scan_state(row: CaptureScan, parent: Entity | None = None) -> ScanState: completed_at=row.completed_at, parent_visibility_epoch=row.parent_visibility_epoch, membership_attempt_id=row.membership_attempt_id, + execution=ScopeExecution.model_validate(row.execution_state) + if row.execution_state + else None, ) @@ -111,6 +120,98 @@ def page_continuation(request: CommitScanPage, mode: CaptureMode) -> dict[str, J return continuation +def effective_mode(row: CaptureScan, cycle: CaptureCycle) -> str: + """A mixed cycle never supplies a default cleanup mode for its scopes.""" + if cycle.mode != "mixed": + return cycle.mode + if row.execution_state is None and row.cycle_id != cycle.version.cycle_id: + return "full" # Legacy evidence is unknown; prepare a new full scope. + if row.execution_state is None: + raise ScanConflict("Mixed scope is missing its immutable execution plan") + return ScopeExecution.model_validate(row.execution_state).plan.mode + + +def publish_scope(row: CaptureScan, cycle: CaptureCycle) -> None: + """Completed evidence is derived only from this row's committed terminal state.""" + if cycle.mode != "mixed" or row.phase != "complete": + return + execution = ScopeExecution.model_validate(row.execution_state) + checkpoint = row.continuation.get(TERMINAL_CHECKPOINT_KEY) + evidence = ScopeEvidence( + cycle_id=row.cycle_id, + sweep_id=row.sweep_id, + completed_at=row.completed_at, + parent_visibility_epoch=row.parent_visibility_epoch, + request_context=execution.plan.request_context, + policy=cycle.configuration.policy(row.record_type), + checkpoint=ProviderCheckpoint.model_validate(checkpoint) if checkpoint else None, + ) + row.execution_state = execution.model_copy( + update={ + "published": evidence, + "last_full": evidence if execution.plan.mode == "full" else execution.last_full, + } + ).model_dump(mode="json") + + +def begin_execution( + row: CaptureScan, cycle: CaptureCycle, request: BeginScan, epoch: int | None +) -> None: + """Previous completion survives replacement; mismatching owner/request invalidates it.""" + if cycle.mode != "mixed": + if request.plan is not None: + raise ScanConflict("Scope plans require mixed execution") + row.execution_state = None + return + plan = request.plan + if plan is None: + raise ScanConflict("Mixed capture requires an explicit scope plan") + previous = ( + ScopeExecution.model_validate(row.execution_state) + if row.execution_state and row.fingerprint == cycle.configuration.fingerprint + else None + ) + full = previous.last_full if previous else None + published = previous.published if previous else None + policy = cycle.configuration.policy(request.scope.record_type) + if full and (not full.matches(plan, epoch) or full.policy != policy): + full = None + if published and (not published.matches(plan, epoch) or published.policy != policy): + published = None + validate_scope_changes(cycle, request.scope.record_type, plan, full, published) + row.execution_state = ScopeExecution(plan=plan, last_full=full, published=published).model_dump( + mode="json" + ) + + +def validate_scope_changes( + cycle: CaptureCycle, + kind: str, + plan: ScopePlan, + full: ScopeEvidence | None, + published: ScopeEvidence | None, +) -> None: + """Only compatible completed evidence authorizes an incremental scope.""" + if plan.mode == "changes": + if cycle.force_full_scopes: + raise ScanConflict("This cycle requires full scope scans") + if kind not in cycle.configuration.scope_changes: + raise ScanConflict("Scope kind does not permit native changes") + if full is None or published is None or published.checkpoint is None: + raise ScanConflict("Scope changes require compatible full evidence and checkpoint") + if plan.starting_checkpoint != published.checkpoint: + raise ScanConflict("Scope changes must start at its last published checkpoint") + + +def attest_planned_owner(cycle: CaptureCycle, request: BeginScan, parent: Entity | None) -> None: + """The post-I/O plan must still describe the captured owner used to select it.""" + if cycle.mode == "mixed" and ( + request.expected_parent_epoch != (parent.visibility_epoch if parent else None) + or request.expected_parent_revision != (parent.record_revision if parent else None) + ): + raise ScanConflict("Scope owner changed during plan selection") + + class CanonicalScanStore: """Reuse record capture, lock ordering and journal semantics for recoverable scans.""" @@ -161,7 +262,7 @@ async def _state(self, db: AsyncSession, row: CaptureScan) -> ScanState: "completion_policy": cycle.configuration.policy(row.record_type) if cycle else "exhaustive", - "mode": cycle.mode if cycle else "full", + "mode": effective_mode(row, cycle) if cycle else "full", } ) @@ -195,6 +296,7 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) await attest_scope(db, request.fence, cycle, request.scope) parent = await scope_owner(db, request.fence, cycle, request.scope) + attest_planned_owner(cycle, request, parent) inventory = bool(cycle.configuration.children_of(request.scope.record_type)) if request.fingerprint != cycle.configuration.fingerprint: raise CycleConflict("Scan configuration differs from the active cycle") @@ -229,12 +331,18 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: and not request.restart ): self._check_parent_epoch(row, parent) + if request.plan is not None and ( + not row.execution_state + or ScopeExecution.model_validate(row.execution_state).plan != request.plan + ): + raise ScanConflict("Active scope plan changed; explicit restart required") return await self._state(db, row) if request.expected is None: raise ScanConflict("Replacing a scan requires its current version") if row.cycle_id == request.cycle_id and not request.restart: raise ScanConflict("Changed scope configuration requires an explicit restart") row.revision += 1 + begin_execution(row, cycle, request, parent.visibility_epoch if parent else None) row.parent_visibility_epoch = parent.visibility_epoch if parent else None row.membership_attempt_id = request.fence.attempt_id if inventory else None row.cycle_id = request.cycle_id @@ -262,11 +370,12 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: """Capture and advance the page as one transaction, never a partial acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) - continuation = page_continuation(request, cycle.mode) await attest_scope(db, request.fence, cycle, request.scope) row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) + mode = effective_mode(row, cycle) + continuation = page_continuation(request, mode) parent = await scope_owner(db, request.fence, cycle, request.scope) if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): raise ScanConflict("Scope owner changed; restart from its current epoch") @@ -323,8 +432,9 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: row.continuation = continuation row.revision += 1 if request.final: - row.phase = "complete" if cycle.mode == "changes" else "reconciling" - row.completed_at = datetime.now(timezone.utc) if cycle.mode == "changes" else None + row.phase = "complete" if mode == "changes" else "reconciling" + row.completed_at = datetime.now(timezone.utc) if mode == "changes" else None + publish_scope(row, cycle) await db.flush() return ScanResult(state=await self._state(db, row), capture=captured) @@ -384,6 +494,8 @@ async def missing( raise ScanConflict("Changes scans cannot reconcile or validate enumeration absence") await attest_scope(db, fence, cycle, state.scope) row = self._expect(await self._row(db, fence, state.scope), state.version, state.cycle_id) + if effective_mode(row, cycle) == "changes": + raise ScanConflict("Changes scans cannot validate enumeration absence") parent = await scope_owner(db, fence, cycle, state.scope) if ( row.phase != "reconciling" @@ -473,6 +585,8 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul row = self._expect( await self._row(db, request.fence, request.scope), request.expected, request.cycle_id ) + if effective_mode(row, cycle) == "changes": + raise ScanConflict("Changes scans cannot reconcile enumeration absence") parent = await scope_owner(db, request.fence, cycle, request.scope) if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): raise ScanConflict("Scope owner changed; restart from its current epoch") @@ -493,6 +607,7 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul row.phase = "complete" row.completed_at = datetime.now(timezone.utc) row.revision += 1 + publish_scope(row, cycle) await db.flush() return ScanResult( state=await self._state(db, row), @@ -521,5 +636,6 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul if not result.has_more: row.phase = "complete" row.completed_at = datetime.now(timezone.utc) + publish_scope(row, cycle) await db.flush() return ScanResult(state=await self._state(db, row), capture=result.capture) diff --git a/backend/airweave/domains/entities/canonical/scope_execution.py b/backend/airweave/domains/entities/canonical/scope_execution.py new file mode 100644 index 0000000000..b148265faf --- /dev/null +++ b/backend/airweave/domains/entities/canonical/scope_execution.py @@ -0,0 +1,55 @@ +"""Immutable scope plans and successful evidence survive page-continuation replacement.""" + +import json +from typing import Literal +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, field_validator + +from airweave.domains.entities.canonical.cycle_models import CompletionPolicy, ProviderCheckpoint + + +class ScopePlan(BaseModel): + """Effective provider parameters, not original content or a second page cursor.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + mode: Literal["full", "changes"] = "full" + request_context: dict[str, JsonValue] = Field(default_factory=dict) + starting_checkpoint: ProviderCheckpoint | None = None + + @field_validator("request_context") + @classmethod + def bounded(cls, value): + """Bound immutable source parameters independently of provider progress.""" + if len(json.dumps(value, allow_nan=False).encode()) > 65536: + raise ValueError("Scope request context exceeds64KiB") + return value + + +class ScopeEvidence(BaseModel): + """Only SQL completion can issue evidence for this exact owner and request.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + sweep_id: UUID + completed_at: AwareDatetime + parent_visibility_epoch: int | None + request_context: dict[str, JsonValue] + policy: CompletionPolicy + checkpoint: ProviderCheckpoint | None = None + + def matches(self, plan: ScopePlan, epoch: int | None) -> bool: + """Revival or changed native request parameters cannot reuse old authority.""" + return ( + self.parent_visibility_epoch == epoch and self.request_context == plan.request_context + ) + + +class ScopeExecution(BaseModel): + """One current plan, with independent previously completed publication and baseline.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + schema_version: Literal[1] = 1 + plan: ScopePlan + last_full: ScopeEvidence | None = None + published: ScopeEvidence | None = None diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 76369741d3..757c9ce474 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -277,6 +277,7 @@ async def database(request): ): await connection.run_sync(seed_flat_scan, request.param) await connection.run_sync(migrate, "0007_scan_scope_owner.py") + await connection.run_sync(migrate, "0008_scope_execution.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() 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..7129471b0a --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py @@ -0,0 +1,497 @@ +"""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 +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" + (tmp_path / "mixed-coverage.json").write_text( + json.dumps( + [item.model_dump(mode="json") for item in (unknown, partial, complete)], 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 diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 817e288fd8..e434dd109c 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -479,6 +479,7 @@ async def test_capture_discovery_coverage_is_separate_from_retrieval( assert body["sources"][0]["capture"] == { "phase": "active", "mode": "full", + "scope_summary": None, "last_full_capture": None, "provider_checkpoint_promoted_at": None, "policies": {"event": "discovery_only"}, diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index e052758303..d475eeff47 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -18,6 +18,7 @@ from airweave.domains.entities.canonical.page_source import ( CanonicalPageSource, CheckpointedPageSource, + ScopedPageSource, ) from airweave.domains.entities.canonical.requests import ( CaptureBatch, @@ -73,6 +74,7 @@ def __init__( container_parents=self._container_parents, completion_policies=config.completion_policies, known_object_validation=config.known_object_validation, + scope_changes=config.scope_changes, ) if config != declared: raise ValueError("Page source cycle must match its declared container topology") @@ -138,7 +140,7 @@ async def progress(result: CaptureResult, records: tuple[CaptureRecord, ...]) -> force_full=( sync_context.force_full_sync or sync_context.execution_config.cursor.skip_load ) - if isinstance(self.page_source, CheckpointedPageSource) + if isinstance(self.page_source, (CheckpointedPageSource, ScopedPageSource)) else False, ).run() diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 9445873584..6bfc9a346a 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -16,8 +16,10 @@ CheckpointedPageSource, InvalidCaptureCheckpoint, InvalidScanContinuation, + InvalidScopeCheckpoint, KnownObjectSource, ScopeAccessLost, + ScopedPageSource, ScopeRemovalReason, ) from airweave.domains.entities.canonical.requests import ( @@ -33,6 +35,7 @@ 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 @@ -71,19 +74,27 @@ async def run(self) -> CaptureCycle: expected = None if ( self.force_full - and isinstance(self.source, CheckpointedPageSource) + and isinstance(self.source, (CheckpointedPageSource, ScopedPageSource)) and current is not None and current.phase == "active" - and (current.mode == "changes" or current.configuration != configuration) + 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) + 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) + if isinstance(self.source, (CheckpointedPageSource, ScopedPageSource)) else CapturePlan() ) ) @@ -96,13 +107,23 @@ async def run(self) -> CaptureCycle: 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): + if restart or not isinstance( + self.source, (CheckpointedPageSource, ScopedPageSource) + ): raise await self.check_limits() cycle = await self.fresh_full(cycle) @@ -110,10 +131,12 @@ async def run(self) -> CaptureCycle: 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): + 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": + 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) @@ -127,6 +150,8 @@ async def fresh_full(self, previous: CaptureCycle) -> CaptureCycle: 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", ), ) @@ -179,29 +204,7 @@ async def scan( parent_epoch: int | None = None, ) -> None: """Only a provider invalid-cursor response permits one explicit sweep restart.""" - 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 ( - (refresh_membership and previous.membership_attempt_id != self.fence.attempt_id) - or previous.parent_visibility_epoch != parent_epoch - ) - ) - async with self.sessions() as db: - state = await self.service.begin_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=self.initial_continuation(cycle), - ), - ) + state = await self.begin_scope(cycle, scope, parent, parent_epoch, refresh_membership) restarts = 0 while state.phase == "collecting": await self.check_limits() @@ -209,23 +212,11 @@ async def scan( page = await self.source.capture_page( scope, state.continuation, files=self.files, parent=parent ) - except InvalidScanContinuation: + except (InvalidScanContinuation, InvalidScopeCheckpoint) as error: if restarts >= 1: raise restarts += 1 - async with self.sessions() as db: - state = 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=self.initial_continuation(cycle), - ), - ) + 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: @@ -254,6 +245,115 @@ async def scan( 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, + refresh_membership: bool, + ) -> ScanState: + """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 ( + (refresh_membership 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 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), + ) + 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=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 if cycle.mode == "mixed" else None, + expected_parent_revision=parent.revision + if parent and cycle.mode == "mixed" + else None, + ), + ) + + 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 + if cycle.mode == "mixed" + else None, + expected_parent_revision=parent.revision + if parent and cycle.mode == "mixed" + 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 ( diff --git a/backend/airweave/models/capture_scan.py b/backend/airweave/models/capture_scan.py index fe57aab52c..c30dabb0ac 100644 --- a/backend/airweave/models/capture_scan.py +++ b/backend/airweave/models/capture_scan.py @@ -38,6 +38,7 @@ class CaptureScan(OrganizationBase): 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)) 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") From 323b8f45c768b035e9960eb06bc1f7a7d8383974 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:05:16 -0700 Subject: [PATCH 104/266] Record live Drive delta proof and migrate lifecycle harnesses --- backend/tests/live/canonical_capture.py | 1 + .../evidence/drive-page-cycles-20260930.json | 39 +++++++++++++++++++ backend/tests/live/provider_lifecycle.py | 1 + 3 files changed, 41 insertions(+) create mode 100644 backend/tests/live/evidence/drive-page-cycles-20260930.json diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index a7c42c2419..e1be77ea1e 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -301,6 +301,7 @@ async def main(): "0005_capture_scan.py", "0006_record_visibility.py", "0007_scan_scope_owner.py", + "0008_scope_execution.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: diff --git a/backend/tests/live/evidence/drive-page-cycles-20260930.json b/backend/tests/live/evidence/drive-page-cycles-20260930.json new file mode 100644 index 0000000000..203cb4433d --- /dev/null +++ b/backend/tests/live/evidence/drive-page-cycles-20260930.json @@ -0,0 +1,39 @@ +{ + "commit": "46172672bcc721fd8bd8b0761d3d28dd4c7d97e1", + "provider": "google_drive", + "verification": "real provider read-only; actual production capture pipeline; disposable local PostgreSQL; two fresh processes", + "first": { + "job_status": "completed", + "provider_requests": 140, + "records_observed": 55, + "stored_records": 55, + "partial_records": 3, + "blob_bytes_written": 207618914, + "observed_change_sequence": 55, + "loaded_durable_checkpoint": false + }, + "second": { + "job_status": "completed", + "provider_requests": 2, + "records_observed": 0, + "stored_records": 55, + "partial_records": 3, + "blob_bytes_written": 0, + "observed_change_sequence": 55, + "loaded_durable_checkpoint": true, + "resumed_changes_requests": 1, + "identical_payload_revision_digest": true + }, + "harness_exit_code": 0, + "independently_verified_cleanup": { + "remaining_private_schemas": 0, + "remaining_private_directories": 0, + "archive_removed": true + }, + "limitations": [ + "Three records remain content-incomplete", + "No live induced edit/deletion/revocation", + "No hosted Almanac authorization or consumer proof", + "No staging deployment" + ] +} diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index d3daa83cd4..676b62b7a7 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -836,6 +836,7 @@ async def main(): "0005_capture_scan.py", "0006_record_visibility.py", "0007_scan_scope_owner.py", + "0008_scope_execution.py", ): await connection.run_sync(harness.migrate, migration) organization_id, sync_id = uuid4(), uuid4() From 7bbcdf61a51bc2baab6eb5b833583e38e43ff54d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:12:10 -0700 Subject: [PATCH 105/266] Export SQL mixed coverage after completed change scopes --- .../canonical/tests/test_mixed_scopes.py | 48 ++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py index 7129471b0a..10e8681536 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py +++ b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py @@ -335,9 +335,55 @@ async def test_mixed_coverage_wire_from_sql(database, source, tmp_path): 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)], indent=2 + [ + item.model_dump(mode="json") + for item in (unknown, partial, complete, partial_delta, complete_delta) + ], + indent=2, ) ) From 6d9cabf6858ed92ec5898db95f4dcd7102fde08a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:32:31 -0700 Subject: [PATCH 106/266] Capture Calendar scopes durably and publish authorized observed ranges --- .../domains/entities/canonical/calendar.py | 10 + .../entities/canonical/calendar_query.py | 128 ++++-- .../domains/entities/canonical/page_source.py | 4 + .../domains/entities/canonical/requests.py | 27 +- .../domains/entities/canonical/scan_store.py | 10 +- .../entities/canonical/scope_execution.py | 2 + .../domains/entities/canonical/store.py | 19 +- .../tests/test_calendar_exclusions.py | 53 --- .../canonical/tests/test_calendar_range.py | 124 +++--- .../canonical/tests/test_calendar_recovery.py | 416 ++++++++++++++++++ .../canonical/tests/test_record_forest.py | 89 +++- .../domains/sync_pipeline/canonical_scan.py | 3 + .../platform/http_client/bounded_response.py | 22 + backend/airweave/platform/sources/gmail.py | 18 +- .../platform/sources/google_calendar.py | 196 ++++++--- .../sources/records/calendar_pages.py | 312 +++++++++++++ .../sources/records/google_calendar.py | 295 ++----------- .../records/test_calendar_transport.py | 137 ++++++ .../sources/records/test_google_calendar.py | 376 +++++----------- 19 files changed, 1456 insertions(+), 785 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py create mode 100644 backend/airweave/platform/http_client/bounded_response.py create mode 100644 backend/airweave/platform/sources/records/calendar_pages.py create mode 100644 backend/tests/unit/platform/sources/records/test_calendar_transport.py diff --git a/backend/airweave/domains/entities/canonical/calendar.py b/backend/airweave/domains/entities/canonical/calendar.py index d8ca49d279..849b7e9928 100644 --- a/backend/airweave/domains/entities/canonical/calendar.py +++ b/backend/airweave/domains/entities/canonical/calendar.py @@ -2,6 +2,8 @@ from collections.abc import Mapping +from pydantic import BaseModel, ConfigDict + def is_cancelled_recurring_event(payload: Mapping[str, object]) -> bool: """Recognize the provider-guaranteed identity of a cancelled series occurrence.""" @@ -17,3 +19,11 @@ def is_cancelled_recurring_event(payload: Mapping[str, object]) -> bool: isinstance(original.get(key), str) and original[key] for key in ("date", "dateTime") ) ) + + +class CalendarScopeContext(BaseModel): + """Effective native parameters and timezone used by the same range publication.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + parameters: dict[str, str | int] + timezone: str | None = None diff --git a/backend/airweave/domains/entities/canonical/calendar_query.py b/backend/airweave/domains/entities/canonical/calendar_query.py index abde4bd7cd..816f639d01 100644 --- a/backend/airweave/domains/entities/canonical/calendar_query.py +++ b/backend/airweave/domains/entities/canonical/calendar_query.py @@ -11,9 +11,11 @@ from sqlalchemy import String, and_, cast, func, select from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.checkpoint import CanonicalCheckpoint +from airweave.domains.entities.canonical.calendar import CalendarScopeContext +from airweave.domains.entities.canonical.cycle_models import CYCLE_KEY, CaptureCycle from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query import InvalidRecordCursor, RecordNotFound +from airweave.domains.entities.canonical.scope_execution import ScopeExecution from airweave.domains.entities.canonical.store import ( CanonicalStoreError, SourceNotFound, @@ -21,7 +23,8 @@ source_record, ) from airweave.models import Entity, Sync, SyncCursor, SyncJob -from airweave.platform.cursors.google_calendar import CalendarWindowCoverage, GoogleCalendarCursor +from airweave.models.capture_scan import CaptureScan +from airweave.platform.cursors.google_calendar import CalendarWindowCoverage class CalendarRangeError(CanonicalStoreError): @@ -158,63 +161,25 @@ async def read( ) if calendar is None: raise RecordNotFound("Calendar is not currently available") - stored_cursor = await db.scalar( - select(SyncCursor).where( - SyncCursor.sync_id == sync_id, SyncCursor.organization_id == organization_id - ) - ) - state = GoogleCalendarCursor.model_validate( - stored_cursor.cursor_data if stored_cursor else {} - ) - coverage = state.occurrence_coverage.get(calendar_id) - stamp_data = ( - stored_cursor.cursor_data.get("canonical_checkpoint") if stored_cursor else None - ) - if ( - coverage is None - or stamp_data is None - or query.start < coverage.start - or query.end > coverage.end - ): - raise CalendarRangeNotCaptured(query.start, query.end) - stamp = CanonicalCheckpoint.model_validate(stamp_data) - if calendar.source_payload.get("timeZone") != coverage.timezone: - raise CalendarChanged("Calendar timezone changed; refresh captured window") + coverage, refresh_state = await self._coverage(db, sync, calendar, query) sequence = sync.observed_change_sequence position = self._position(query, organization_id, sync_id, calendar_id, coverage, sequence) rows = await self._bounded_rows(db, organization_id, sync_id, calendar_id) matches = self._matching_rows(rows, coverage, query, position) # Recheck both journal and checkpoint: no snapshot or mixed-page claim. + calendar_record_id = calendar.id db.expire_all() latest = await db.scalar( select(Sync).where(Sync.id == sync_id, Sync.organization_id == organization_id) ) - latest_cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == sync_id)) - latest_scan = ( - latest_cursor.cursor_data.get("occurrence_coverage", {}) - .get(calendar_id, {}) - .get("scan_id") - if latest_cursor - else None + latest_calendar = await db.scalar( + select(Entity).where(Entity.id == calendar_record_id, content_is_available()) ) - if ( - latest is None - or latest.observed_change_sequence != sequence - or latest_scan != str(coverage.scan_id) - ): + if latest is None or latest_calendar is None or latest.observed_change_sequence != sequence: raise CalendarChanged("Calendar changed while reading; restart range query") - refresh_state = "last_completed" - if ( - latest.writer_attempt_id != stamp.writer_attempt_id - or sequence != stamp.observed_change_sequence - ): - job = await db.get(SyncJob, latest.writer_job_id) if latest.writer_job_id else None - refresh_state = ( - "refresh_in_progress" - if job - and str(getattr(job.status, "value", job.status)).lower() in {"running", "pending"} - else "refresh_incomplete" - ) + latest_coverage, refresh_state = await self._coverage(db, latest, latest_calendar, query) + if latest_coverage != coverage: + raise CalendarChanged("Calendar coverage changed while reading; restart range query") more = len(matches) > query.limit page = matches[: query.limit] next_cursor = None @@ -245,6 +210,73 @@ async def read( scanned_records=len(rows), ) + @staticmethod + async def _coverage( + db: AsyncSession, sync: Sync, calendar: Entity, query: CalendarRange + ) -> tuple[CalendarWindowCoverage, str]: + """Only completed exact-scope publication attests this native occurrence window.""" + cursor = await db.scalar( + select(SyncCursor).where( + SyncCursor.sync_id == sync.id, SyncCursor.organization_id == sync.organization_id + ) + ) + raw_cycle = cursor.cursor_data.get(CYCLE_KEY) if cursor else None + scan = await db.scalar( + select(CaptureScan).where( + CaptureScan.sync_id == sync.id, + CaptureScan.organization_id == sync.organization_id, + CaptureScan.record_type == "event_occurrence", + CaptureScan.container_id == calendar.native_id, + CaptureScan.parent_record_id == calendar.id, + CaptureScan.parent_visibility_epoch == calendar.visibility_epoch, + ) + ) + if not raw_cycle or scan is None or not scan.execution_state: + raise CalendarRangeNotCaptured(query.start, query.end) + cycle = CaptureCycle.model_validate(raw_cycle) + execution = ScopeExecution.model_validate(scan.execution_state) + published = execution.published + if ( + cycle.mode != "mixed" + or scan.fingerprint != cycle.configuration.fingerprint + or published is None + or published.parent_visibility_epoch != calendar.visibility_epoch + or published.policy != "exhaustive" + ): + raise CalendarRangeNotCaptured(query.start, query.end) + context = CalendarScopeContext.model_validate(published.request_context) + try: + coverage = CalendarWindowCoverage( + start=context.parameters["timeMin"], + end=context.parameters["timeMax"], + timezone=context.timezone, + completed_at=published.completed_at, + scan_id=published.sweep_id, + ) + except (KeyError, ValidationError) as error: + raise CalendarRangeNotCaptured(query.start, query.end) from error + if query.start < coverage.start or query.end > coverage.end: + raise CalendarRangeNotCaptured(query.start, query.end) + if calendar.source_payload.get("timeZone") != coverage.timezone: + raise CalendarChanged("Calendar timezone changed; refresh captured window") + refresh_state = "last_completed" + if published.writer_attempt_id is None or published.observed_change_sequence is None: + return coverage, "refresh_incomplete" + if ( + published.writer_attempt_id != sync.writer_attempt_id + or published.observed_change_sequence != sync.observed_change_sequence + or scan.cycle_id != cycle.version.cycle_id + or scan.phase != "complete" + or scan.sweep_id != published.sweep_id + ): + job = await db.get(SyncJob, sync.writer_job_id) if sync.writer_job_id else None + refresh_state = ( + "refresh_in_progress" + if job and job.status in {"running", "pending"} + else "refresh_incomplete" + ) + return coverage, refresh_state + @staticmethod async def _bounded_rows( db: AsyncSession, organization_id: UUID, sync_id: UUID, calendar_id: str diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index 023cdebeb7..8181134055 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -81,6 +81,10 @@ def __init__( 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.""" diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 5984e403ed..972ac36cf0 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -5,7 +5,15 @@ from typing import Annotated, Literal from uuid import UUID -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator +from pydantic import ( + AwareDatetime, + BaseModel, + ConfigDict, + Field, + JsonValue, + field_validator, + model_validator, +) ScopeRemovalReason = Literal["scope_removed", "access_revoked"] RecordKind = Annotated[str, Field(min_length=1, max_length=200)] @@ -71,6 +79,13 @@ class CaptureRecord(BaseModel): default=False, description="Audited provider evidence that this same object moved; never a retry override", ) + descendant_visibility_fields: tuple[ + Annotated[str, Field(min_length=1, max_length=200)], ... + ] = Field( + default=(), + max_length=8, + description="Audited native top-level fields whose change invalidates descendant access", + ) payload: dict[str, JsonValue] payload_schema_version: int = Field(default=1, ge=1) kind: Literal["upsert", "delete"] = "upsert" @@ -84,6 +99,16 @@ class CaptureRecord(BaseModel): observed_at: AwareDatetime blobs: tuple[BlobReference, ...] = () + @field_validator("descendant_visibility_fields") + @classmethod + def unique_visibility_fields(cls, value): + """Source declarations are bounded native names, never paths or expressions.""" + if len(set(value)) != len(value) or any( + not key.strip() or key != key.strip() for key in value + ): + raise ValueError("Descendant visibility fields require unique nonblank native names") + return value + @model_validator(mode="after") def validate_removal(self) -> "CaptureRecord": """Keep provider deletion distinct from inaccessible or deselected scope.""" diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 596d308859..b67eff7d39 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -131,7 +131,7 @@ def effective_mode(row: CaptureScan, cycle: CaptureCycle) -> str: return ScopeExecution.model_validate(row.execution_state).plan.mode -def publish_scope(row: CaptureScan, cycle: CaptureCycle) -> None: +def publish_scope(row: CaptureScan, cycle: CaptureCycle, sequence: int, attempt_id: UUID) -> None: """Completed evidence is derived only from this row's committed terminal state.""" if cycle.mode != "mixed" or row.phase != "complete": return @@ -144,6 +144,8 @@ def publish_scope(row: CaptureScan, cycle: CaptureCycle) -> None: parent_visibility_epoch=row.parent_visibility_epoch, request_context=execution.plan.request_context, policy=cycle.configuration.policy(row.record_type), + observed_change_sequence=sequence, + writer_attempt_id=attempt_id, checkpoint=ProviderCheckpoint.model_validate(checkpoint) if checkpoint else None, ) row.execution_state = execution.model_copy( @@ -434,7 +436,7 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: if request.final: row.phase = "complete" if mode == "changes" else "reconciling" row.completed_at = datetime.now(timezone.utc) if mode == "changes" else None - publish_scope(row, cycle) + publish_scope(row, cycle, sync.observed_change_sequence, request.fence.attempt_id) await db.flush() return ScanResult(state=await self._state(db, row), capture=captured) @@ -607,7 +609,7 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul row.phase = "complete" row.completed_at = datetime.now(timezone.utc) row.revision += 1 - publish_scope(row, cycle) + publish_scope(row, cycle, sync.observed_change_sequence, request.fence.attempt_id) await db.flush() return ScanResult( state=await self._state(db, row), @@ -636,6 +638,6 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul if not result.has_more: row.phase = "complete" row.completed_at = datetime.now(timezone.utc) - publish_scope(row, cycle) + publish_scope(row, cycle, sync.observed_change_sequence, request.fence.attempt_id) await db.flush() return ScanResult(state=await self._state(db, row), capture=result.capture) diff --git a/backend/airweave/domains/entities/canonical/scope_execution.py b/backend/airweave/domains/entities/canonical/scope_execution.py index b148265faf..4732fa762d 100644 --- a/backend/airweave/domains/entities/canonical/scope_execution.py +++ b/backend/airweave/domains/entities/canonical/scope_execution.py @@ -37,6 +37,8 @@ class ScopeEvidence(BaseModel): request_context: dict[str, JsonValue] policy: CompletionPolicy checkpoint: ProviderCheckpoint | None = None + observed_change_sequence: int | None = Field(default=None, ge=0) + writer_attempt_id: UUID | None = None def matches(self, plan: ScopePlan, epoch: int | None) -> bool: """Revival or changed native request parameters cannot reuse old authority.""" diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 0af5036e15..7b53a34a10 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -61,7 +61,9 @@ class WriterBusy(CanonicalStoreError): def capture_fingerprint(record: CaptureRecord) -> str: """Observation times never create changes; sparse tombstone payloads do.""" - material = record.model_dump(mode="json", exclude={"observed_at", "allow_reparent"}) + material = record.model_dump( + mode="json", exclude={"observed_at", "allow_reparent", "descendant_visibility_fields"} + ) serialized = json.dumps( material, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False ) @@ -395,10 +397,23 @@ async def _capture_locked( attestation_changed = ( entity is not None and entity.parent_visibility_epoch != parent_epoch ) + access_changed = ( + entity is not None + and bool(observation.descendant_visibility_fields) + and ( + entity.source_payload is None + or any( + (key in entity.source_payload, entity.source_payload.get(key)) + != (key in observation.payload, observation.payload.get(key)) + for key in observation.descendant_visibility_fields + ) + ) + ) if ( entity is not None and entity.capture_hash == fingerprint and not attestation_changed + and not access_changed and not (observation.kind == "upsert" and not was_available) ): if mark_seen: @@ -418,7 +433,7 @@ async def _capture_locked( ) db.add(entity) if observation.kind == "upsert" and ( - was_deleted or not was_available or parent_changed + was_deleted or not was_available or parent_changed or access_changed ): entity.visibility_epoch += 1 entity.parent_visibility_epoch = parent_epoch diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index 342dd4228e..302d46cde4 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -185,56 +185,3 @@ async def test_operational_occurrences_publish_zero_documents_without_retry(data for w in await projections.pending(db, fence.organization_id, fence.sync_id) ) assert (await db.get(Entity, row.id)).indexed_chunk_count == 0 - - -async def test_occurrence_access_loss_then_full_reacquisition_restores_original(database, source): - from airweave.domains.entities.canonical.tests.test_capture_pipeline import components - from airweave.domains.sources.exceptions import SourceEntityNotFoundError - from airweave.domains.sync_pipeline.canonical_capture import CanonicalCapturePipeline - from airweave.domains.sync_pipeline.capture_attempt import CaptureAttempt - from airweave.domains.syncs.cursors.cursor import SyncCursor - from airweave.platform.configs.config import CalendarOccurrenceWindow - from airweave.platform.sources.records.google_calendar import generate_calendar_observations - - capture, fence = source - state = SyncCursor(fence.sync_id) - window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") - native = {"id": "unchanged", "status": "confirmed", "summary": "Original"} - - async def run(restored, attempt): - async def get(url, params): - if url.endswith("calendarList"): - return {"items": [{"id": "cal", "timeZone": "UTC"}]} - if params["singleEvents"] == "true": - if not restored: - raise SourceEntityNotFoundError("synthetic access loss") - return {"items": []} - assert "syncToken" not in params - return {"items": [native], "nextSyncToken": "restored" if restored else "discard"} - - ctx, _, runtime, bus = components(database, source) - pipeline = CanonicalCapturePipeline( - capture, - database, - bus, - ("calendar", "event", "event_occurrence"), - CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), - {"event": "calendar", "event_occurrence": "calendar"}, - ) - await pipeline.start(ctx) - async for observation in generate_calendar_observations(get, state, window): - await pipeline.process([observation], ctx, runtime) - await pipeline.cleanup_orphaned_entities(ctx, runtime) - - await run(False, 1) - async with database() as db: - entity = await db.scalar(select(Entity).where(Entity.native_id == "unchanged")) - record_id = entity.id - lost = await capture.store.read(db, fence.organization_id, fence.sync_id, record_id) - assert lost.content_access == "unavailable" and lost.payload == {} - assert state.data["calendar_tokens"] == {} - await run(True, 2) - async with database() as db: - restored = await capture.store.read(db, fence.organization_id, fence.sync_id, record_id) - assert restored.content_access == "available" - assert restored.deleted_at is None and restored.payload == native diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py index 401888ff9f..9c377b066b 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_range.py @@ -1,6 +1,5 @@ """Stored provider instances support bounded ranges without recurrence guessing.""" -from datetime import datetime, timezone from types import SimpleNamespace from uuid import uuid4 @@ -23,7 +22,6 @@ from airweave.domains.entities.canonical.store import CanonicalStoreError from airweave.models import SyncCursor from airweave.models.source_connection import SourceConnection -from airweave.platform.cursors.google_calendar import CalendarWindowCoverage from airweave.platform.sources.records.google_calendar import record @@ -50,34 +48,31 @@ async def seed(database, source): }, {"id": "ends-at-boundary", "start": {"date": "2026-03-07"}, "end": {"date": "2026-03-08"}}, ] - observations = ( - parent, - *( - record("event_occurrence", item, "cal").model_copy(update={"parent": parent.identity}) - for item in events - ), + from airweave.domains.entities.canonical.tests.test_calendar_recovery import ( + NativeHTTP, + run, + setup, ) - async with database() as db: - await capture.capture(db, CaptureBatch(fence=fence, records=observations)) - coverage = CalendarWindowCoverage( - start="2026-03-01T00:00:00Z", - end="2026-04-01T00:00:00Z", - timezone="America/Los_Angeles", - completed_at=datetime.now(timezone.utc), - scan_id=uuid4(), + from airweave.platform.configs.config import GoogleCalendarConfig + + native = NativeHTTP( + [ + ("/users/me/calendarList", {"items": [parent.payload]}), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "seed"}), + ("/calendars/cal/events", {"items": events}), + ] ) + config = GoogleCalendarConfig( + occurrence_window={"start": "2026-03-01T00:00:00Z", "end": "2026-04-01T00:00:00Z"} + ) + pipeline, ctx, runtime, client = await setup(database, source, native, config=config) + await run(pipeline, ctx, runtime) + await client.aclose() async with database() as db: - await capture.save_checkpoint( - db, - fence, - { - "occurrence_coverage": {"cal": coverage.model_dump(mode="json")}, - "canonical_checkpoint": { - "writer_attempt_id": str(uuid4()), - "observed_change_sequence": 999999, - }, - }, + result = await CalendarRangeService("secret").read( + db, fence.organization_id, fence.sync_id, "cal", window() ) + coverage = result.coverage return parent, coverage @@ -257,12 +252,12 @@ async def test_budget_rejects_oversized_row_before_orm_materialization(database, async def test_empty_selection_reconciles_partial_parent_without_saved_cursor(database, source): - from airweave.domains.entities.canonical.requests import ( - CompletedScope, - ReconcileScope, - StartedScope, + from airweave.domains.entities.canonical.tests.test_calendar_recovery import ( + NativeHTTP, + run, + setup, ) - from airweave.platform.sources.records.google_calendar import generate_calendar_observations + from airweave.platform.configs.config import GoogleCalendarConfig capture, fence = source parent = record("calendar", {"id": "partial", "timeZone": "UTC"}) @@ -272,36 +267,26 @@ async def test_empty_selection_reconciles_partial_parent_without_saved_cursor(da async with database() as db: result = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, child))) child_id = result.changes[-1].record.id - - async def never_get(*args, **kwargs): - raise AssertionError("Empty selection must not fetch provider") - - # Same-attempt retry is intentionally harder: StartedScope must clear old sightings. - async for item in generate_calendar_observations(never_get, None, calendar_ids=()): - if isinstance(item, StartedScope): - async with database() as db: - await capture.start_scope(db, fence, item) - elif type(item) is CompletedScope: - async with database() as db: - await capture.reconcile_scope( - db, - ReconcileScope(fence=fence, scope=item, observed_at=datetime.now(timezone.utc)), - ) + native = NativeHTTP([]) + pipeline, ctx, runtime, client = await setup( + database, source, native, config=GoogleCalendarConfig(calendar_ids=()) + ) + await run(pipeline, ctx, runtime) async with database() as db: current = await capture.store.read(db, fence.organization_id, fence.sync_id, child_id) assert current.content_access == "unavailable" and current.payload == {} - await db.rollback() - removed = await capture.reconcile_parents(db, fence) - assert removed.capture.changes[0].record.removal_reason != "provider_deleted" + assert current.removal_reason != "provider_deleted" + assert not native.calls + await client.aclose() async def test_missing_selected_calendar_revokes_partial_parent_without_cursor(database, source): - from airweave.domains.entities.canonical.requests import ( - CaptureRecord, - RemovedScope, - StartedScope, + from airweave.domains.entities.canonical.tests.test_calendar_recovery import ( + NativeHTTP, + run, + setup, ) - from airweave.platform.sources.records.google_calendar import generate_calendar_observations + from airweave.platform.configs.config import GoogleCalendarConfig capture, fence = source parent = record("calendar", {"id": "missing", "timeZone": "UTC"}) @@ -311,27 +296,18 @@ async def test_missing_selected_calendar_revokes_partial_parent_without_cursor(d async with database() as db: initial = await capture.capture(db, CaptureBatch(fence=fence, records=(parent, child))) child_id = initial.changes[-1].record.id - - async def empty_membership(*args, **kwargs): - return {"items": []} - + native = NativeHTTP( + [("/users/me/calendarList", {"items": []}), ("/users/me/calendarList/missing", (404, {}))] + ) + pipeline, ctx, runtime, client = await setup( + database, source, native, config=GoogleCalendarConfig(calendar_ids=("missing",)) + ) with pytest.raises(ValueError, match="not accessible"): - async for item in generate_calendar_observations( - empty_membership, None, calendar_ids=("missing",) - ): - async with database() as db: - if isinstance(item, CaptureRecord): - await capture.capture(db, CaptureBatch(fence=fence, records=(item,))) - elif isinstance(item, StartedScope): - await capture.start_scope(db, fence, item) - elif isinstance(item, RemovedScope): - await capture.remove_scope(db, fence, item) - else: - raise AssertionError("Missing requested scope cannot complete") + await run(pipeline, ctx, runtime) async with database() as db: current = await capture.store.read(db, fence.organization_id, fence.sync_id, child_id) assert current.content_access == "unavailable" and current.payload == {} - assert current.removal_reason == "access_revoked" - assert ( - await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) is None - ) + # Child cleanup may follow a later retry; recursive parent visibility hides it now. + cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + assert cursor.cursor_data["canonical_cycle"]["phase"] == "active" + await client.aclose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py new file mode 100644 index 0000000000..0cb932a9ca --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py @@ -0,0 +1,416 @@ +"""Real Calendar source and SQL page lifecycle; only provider HTTP is synthetic.""" + +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.calendar_query import CalendarRange, CalendarRangeService +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components +from airweave.domains.entities.canonical.tests.test_mixed_scopes import next_job +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.platform.configs.config import GoogleCalendarConfig +from airweave.platform.sources.google_calendar import GoogleCalendarSource +from airweave.platform.sources.records.google_calendar import BASE + +CONFIG = GoogleCalendarConfig( + occurrence_window={"start": "2026-03-01T00:00:00Z", "end": "2026-04-01T00:00:00Z"} +) +MEMBERS = {"items": [{"id": "cal", "timeZone": "UTC"}]} + + +def event(native_id): + return { + "id": native_id, + "start": {"dateTime": "2026-03-08T09:00:00Z"}, + "end": {"dateTime": "2026-03-08T10:00:00Z"}, + } + + +class NativeHTTP: + def __init__(self, replies): + self.replies, self.calls = list(replies), [] + + async def handle(self, request): + self.calls.append((request.url.path, dict(request.url.params))) + path, value = self.replies.pop(0) + assert str(request.url).split("?")[0] == BASE + path + if isinstance(value, Exception): + raise value + status, payload = value if isinstance(value, tuple) else (200, value) + return httpx.Response(status, json=payload) + + +async def setup(database, source, native, attempt=1, config=CONFIG): + service, fence = source + client = httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) + connector = await GoogleCalendarSource.create( + auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=client, config=config + ) + 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_native_calendar_resume_delta_and_published_range(database, source): + # Pre-page-engine tokens cannot authorize a delta without exact scope evidence. + async with database() as db: + await source[0].save_checkpoint(db, source[1], {"calendar_tokens": {"cal": "legacy"}}) + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("a")], "nextPageToken": "p2"}), + ("/calendars/cal/events", ConnectionError("interrupted")), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("b")], "nextSyncToken": "full-token"}), + ("/calendars/cal/events", {"items": [event("occurrence")]}), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "delta-token"}), + ("/calendars/cal/events", {"items": [event("occurrence")]}), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + with pytest.raises(ConnectionError): + await run(pipeline, ctx, runtime) + await client.aclose() + pipeline, ctx, runtime, client = await setup(database, source, native, 2) + await run(pipeline, ctx, runtime) + assert native.calls[4][1]["pageToken"] == "p2" + async with database() as db: + result = await CalendarRangeService("secret").read( + db, + source[1].organization_id, + source[1].sync_id, + "cal", + CalendarRange(start="2026-03-08T00:00:00Z", end="2026-03-09T00:00:00Z", timezone="UTC"), + ) + assert result.records[0].identity.native_id == "occurrence" + assert result.refresh_state == "last_completed" + await client.aclose() + source = await next_job(database, source) + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + assert native.calls[7][1]["syncToken"] == "full-token" + async with database() as db: + masters = ( + await db.scalars( + select(Entity).where( + Entity.entity_definition_short_name == "event", Entity.deleted_at.is_(None) + ) + ) + ).all() + assert {row.native_id for row in masters} == {"a", "b"} + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + assert scan.execution_state["last_full"]["checkpoint"]["value"] == { + "sync_token": "full-token" + } + assert scan.execution_state["published"]["checkpoint"]["value"] == { + "sync_token": "delta-token" + } + assert not native.replies + await client.aclose() + + +async def test_native_410_restarts_only_one_master_scope(database, source): + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("kept")], "nextSyncToken": "old"}), + ("/calendars/cal/events", {"items": []}), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("delta-only")], "nextPageToken": "delta2"}), + ("/calendars/cal/events", (410, {"error": {"message": "Expired"}})), + ("/calendars/cal/events", {"items": [event("kept")], "nextSyncToken": "fresh"}), + ("/calendars/cal/events", {"items": []}), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + await client.aclose() + source = await next_job(database, source) + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + assert native.calls[5][1]["syncToken"] == "old" + assert "syncToken" not in native.calls[6][1] and "pageToken" not in native.calls[6][1] + async with database() as db: + stale = await db.scalar(select(Entity).where(Entity.native_id == "delta-only")) + assert stale.deleted_at is not None and stale.removal_reason == "absent" + await client.aclose() + + +async def test_config_change_selects_full_instead_of_rejecting_forever(database, source): + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "old"}), + ("/calendars/cal/events", {"items": []}), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "new"}), + ("/calendars/cal/events", {"items": []}), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + await client.aclose() + source = await next_job(database, source) + changed = CONFIG.model_copy(update={"calendar_ids": ("cal",)}) + pipeline, ctx, runtime, client = await setup(database, source, native, config=changed) + await run(pipeline, ctx, runtime) + assert "syncToken" not in native.calls[4][1] + await client.aclose() + + +@pytest.mark.parametrize("required", [False, True]) +async def test_occurrence_loss_withdraws_then_restoration_reacquires_unchanged_master( + database, source, required +): + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("unchanged")], "nextSyncToken": "discard"}), + ("/calendars/cal/events", (404, {})), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("unchanged")], "nextSyncToken": "restored"}), + ("/calendars/cal/events", {"items": []}), + ] + ) + from airweave.domains.entities.canonical.page_source import RequiredScopeAccessLost + + config = CONFIG.model_copy(update={"calendar_ids": ("cal",)}) if required else CONFIG + pipeline, ctx, runtime, client = await setup(database, source, native, config=config) + if required: + with pytest.raises(RequiredScopeAccessLost): + await run(pipeline, ctx, runtime) + else: + await run(pipeline, ctx, runtime) + async with database() as db: + entity = await db.scalar(select(Entity).where(Entity.native_id == "unchanged")) + record_id = entity.id + result = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, record_id + ) + assert result.content_access == "unavailable" and result.payload == {} + await client.aclose() + if not required: + source = await next_job(database, source) + pipeline, ctx, runtime, client = await setup( + database, source, native, 2 if required else 1, config + ) + await run(pipeline, ctx, runtime) + assert "syncToken" not in native.calls[4][1] + async with database() as db: + result = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, record_id + ) + assert result.content_access == "available" and result.payload == event("unchanged") + await client.aclose() + + +async def test_missing_selection_withdrawal_survives_interruption_then_restores(database, source): + native = NativeHTTP( + [ + ("/users/me/calendarList", {"items": []}), + ("/users/me/calendarList/cal", (404, {})), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("fresh")], "nextSyncToken": "new"}), + ("/calendars/cal/events", {"items": []}), + ] + ) + config = CONFIG.model_copy(update={"calendar_ids": ("cal",)}) + pipeline, ctx, runtime, client = await setup(database, source, native, config=config) + + async def stop_after_committed_page(): + async with database() as db: + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "calendar")) + if scan and scan.continuation.get("selection_unavailable"): + raise InterruptedError("Synthetic postcommit interruption") + + with pytest.raises(InterruptedError): + await pipeline.run_scans(ctx, runtime, stop_after_committed_page) + async with database() as db: + root = await db.scalar(select(Entity).where(Entity.native_id == "cal")) + assert root.deleted_at is not None and root.removal_reason == "scope_removed" + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "calendar")) + assert scan.phase == "collecting" and scan.continuation["selection_unavailable"] + await client.aclose() + pipeline, ctx, runtime, client = await setup(database, source, native, 2, config) + await run(pipeline, ctx, runtime) + assert "syncToken" not in native.calls[3][1] + assert not native.replies + await client.aclose() + + +async def test_failed_refresh_keeps_publication_and_new_writer_is_not_fresh(database, source): + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "old"}), + ("/calendars/cal/events", {"items": [event("old-occurrence")]}), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [], "nextSyncToken": "new"}), + ("/calendars/cal/events", {"items": [event("new-occurrence")], "nextPageToken": "p2"}), + ("/calendars/cal/events", ConnectionError("interrupted")), + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": []}), + ] + ) + query = CalendarRange(start="2026-03-08T00:00:00Z", end="2026-03-09T00:00:00Z", timezone="UTC") + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + async with database() as db: + scan = await db.scalar( + select(CaptureScan).where(CaptureScan.record_type == "event_occurrence") + ) + published = scan.execution_state["published"] + await client.aclose() + source = await next_job(database, source) + async with database() as db: + result = await CalendarRangeService("secret").read( + db, source[1].organization_id, source[1].sync_id, "cal", query + ) + assert result.refresh_state == "refresh_in_progress" + pipeline, ctx, runtime, client = await setup(database, source, native) + with pytest.raises(ConnectionError): + await run(pipeline, ctx, runtime) + async with database() as db: + scan = await db.scalar( + select(CaptureScan).where(CaptureScan.record_type == "event_occurrence") + ) + assert scan.execution_state["published"] == published + result = await CalendarRangeService("secret").read( + db, source[1].organization_id, source[1].sync_id, "cal", query + ) + assert result.refresh_state == "refresh_in_progress" + assert {item.identity.native_id for item in result.records} == { + "old-occurrence", + "new-occurrence", + } + await client.aclose() + pipeline, ctx, runtime, client = await setup(database, source, native, 2) + await run(pipeline, ctx, runtime) + async with database() as db: + result = await CalendarRangeService("secret").read( + db, source[1].organization_id, source[1].sync_id, "cal", query + ) + assert result.refresh_state == "last_completed" + assert {item.identity.native_id for item in result.records} == {"new-occurrence"} + scan = await db.scalar( + select(CaptureScan).where(CaptureScan.record_type == "event_occurrence") + ) + assert ( + scan.execution_state["published"]["observed_change_sequence"] + > published["observed_change_sequence"] + ) + await client.aclose() + + +async def test_membership_role_downgrade_hides_prior_event_details(database, source): + config = CONFIG.model_copy(update={"calendar_ids": ("cal",)}) + native = NativeHTTP( + [ + ("/users/me/calendarList", MEMBERS), + ("/calendars/cal/events", {"items": [event("private-detail")], "nextSyncToken": "old"}), + ("/calendars/cal/events", {"items": []}), + ( + "/users/me/calendarList", + {"items": [{"id": "cal", "timeZone": "UTC", "accessRole": "freeBusyReader"}]}, + ), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native, config=config) + await run(pipeline, ctx, runtime) + await client.aclose() + source = await next_job(database, source) + pipeline, ctx, runtime, client = await setup(database, source, native, config=config) + with pytest.raises(ValueError, match="not accessible"): + await run(pipeline, ctx, runtime) + async with database() as db: + original = await db.scalar(select(Entity).where(Entity.native_id == "private-detail")) + result = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, original.id + ) + assert result.content_access == "unavailable" and result.payload == {} + parent = await db.scalar(select(Entity).where(Entity.native_id == "cal")) + assert parent.removal_reason == "access_revoked" + assert not native.replies and len(native.calls) == 4 + await client.aclose() + + +async def test_reader_role_change_hides_children_until_same_identity_recaptured(database, source): + owner = {"items": [{"id": "cal", "timeZone": "UTC", "accessRole": "owner"}]} + reader = {"items": [{"id": "cal", "timeZone": "UTC", "accessRole": "reader"}]} + secret = {**event("private"), "visibility": "private", "description": "Formerly visible"} + redacted = {**event("private"), "visibility": "private"} + native = NativeHTTP( + [ + ("/users/me/calendarList", owner), + ("/calendars/cal/events", {"items": [secret], "nextSyncToken": "owner-token"}), + ("/calendars/cal/events", {"items": []}), + ("/users/me/calendarList", reader), + ("/users/me/calendarList", reader), + ("/calendars/cal/events", {"items": [redacted], "nextSyncToken": "reader-token"}), + ("/calendars/cal/events", {"items": []}), + ] + ) + pipeline, ctx, runtime, client = await setup(database, source, native) + await run(pipeline, ctx, runtime) + async with database() as db: + child = await db.scalar(select(Entity).where(Entity.native_id == "private")) + child_id = child.id + parent = await db.scalar(select(Entity).where(Entity.native_id == "cal")) + old_epoch = parent.visibility_epoch + await client.aclose() + source = await next_job(database, source) + pipeline, ctx, runtime, client = await setup(database, source, native) + + async def stop_after_role_commit(): + async with database() as db: + parent = await db.scalar(select(Entity).where(Entity.native_id == "cal")) + if parent.visibility_epoch != old_epoch: + raise InterruptedError("Synthetic role-change commit boundary") + + with pytest.raises(InterruptedError): + await pipeline.run_scans(ctx, runtime, stop_after_role_commit) + async with database() as db: + hidden = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, child_id + ) + assert hidden.content_access == "unavailable" and hidden.payload == {} + await client.aclose() + pipeline, ctx, runtime, client = await setup(database, source, native, 2) + await run(pipeline, ctx, runtime) + assert "syncToken" not in native.calls[5][1] + async with database() as db: + current = await source[0].store.read( + db, source[1].organization_id, source[1].sync_id, child_id + ) + assert current.content_access == "available" and current.payload == redacted + parent = await db.scalar(select(Entity).where(Entity.native_id == "cal")) + assert parent.visibility_epoch == old_epoch + 1 + await client.aclose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_record_forest.py b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py index 26b9d5f0a8..2a7c0e1d06 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_record_forest.py +++ b/backend/airweave/domains/entities/canonical/tests/test_record_forest.py @@ -225,8 +225,9 @@ async def test_real_recursive_query_plan(shape, size, database, source): ) +@pytest.mark.parametrize("withdrawal", ["delete", "access_change"]) async def test_nested_withdrawal_gates_original_blob_history_search_and_publication( - database, source + database, source, withdrawal ): from types import SimpleNamespace from unittest.mock import AsyncMock @@ -246,7 +247,9 @@ async def test_nested_withdrawal_gates_original_blob_history_search_and_publicat from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata service, fence = source - root = node("root") + 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),)} @@ -305,18 +308,27 @@ async def test_nested_withdrawal_gates_original_blob_history_search_and_publicat database, service, fence, - root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}), + 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 - assert not await query.queries.list_records( + 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) + 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 @@ -324,3 +336,70 @@ async def test_nested_withdrawal_gates_original_blob_history_search_and_publicat 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/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 6bfc9a346a..c0409ef7fb 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -18,6 +18,7 @@ InvalidScanContinuation, InvalidScopeCheckpoint, KnownObjectSource, + RequiredScopeAccessLost, ScopeAccessLost, ScopedPageSource, ScopeRemovalReason, @@ -188,6 +189,8 @@ async def run_cycle(self, cycle: CaptureCycle) -> CaptureCycle: await self.withdraw_parent( work.parent, work.parent_visibility_epoch, error.removal_reason ) + if isinstance(error, RequiredScopeAccessLost): + raise async with self.sessions() as db: current = await self.service.read_cycle(db, self.fence) if current is None: 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/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index c9777a6da9..50e861bea3 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -10,7 +10,6 @@ import base64 import hashlib import json -from contextlib import aclosing from datetime import datetime from typing import Any, AsyncGenerator, Dict, List, Optional, Set @@ -57,6 +56,7 @@ GmailThreadEntity, ) 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.retry_helpers import ( wait_rate_limit_with_backoff, ) @@ -439,21 +439,7 @@ async def _get_mime_json( @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.""" - if response.is_stream_consumed: - if len(response.content) > maximum: - raise FileSkippedException( - "Gmail JSON response exceeds size limit", "Gmail response" - ) - 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( - "Gmail JSON response exceeds size limit", "Gmail response" - ) - body.extend(chunk) - return bytes(body) + return await bounded_response_bytes(response, maximum, label="Gmail response") # ----------------------- # Cursor helper diff --git a/backend/airweave/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index 447ae58aaa..52737608e4 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -1,48 +1,42 @@ -"""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 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.entities.canonical.source import SourceObservation +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, @@ -60,16 +54,32 @@ GoogleCalendarListEntity, ) 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.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.google_calendar import generate_calendar_observations +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) + + @source( name="Google Calendar", short_name="google_calendar", @@ -99,23 +109,61 @@ class GoogleCalendarSource(BaseSource): canonical_record_types = ("calendar", "event", "event_occurrence") canonical_container_parents = {"event": "calendar", "event_occurrence": "calendar"} - async def generate_observations( + @property + def capture_cycle_configuration(self) -> CycleConfiguration: + """Calendar membership owns exact native event and occurrence scopes.""" + return CalendarPages(self._get_capture_json, self.calendar_config).configuration + + async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: + """Resolve the expansion window once per persisted cycle.""" + return CalendarPages(self._get_capture_json, self.calendar_config).prepare() + + async def prepare_scope( self, + scope: CompletedScope, + cycle: CaptureCycle, + previous: ScanState | None, *, - cursor: SyncCursor | None = None, - files: FileService | None = None, - node_selections: list[NodeSelectionData] | None = None, - ) -> AsyncGenerator[SourceObservation, None]: - """Capture originals and bounded provider-expanded instances with real IDs.""" - if node_selections: - raise ValueError("Calendar selection scopes are not supported by canonical capture") - async for observation in generate_calendar_observations( - self._get, - cursor, - self.calendar_config.resolved_window(), - calendar_ids=self.calendar_config.calendar_ids, - ): - yield observation + parent: SourceRecord | None, + force_full: bool, + ) -> ScopePlan: + """Choose native delta only from compatible completed exact-scope evidence.""" + return CalendarPages(self._get_capture_json, self.calendar_config).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.""" + 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 CalendarPages(self._get_capture_json, self.calendar_config).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 CalendarPages(self._get_capture_json, self.calendar_config).confirm_member_absent( + record.identity.native_id + ) # ----------------------- # Construction / Config @@ -154,6 +202,56 @@ async def _refresh_and_get_headers(self) -> Dict[str, str]: """Force-refresh the token and return updated headers.""" 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), retry=retry_if_rate_limit_or_timeout, 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..7ad52ead89 --- /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": 2, "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 index 61aca399db..7651f6f31a 100644 --- a/backend/airweave/platform/sources/records/google_calendar.py +++ b/backend/airweave/platform/sources/records/google_calendar.py @@ -1,31 +1,12 @@ -"""Original Calendar resources with per-calendar incremental synchronization.""" +"""Native Calendar resource validation; page progress belongs to the canonical engine.""" -import json -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator -from contextlib import aclosing +from collections.abc import Awaitable, Callable from datetime import datetime, timezone -from typing import Literal -from urllib.parse import quote -from uuid import uuid4 -from pydantic import BaseModel, ConfigDict, Field, JsonValue +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, - CompletedScope, - RecordIdentity, - RemovedScope, - StartedScope, -) -from airweave.domains.entities.canonical.source import SourceObservation -from airweave.domains.sources.exceptions import ( - SourceEntityNotFoundError, - SourceGoneError, -) -from airweave.domains.syncs.cursors.cursor import SyncCursor -from airweave.platform.configs.config import CalendarOccurrenceWindow -from airweave.platform.cursors.google_calendar import CalendarWindowCoverage, GoogleCalendarCursor +from airweave.domains.entities.canonical.requests import CaptureRecord, RecordIdentity GetJSON = Callable[..., Awaitable[dict]] BASE = "https://www.googleapis.com/calendar/v3" @@ -63,252 +44,34 @@ def record( 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. - deleted = (payload.get("status") == "cancelled" and not recurring_cancellation) or payload.get( - "deleted" - ) is True + 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="provider_deleted" if deleted else None, - observed_at=datetime.now(timezone.utc), - ) - - -async def capture_calendar( - get: GetJSON, calendar_id: str, token: str | None, tokens: dict[str, str] -) -> AsyncGenerator[SourceObservation, None]: - """A 410 resets only this calendar, including sightings from earlier delta pages.""" - full = token is None - if full: - yield StartedScope(record_type="event", container_id=calendar_id) - params: dict[str, str | int] = { - "maxResults": 2500, - "showDeleted": "true", - "singleEvents": "false", - } - if token: - params["syncToken"] = token - seen_pages: set[str] = set() - while True: - try: - page = CalendarPage.model_validate( - await get(f"{BASE}/calendars/{quote(calendar_id, safe='')}/events", params=params) - ) - except SourceGoneError: - if full: - raise - full = True - params.pop("syncToken", None) - params.pop("pageToken", None) - seen_pages.clear() - yield StartedScope(record_type="event", container_id=calendar_id) - continue - for item in page.items: - yield record("event", item, calendar_id) - if not page.nextPageToken: - tokens[calendar_id] = page.completed_token() - if full: - yield CompletedScope(record_type="event", container_id=calendar_id) - return - if page.nextPageToken in seen_pages: - raise ValueError("Calendar returned a repeated page token") - seen_pages.add(page.nextPageToken) - params["pageToken"] = page.nextPageToken - - -async def generate_calendar_observations( - get: GetJSON, - cursor: SyncCursor | None, - window: CalendarOccurrenceWindow | None = None, - *, - calendar_ids: tuple[str, ...] | None = None, -) -> AsyncGenerator[SourceObservation, None]: - """Refresh calendar membership; capture each accessible calendar independently.""" - previous = GoogleCalendarCursor.model_validate(cursor.data if cursor else {}) - selection = set(calendar_ids) if calendar_ids is not None else None - withdrawn = previous.calendar_tokens.keys() - selection if selection is not None else set() - for calendar_id in withdrawn: - for observation in removed_calendar(calendar_id, "scope_removed"): - yield observation - tokens: dict[str, str] = {} - coverage: dict[str, CalendarWindowCoverage] = {} - seen_calendars: set[str] = set() - yield StartedScope(record_type="calendar") - async with aclosing(selected_members(get, selection)) as members: - async for calendar_record in members: - seen_calendars.add(calendar_record.identity.native_id) - yield calendar_record - async for observation in _capture_member( - get, calendar_record, calendar_record.payload, previous, window, tokens, coverage - ): - yield observation - for observation in missing_members( - set(previous.calendar_tokens), seen_calendars, withdrawn, selection - ): - yield observation - yield CompletedScope(record_type="calendar") - if cursor is not None: - cursor.update(calendar_tokens=tokens) - if window is not None: - cursor.update(occurrence_coverage=coverage) - - -async def capture_occurrences( - get: GetJSON, calendar_id: str, window: CalendarOccurrenceWindow -) -> AsyncGenerator[SourceObservation, None]: - """Only a fully exhausted fixed horizon may reconcile this distinct scope.""" - yield StartedScope(record_type="event_occurrence", container_id=calendar_id) - params = { - "singleEvents": "true", - "showDeleted": "true", - "maxResults": 2500, - "timeMin": window.start.isoformat(), - "timeMax": window.end.isoformat(), - } - pages: set[str] = set() - identities: set[str] = set() - byte_count = 0 - while True: - page = CalendarPage.model_validate( - await get(f"{BASE}/calendars/{quote(calendar_id, safe='')}/events", params=params) + removal_reason=( + "access_revoked" + if unavailable_role + else "scope_removed" + if kind == "calendar" + else "provider_deleted" ) - for item in page.items: - byte_count += len(json.dumps(item).encode()) - observed = record("event_occurrence", item, calendar_id) - if observed.identity.native_id in identities: - raise ValueError("Expanded Calendar listing repeated an occurrence identity") - identities.add(observed.identity.native_id) - if len(identities) > 10000 or byte_count > 20 * 1024 * 1024: - raise ValueError("Expanded Calendar capture exceeds bounded horizon budget") - yield observed - if not page.nextPageToken: - yield CompletedScope(record_type="event_occurrence", container_id=calendar_id) - return - if page.nextPageToken in pages: - raise ValueError("Expanded Calendar listing repeated a page token") - pages.add(page.nextPageToken) - if len(pages) >= 100: - raise ValueError("Expanded Calendar capture exceeds page budget") - params["pageToken"] = page.nextPageToken - - -async def _capture_member( - get: GetJSON, - calendar_record: CaptureRecord, - item: dict[str, JsonValue], - previous: GoogleCalendarCursor, - window: CalendarOccurrenceWindow | None, - tokens: dict[str, str], - coverage: dict[str, CalendarWindowCoverage], -) -> AsyncGenerator[SourceObservation, None]: - """Capture one member or record explicit access loss without claiming its window.""" - calendar_id = calendar_record.identity.native_id - try: - async for observation in capture_calendar( - get, calendar_id, previous.calendar_tokens.get(calendar_id), tokens - ): - yield observation - if window is not None: - async for observation in capture_occurrences(get, calendar_id, window): - yield observation - calendar_timezone = item.get("timeZone") - if not isinstance(calendar_timezone, str) or not calendar_timezone: - raise ValueError("Calendar has no timezone for occurrence coverage") - coverage[calendar_id] = CalendarWindowCoverage( - start=window.start, - end=window.end, - timezone=calendar_timezone, - completed_at=datetime.now(timezone.utc), - scan_id=uuid4(), - ) - except SourceEntityNotFoundError: - # Master enumeration can finish before expanded instances lose access. - # Its token cannot survive withdrawing those originals: restoration must - # reacquire unchanged masters, which an incremental response may omit. - tokens.pop(calendar_id, None) - coverage.pop(calendar_id, None) - yield CaptureRecord( - identity=RecordIdentity(record_type="calendar", native_id=calendar_id), - payload={"id": calendar_id}, - kind="delete", - removal_reason="access_revoked", - observed_at=datetime.now(timezone.utc), - ) - yield RemovedScope( - record_type="event_occurrence", - container_id=calendar_id, - removal_reason="access_revoked", - observed_at=datetime.now(timezone.utc), - ) - yield RemovedScope( - record_type="event", - container_id=calendar_id, - removal_reason="access_revoked", - observed_at=datetime.now(timezone.utc), - ) - - -def removed_calendar( - calendar_id: str, reason: Literal["scope_removed", "access_revoked"] -) -> Iterator[SourceObservation]: - """Withdraw parent visibility and exact child scopes without alleging provider deletion.""" - now = datetime.now(timezone.utc) - yield CaptureRecord( - identity=RecordIdentity(record_type="calendar", native_id=calendar_id), - payload={"id": calendar_id}, - kind="delete", - removal_reason=reason, - observed_at=now, + 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), ) - for record_type in ("event_occurrence", "event"): - yield RemovedScope( - record_type=record_type, - container_id=calendar_id, - removal_reason=reason, - observed_at=now, - ) - - -async def selected_members( - get: GetJSON, selection: set[str] | None -) -> AsyncGenerator[CaptureRecord, None]: - """Enumerate membership completely, but fetch content only for configured native IDs.""" - seen_pages: set[str] = set() - seen_calendars: set[str] = set() - params: dict[str, str | int] = {"maxResults": 250, "showHidden": "true"} - while selection != set(): - page = CalendarPage.model_validate( - await get(f"{BASE}/users/me/calendarList", params=params) - ) - for item in page.items: - member = record("calendar", item) - calendar_id = member.identity.native_id - if calendar_id in seen_calendars: - raise ValueError("Calendar list repeated an identity across pages") - seen_calendars.add(calendar_id) - if selection is None or calendar_id in selection: - yield member - if not page.nextPageToken: - return - if page.nextPageToken in seen_pages: - raise ValueError("Calendar list returned a repeated page token") - seen_pages.add(page.nextPageToken) - params["pageToken"] = page.nextPageToken - - -def missing_members( - previous: set[str], seen_calendars: set[str], withdrawn: set[str], selection: set[str] | None -) -> Iterator[SourceObservation]: - """Revoke membership-proven absence independently of saved cursor history.""" - missing_requested = selection - seen_calendars if selection is not None else set() - for missing in previous - seen_calendars - withdrawn - missing_requested: - for observation in removed_calendar(missing, "scope_removed"): - yield observation - for missing in missing_requested: - # Full membership exhaustion proves absence even when an earlier partial - # capture never saved this calendar's token/checkpoint. - for observation in removed_calendar(missing, "access_revoked"): - yield observation - if missing_requested: - raise ValueError("A requested calendar is not accessible in current membership") diff --git a/backend/tests/unit/platform/sources/records/test_calendar_transport.py b/backend/tests/unit/platform/sources/records/test_calendar_transport.py new file mode 100644 index 0000000000..e60d7ab91e --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_calendar_transport.py @@ -0,0 +1,137 @@ +"""Actual decorated native boundary with synthetic HTTP and managed-provider failures.""" + +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from airweave.domains.auth_provider.exceptions import ( + AuthProviderRateLimitError, + AuthProviderServerError, +) +from airweave.domains.entities.canonical.page_source import InvalidScanContinuation +from airweave.domains.sources.exceptions import SourceError +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.storage.exceptions import FileSkippedException +from airweave.platform.configs.config import GoogleCalendarConfig +from airweave.platform.sources.google_calendar import GoogleCalendarSource + + +async def source(client, auth=None): + return await GoogleCalendarSource.create( + auth=auth or StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=client, + config=GoogleCalendarConfig(), + ) + + +@pytest.mark.parametrize( + "error", [AuthProviderServerError(status_code=503), AuthProviderRateLimitError(retry_after=60)] +) +async def test_managed_transient_retries_and_honors_delay(error): + replies = [error, httpx.Response(200, json={"items": []})] + calls = [] + + async def handler(request): + calls.append(request) + value = replies.pop(0) + if isinstance(value, Exception): + raise value + return value + + sleep = AsyncMock() + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + instance = await source(client) + call = GoogleCalendarSource._get_capture_json.retry_with(sleep=sleep) + assert await call(instance, "https://example.test") == {"items": []} + assert len(calls) == 2 + if isinstance(error, AuthProviderRateLimitError): + sleep.assert_awaited_once_with(60) + + +async def test_long_managed_retry_after_is_deferred_without_early_request(): + calls = [] + + async def handler(request): + calls.append(request) + raise AuthProviderRateLimitError(retry_after=180) + + sleep = AsyncMock() + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + instance = await source(client) + with pytest.raises(AuthProviderRateLimitError): + await GoogleCalendarSource._get_capture_json.retry_with(sleep=sleep)( + instance, "https://example.test" + ) + assert len(calls) == 1 + sleep.assert_not_awaited() + + +async def test_one_auth_refresh_and_only_explicit_invalid_token_resets(): + 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") + replies = [ + httpx.Response(401, json={}), + httpx.Response(200, json={"items": []}), + httpx.Response( + 400, + json={ + "error": { + "errors": [ + {"reason": "invalid", "location": "pageToken", "locationType": "parameter"} + ] + } + }, + ), + httpx.Response(400, json={"error": {"message": "invalid unrelated parameter"}}), + ] + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: replies.pop(0)) + ) as client: + instance = await source(client, auth) + assert await instance._get_capture_json("https://example.test") == {"items": []} + with pytest.raises(InvalidScanContinuation): + await instance._get_capture_json("https://example.test", params={"pageToken": "old"}) + with pytest.raises(SourceError): + await instance._get_capture_json("https://example.test", params={"pageToken": "old"}) + assert auth.refreshes == 1 and not replies + + +async def test_oversized_native_body_closes_before_full_download(): + 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(40): + self.observed += 1 + yield b"x" * 1024 * 1024 + + 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: + instance = await source(client) + with pytest.raises(FileSkippedException): + await instance._get_capture_json("https://example.test") + assert chunks.observed == 33 and chunks.closed diff --git a/backend/tests/unit/platform/sources/records/test_google_calendar.py b/backend/tests/unit/platform/sources/records/test_google_calendar.py index 8dd26ed4d7..32cd363bc8 100644 --- a/backend/tests/unit/platform/sources/records/test_google_calendar.py +++ b/backend/tests/unit/platform/sources/records/test_google_calendar.py @@ -1,20 +1,20 @@ -"""Calendar races and resets, with no provider mutation or live credentials.""" +"""Calendar native pagination and resource rules; recovery is verified with real SQL separately.""" from copy import deepcopy -from uuid import uuid4 import pytest -from airweave.domains.entities.canonical.requests import ( - CaptureRecord, - CompletedScope, - RemovedScope, - StartedScope, +from airweave.domains.entities.canonical.calendar import CalendarScopeContext +from airweave.domains.entities.canonical.page_source import InvalidScopeCheckpoint, ScopeAccessLost +from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.sources.exceptions import ( + SourceEntityNotFoundError, + SourceGoneError, + SourceServerError, ) -from airweave.domains.sources.exceptions import SourceGoneError, SourceServerError -from airweave.domains.syncs.cursors.cursor import SyncCursor -from airweave.platform.cursors.google_calendar import GoogleCalendarCursor -from airweave.platform.sources.records.google_calendar import generate_calendar_observations +from airweave.platform.configs.config import GoogleCalendarConfig +from airweave.platform.sources.records.calendar_pages import CalendarPages, CalendarProgress +from airweave.platform.sources.records.google_calendar import record class ScriptedGet: @@ -22,7 +22,7 @@ def __init__(self, replies): self.replies = iter(replies) self.calls = [] - async def __call__(self, url, params): + async def __call__(self, url, params=None): self.calls.append((url, deepcopy(params))) value = next(self.replies) if isinstance(value, Exception): @@ -30,14 +30,20 @@ async def __call__(self, url, params): return value -def cursor(tokens=None): - return SyncCursor(uuid4(), GoogleCalendarCursor, {"calendar_tokens": tokens or {}}) +def progress(kind="event", **updates): + params = {"maxResults": 250, "showDeleted": "true", "singleEvents": "false"} + if kind == "event_occurrence": + params.update( + singleEvents="true", timeMin="2026-03-01T00:00:00Z", timeMax="2026-04-01T00:00:00Z" + ) + return CalendarProgress( + mode="full", context=CalendarScopeContext(parameters=params), **updates + ).continuation() async def test_full_pagination_retains_sparse_cancellation_and_container(): get = ScriptedGet( [ - {"items": [{"id": "team/a@example.com"}]}, { "items": [{"id": "series", "recurrence": ["RRULE:FREQ=WEEKLY"]}], "nextPageToken": "page2", @@ -55,77 +61,35 @@ async def test_full_pagination_retains_sparse_cancellation_and_container(): }, ] ) - state = cursor() - results = [item async for item in generate_calendar_observations(get, state)] - events = [ - item - for item in results - if isinstance(item, CaptureRecord) and item.identity.record_type == "event" - ] - assert all(item.identity.container_id == "team/a@example.com" for item in events) - assert events[1].kind == "upsert" and events[1].payload["recurringEventId"] == "series" - assert "%2F" in get.calls[1][0] - assert get.calls[2][1]["pageToken"] == "page2" - assert state.data["calendar_tokens"] == {"team/a@example.com": "boundary"} - assert any(isinstance(item, CompletedScope) and item.record_type == "event" for item in results) - - -async def test_410_after_delta_page_resets_only_that_scope(): - get = ScriptedGet( - [ - {"items": [{"id": "a"}, {"id": "b"}]}, - {"items": [{"id": "stale"}], "nextPageToken": "delta2"}, - SourceGoneError("expired"), - {"items": [{"id": "current"}], "nextSyncToken": "a-new"}, - {"items": [], "nextSyncToken": "b-new"}, - ] + pages = CalendarPages(get, GoogleCalendarConfig()) + scope = CompletedScope(record_type="event", container_id="team/a@example.com") + first = await pages.page(scope, progress()) + assert not first.final and first.provider_checkpoint is None + second = await pages.page(scope, first.continuation) + assert second.final and second.provider_checkpoint.value == {"sync_token": "boundary"} + assert ( + second.records[0].kind == "upsert" + and second.records[0].identity.container_id == scope.container_id ) - state = cursor({"a": "a-old", "b": "b-old"}) - results = [item async for item in generate_calendar_observations(get, state)] - starts = [ - item.container_id - for item in results - if type(item) is StartedScope and item.record_type == "event" - ] - assert starts == ["a"] - assert get.calls[2][1]["syncToken"] == "a-old" - assert "syncToken" not in get.calls[3][1] and "pageToken" not in get.calls[3][1] - assert get.calls[4][1]["syncToken"] == "b-old" - assert state.data["calendar_tokens"] == {"a": "a-new", "b": "b-new"} + assert "%2F" in get.calls[0][0] and get.calls[1][1]["pageToken"] == "page2" -async def test_failure_keeps_original_cursor_and_does_not_complete_scope(): - get = ScriptedGet( - [ - {"items": [{"id": "a"}]}, - {"items": [{"id": "one"}], "nextPageToken": "next"}, - SourceServerError("temporary", status_code=503), - ] +async def test_410_only_delta_is_scope_reset_other_errors_fail(): + scope = CompletedScope(record_type="event", container_id="cal") + pages = CalendarPages(ScriptedGet([SourceGoneError("expired")]), GoogleCalendarConfig()) + delta = CalendarProgress.model_validate(progress().value).model_copy( + update={"mode": "changes", "sync_token": "old"} ) - state = cursor() - observed = [] - with pytest.raises(SourceServerError): - async for item in generate_calendar_observations(get, state): - observed.append(item) - assert state.data["calendar_tokens"] == {} - assert not any(type(item) is CompletedScope for item in observed) - - -async def test_missing_calendar_removes_only_previous_calendar_children(): - get = ScriptedGet([{"items": [{"id": "kept"}]}, {"items": [], "nextSyncToken": "new"}]) - state = cursor({"kept": "old", "removed": "old-removed"}) - results = [item async for item in generate_calendar_observations(get, state)] - removed = [item for item in results if type(item) is RemovedScope] - assert [(item.record_type, item.container_id, item.removal_reason) for item in removed] == [ - ("event_occurrence", "removed", "scope_removed"), - ("event", "removed", "scope_removed"), - ] - assert state.data["calendar_tokens"] == {"kept": "new"} + with pytest.raises(InvalidScopeCheckpoint): + await pages.page(scope, delta.continuation()) + for error in (SourceGoneError("full gone"), SourceServerError("temporary", status_code=503)): + with pytest.raises(type(error)): + await CalendarPages(ScriptedGet([error]), GoogleCalendarConfig()).page( + scope, progress() + ) def test_only_recurring_cancellations_are_retained_and_reinstatement_keeps_identity(): - from airweave.platform.sources.records.google_calendar import record - cancelled = { "id": "exception", "status": "cancelled", @@ -134,8 +98,7 @@ def test_only_recurring_cancellations_are_retained_and_reinstatement_keeps_ident } exclusion = record("event", cancelled, "calendar") active = record("event", {**cancelled, "status": "confirmed"}, "calendar") - assert exclusion.kind == active.kind == "upsert" - assert exclusion.identity == active.identity + assert exclusion.kind == active.kind == "upsert" and exclusion.identity == active.identity assert exclusion.payload == cancelled assert record("event", {"id": "single", "status": "cancelled"}, "calendar").kind == "delete" with pytest.raises(ValueError, match="original occurrence"): @@ -144,208 +107,87 @@ def test_only_recurring_cancellations_are_retained_and_reinstatement_keeps_ident ) -async def test_expanded_window_is_fixed_paged_and_separate_from_master_scope(): - from airweave.platform.configs.config import CalendarOccurrenceWindow - - window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") +async def test_occurrence_window_and_budgets_survive_continuation(): + scope = CompletedScope(record_type="event_occurrence", container_id="cal") get = ScriptedGet( [ - {"items": [{"id": "cal", "timeZone": "America/Los_Angeles"}]}, - { - "items": [{"id": "master", "recurrence": ["RRULE:FREQ=WEEKLY"]}], - "nextSyncToken": "master-token", - }, - { - "items": [ - { - "id": "real-instance", - "recurringEventId": "master", - "originalStartTime": {"date": "2026-03-01"}, - } - ], - "nextPageToken": "p2", - }, - {"items": [{"id": "cancelled-instance", "status": "cancelled"}]}, + {"items": [{"id": "a"}], "nextPageToken": "p2"}, + {"items": [{"id": "b", "status": "cancelled"}]}, ] ) - state = cursor() - observations = [item async for item in generate_calendar_observations(get, state, window)] - occurrences = [ - item - for item in observations - if isinstance(item, CaptureRecord) and item.identity.record_type == "event_occurrence" - ] - assert [item.identity.native_id for item in occurrences] == [ - "real-instance", - "cancelled-instance", - ] - assert occurrences[1].kind == "delete" - assert get.calls[2][1]["timeMin"] == get.calls[3][1]["timeMin"] - assert "syncToken" not in get.calls[2][1] and get.calls[2][1]["singleEvents"] == "true" - assert state.data["occurrence_coverage"]["cal"]["timezone"] == "America/Los_Angeles" - assert any( - type(item) is CompletedScope and item.record_type == "event_occurrence" - for item in observations - ) - - -async def test_failed_expanded_refresh_keeps_prior_coverage_and_does_not_reconcile(): - from airweave.platform.configs.config import CalendarOccurrenceWindow - - window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") - state = cursor({"cal": "old"}) - before = state.data - get = ScriptedGet( - [ - {"items": [{"id": "cal", "timeZone": "UTC"}]}, - {"items": [], "nextSyncToken": "new"}, - {"items": [{"id": "one"}], "nextPageToken": "p2"}, - SourceServerError("temporary", status_code=503), - ] + pages = CalendarPages(get, GoogleCalendarConfig()) + first = await pages.page(scope, progress("event_occurrence")) + second = await pages.page(scope, first.continuation) + assert ( + second.final and second.provider_checkpoint is None and second.records[0].kind == "delete" ) - observations = [] - with pytest.raises(SourceServerError): - async for item in generate_calendar_observations(get, state, window): - observations.append(item) - assert state.data == before - assert not any( - type(item) is CompletedScope and item.record_type == "event_occurrence" - for item in observations + assert get.calls[0][1]["timeMin"] == get.calls[1][1]["timeMin"] + assert second.continuation.value["records"] == 2 + with pytest.raises(ValueError, match="budget"): + await CalendarPages( + ScriptedGet([{"items": [{"id": "excess"}]}]), GoogleCalendarConfig() + ).page(scope, progress("event_occurrence", records=10000)) + + +async def test_repeated_pages_and_oversized_native_pages_fail(): + scope = CompletedScope(record_type="event_occurrence", container_id="cal") + pages = CalendarPages( + ScriptedGet([{"items": [], "nextPageToken": "same"}] * 2), GoogleCalendarConfig() ) - - -def test_explicit_historical_window_and_rolling_defaults_are_bounded(): + first = await pages.page(scope, progress("event_occurrence")) + with pytest.raises(ValueError, match="repeated"): + await pages.page(scope, first.continuation) + with pytest.raises(ValueError, match="bounded page"): + await CalendarPages( + ScriptedGet([{"items": [{"id": str(i)} for i in range(251)]}]), GoogleCalendarConfig() + ).page(scope, progress("event_occurrence")) + + +async def test_exact_missing_membership_and_accessible_omission_are_different(): + pages = CalendarPages(ScriptedGet([SourceEntityNotFoundError("gone")]), GoogleCalendarConfig()) + await pages.confirm_member_absent("missing") + with pytest.raises(ValueError, match="omitted accessible"): + await CalendarPages( + ScriptedGet([{"id": "visible"}]), GoogleCalendarConfig() + ).confirm_member_absent("visible") + with pytest.raises(ScopeAccessLost): + await CalendarPages( + ScriptedGet([SourceEntityNotFoundError("gone")]), GoogleCalendarConfig() + ).page(CompletedScope(record_type="event", container_id="cal"), progress()) + + +def test_explicit_historical_window_rolling_defaults_and_selection_are_bounded(): from datetime import timedelta - from airweave.platform.configs.config import GoogleCalendarConfig - - default = GoogleCalendarConfig().resolved_window() - assert default.end - default.start == timedelta(days=120) - chosen = GoogleCalendarConfig( - occurrence_window={"start": "2001-01-01T00:00:00Z", "end": "2001-02-01T00:00:00Z"} + assert ( + GoogleCalendarConfig().resolved_window().end + - GoogleCalendarConfig().resolved_window().start + == timedelta(days=120) + ) + assert ( + GoogleCalendarConfig( + occurrence_window={"start": "2001-01-01T00:00:00Z", "end": "2001-02-01T00:00:00Z"} + ) + .resolved_window() + .start.year + == 2001 ) - assert chosen.resolved_window().start.year == 2001 with pytest.raises(ValueError): GoogleCalendarConfig( occurrence_window={"start": "2001-01-01T00:00:00Z", "end": "2003-02-01T00:00:00Z"} ) - - -async def test_expanded_repeated_pages_never_claim_completion(): - from airweave.platform.configs.config import CalendarOccurrenceWindow - from airweave.platform.sources.records.google_calendar import capture_occurrences - - window = CalendarOccurrenceWindow(start="2026-01-01T00:00:00Z", end="2026-02-01T00:00:00Z") - get = ScriptedGet( - [{"items": [], "nextPageToken": "repeat"}, {"items": [], "nextPageToken": "repeat"}] + assert ( + GoogleCalendarConfig().calendar_ids is None + and GoogleCalendarConfig(calendar_ids=[]).calendar_ids == () ) - observed = [] - with pytest.raises(ValueError, match="repeated"): - async for item in capture_occurrences(get, "cal", window): - observed.append(item) - assert not any(type(item) is CompletedScope for item in observed) - - -async def test_selection_withdraws_old_scope_and_readd_bootstraps(): - state = cursor({"old": "old-token", "keep": "keep-token"}) - get = ScriptedGet( - [{"items": [{"id": "old"}, {"id": "keep"}]}, {"items": [], "nextSyncToken": "next"}] - ) - result = [ - item async for item in generate_calendar_observations(get, state, calendar_ids=("keep",)) - ] - assert len(get.calls) == 2 and "/keep/events" in get.calls[1][0] - assert get.calls[1][1]["syncToken"] == "keep-token" - assert state.data["calendar_tokens"] == {"keep": "next"} - assert {item.record_type for item in result if isinstance(item, RemovedScope)} == { - "event", - "event_occurrence", - } - assert all( - item.removal_reason == "scope_removed" for item in result if isinstance(item, RemovedScope) - ) - readd = ScriptedGet([{"items": [{"id": "old"}]}, {"items": [], "nextSyncToken": "fresh"}]) - _ = [item async for item in generate_calendar_observations(readd, state, calendar_ids=("old",))] - assert "syncToken" not in readd.calls[1][1] - - -async def test_empty_selection_makes_no_provider_requests_and_completes_membership(): - state = cursor({"old": "token"}) - get = ScriptedGet([]) - result = [item async for item in generate_calendar_observations(get, state, calendar_ids=())] - assert get.calls == [] and state.data["calendar_tokens"] == {} - assert any(type(item) is CompletedScope and item.record_type == "calendar" for item in result) - - -async def test_missing_requested_calendar_revokes_known_scope_without_completion(): - state = cursor({"missing": "token"}) - result = [] - with pytest.raises(ValueError, match="not accessible"): - async for item in generate_calendar_observations( - ScriptedGet([{"items": []}]), state, calendar_ids=("missing",) - ): - result.append(item) - assert state.data["calendar_tokens"] == {"missing": "token"} - assert any( - isinstance(item, RemovedScope) and item.removal_reason == "access_revoked" - for item in result - ) - assert not any(type(item) is CompletedScope for item in result) - - -def test_calendar_selection_config_distinguishes_none_empty_and_invalid(): - from airweave.platform.configs.config import GoogleCalendarConfig - - assert GoogleCalendarConfig().calendar_ids is None - assert GoogleCalendarConfig(calendar_ids=[]).calendar_ids == () for value in (["primary"], ["a", "a"], [""], [" a"]): with pytest.raises(ValueError): GoogleCalendarConfig(calendar_ids=value) -async def test_selected_calendar_loses_access_between_membership_and_events(): - from airweave.domains.sources.exceptions import SourceEntityNotFoundError - - state = cursor({"selected": "old"}) - get = ScriptedGet([{"items": [{"id": "selected"}]}, SourceEntityNotFoundError("gone")]) - result = [ - item - async for item in generate_calendar_observations(get, state, calendar_ids=("selected",)) - ] - revoked = [item for item in result if isinstance(item, RemovedScope)] - assert len(revoked) == 2 and all(item.removal_reason == "access_revoked" for item in revoked) - assert state.data["calendar_tokens"] == {} - - -async def test_occurrence_access_loss_discards_master_token_before_restoration(): - from airweave.domains.sources.exceptions import SourceEntityNotFoundError - from airweave.platform.configs.config import CalendarOccurrenceWindow - - window = CalendarOccurrenceWindow(start="2026-03-01T00:00:00Z", end="2026-04-01T00:00:00Z") - state = cursor({"cal": "old"}) - get = ScriptedGet( - [ - {"items": [{"id": "cal", "timeZone": "UTC"}]}, - {"items": [{"id": "unchanged"}], "nextSyncToken": "must-discard"}, - SourceEntityNotFoundError("synthetic access loss"), - ] - ) - lost = [item async for item in generate_calendar_observations(get, state, window)] - assert state.data["calendar_tokens"] == {} - assert state.data["occurrence_coverage"] == {} - assert {item.record_type for item in lost if isinstance(item, RemovedScope)} == { - "event", - "event_occurrence", - } - restored_get = ScriptedGet( - [ - {"items": [{"id": "cal", "timeZone": "UTC"}]}, - {"items": [{"id": "unchanged"}], "nextSyncToken": "restored"}, - {"items": []}, - ] - ) - restored = [item async for item in generate_calendar_observations(restored_get, state, window)] - assert "syncToken" not in restored_get.calls[1][1] - assert any(isinstance(item, StartedScope) and item.record_type == "event" for item in restored) - assert state.data["calendar_tokens"] == {"cal": "restored"} - assert "cal" in state.data["occurrence_coverage"] +def test_record_stamps_native_dates_without_inventing_capture_dates(): + native = {"id": "event", "created": "2026-03-01T00:00:00Z", "updated": "2026-03-02T00:00:00Z"} + captured = record("event", native, "cal") + assert captured.source_created_at.day == 1 and captured.source_updated_at.day == 2 + assert captured.payload == native + assert record("calendar", {"id": "cal"}).source_created_at is None From 826a2c1db50966509251206ff8a9e19eafdf18ab Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:34:03 -0700 Subject: [PATCH 107/266] Verify Calendar lifecycle from published scoped checkpoints --- .../tests/test_calendar_lifecycle_probe.py | 96 +++++++++++++++++++ backend/tests/live/calendar_lifecycle.py | 80 ++++++++++++++++ backend/tests/live/provider_lifecycle.py | 58 ++++++++--- 3 files changed, 223 insertions(+), 11 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py create mode 100644 backend/tests/live/calendar_lifecycle.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py new file mode 100644 index 0000000000..a3773319c9 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py @@ -0,0 +1,96 @@ +"""The real live harness must preserve scoped capture through fresh child processes.""" + +import json + +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 unittest.mock import MagicMock +from urllib.parse import quote +sys.path.insert(0, 'tests/live') +import conftest +import httpx +import provider_lifecycle as lifecycle +from provider_sample import verify_rest_identity +from airweave.domains.entities.canonical.tests.test_calendar_recovery import NativeHTTP, event +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.platform.sources.google_calendar import GoogleCalendarSource + +@asynccontextmanager +async def synthetic_source(name, account, expected_email, key, fence, **options): + assert name == 'google_calendar' + path = '/calendars/' + quote(expected_email, safe='') + '/events' + first = os.environ['TEST_STAGE'] == 'first' + native = NativeHTTP([ + ('/users/me/calendarList/primary', {'id': expected_email, 'primary': True}), + ('/users/me/calendarList', {'items': [{'id': expected_email, 'timeZone': 'UTC'}]}), + (path, {'items': [event('a')] if first else [], + 'nextSyncToken': 'full' if first else 'delta'}), + (path, {'items': [event('a')]}), + ]) + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle), + event_hooks={'request': [options['request_hook']]}) as client: + connector = await GoogleCalendarSource.create(auth=StaticTokenProvider('synthetic'), + logger=MagicMock(), http_client=client, config=options['calendar_config']) + await verify_rest_identity(name, connector, expected_email) + yield connector, 'provider_email' + assert not native.replies + +lifecycle.rest_source = synthetic_source +lifecycle.harness.test_database_url = lambda: os.environ['CANONICAL_TEST_DATABASE_URL'] +os.environ['COMPOSIO_API_KEY'] = 'synthetic-test-only' +os.environ['LIVE_CALENDAR_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'])))) +""" + + +async def test_calendar_lifecycle_initial_then_published_scope_delta(database, source, tmp_path): + fence = source[1] + async with database() as db: + schema = await db.scalar(text("select current_schema()")) + manifest = { + "provider": "google_calendar", + "schema": schema, + "root": str(tmp_path), + "organization_id": str(fence.organization_id), + "sync_id": str(fence.sync_id), + "record_limit": 10, + "request_limit": 20, + "timeout": 30, + "file_byte_limit": 1024, + "blob_byte_limit": 1024, + "calendar_config": { + "calendar_ids": ["synthetic@example.test"], + "occurrence_window": { + "start": "2026-03-01T00:00:00Z", + "end": "2026-04-01T00:00:00Z", + }, + }, + } + results = [] + for stage in ("first", "second"): + current = dict(manifest) + if stage == "first": + current.update(job_id=str(fence.job_id), attempt_number=2) + code, output, error = await subprocess_script( + SCRIPT, {"TEST_MANIFEST": json.dumps(current), "TEST_STAGE": stage} + ) + if code != 0: + raise AssertionError(output + error) + result = json.loads(output.splitlines()[-1]) + assert not result["failed"] and result["job_status"] == "completed" + assert result["stored_records"] == 3 and result["partial_records"] == 0 + assert result["provider_requests"] == 4 + results.append(result) + first, second = results + assert first["records_observed"] == 3 and second["records_observed"] == 2 + assert not first["loaded_durable_checkpoint"] and second["loaded_durable_checkpoint"] + assert first["sync_token_requests"] == 0 and second["sync_token_requests"] == 1 + assert first["full_scope_completed"] and not second["full_scope_completed"] + assert first["observed_change_sequence"] == second["observed_change_sequence"] + assert first["payload_revision_digest"] == second["payload_revision_digest"] diff --git a/backend/tests/live/calendar_lifecycle.py b/backend/tests/live/calendar_lifecycle.py new file mode 100644 index 0000000000..ab8afe823b --- /dev/null +++ b/backend/tests/live/calendar_lifecycle.py @@ -0,0 +1,80 @@ +"""Private SQL assertions for the selected-calendar lifecycle; no token output.""" + +from datetime import datetime +from urllib.parse import unquote + +from sqlalchemy import select + +from airweave.domains.entities.canonical.coverage import capture_coverage +from airweave.domains.entities.canonical.scope_execution import ScopeExecution +from airweave.models.capture_scan import CaptureScan + + +async def event_checkpoint(db, organization_id, sync_id, calendar_id): + row = await db.scalar( + select(CaptureScan).where( + CaptureScan.organization_id == organization_id, + CaptureScan.sync_id == sync_id, + CaptureScan.record_type == "event", + CaptureScan.container_id == calendar_id, + ) + ) + if row is None or not row.execution_state: + return None + execution = ScopeExecution.model_validate(row.execution_state) + return execution.published.checkpoint if execution.published else None + + +def count_calendar_request(request, calendar_id, previous, counters): + if not request.url.params.get("syncToken"): + return + assert previous is not None + assert unquote(request.url.path).endswith(f"/calendars/{calendar_id}/events") + assert request.url.params["syncToken"] == previous.value["sync_token"] + counters["sync_token_requests"] += 1 + + +async def verify_calendar_scopes(sessions, organization_id, sync_id, manifest, saved, loaded): + """Require real completed scope publication and current SQL public coverage.""" + if manifest["provider"] != "google_calendar": + return + calendar_id = manifest["calendar_config"]["calendar_ids"][0] + async with sessions() as db: + rows = list( + await db.scalars( + select(CaptureScan).where( + CaptureScan.organization_id == organization_id, CaptureScan.sync_id == sync_id + ) + ) + ) + coverage = (await capture_coverage(db, organization_id, (sync_id,)))[sync_id] + assert len(rows) == 3 + assert {row.record_type for row in rows} == {"calendar", "event", "event_occurrence"} + cycle = saved["canonical_cycle"] + for row in rows: + assert row.phase == "complete" and str(row.cycle_id) == cycle["version"]["cycle_id"] + assert row.container_id == (None if row.record_type == "calendar" else calendar_id) + execution = ScopeExecution.model_validate(row.execution_state) + assert execution.last_full is not None and execution.published is not None + assert execution.published.sweep_id == row.sweep_id + if row.record_type == "event": + assert execution.plan.mode == ("changes" if loaded else "full") + assert execution.published.checkpoint.value["sync_token"] + if row.record_type == "event_occurrence": + params = execution.published.request_context["parameters"] + window = manifest["calendar_config"]["occurrence_window"] + + assert datetime.fromisoformat(params["timeMin"]) == datetime.fromisoformat( + window["start"] + ) + assert datetime.fromisoformat(params["timeMax"]) == datetime.fromisoformat( + window["end"] + ) + assert coverage.phase == "complete" and coverage.mode == "mixed" + assert coverage.last_full_capture is None and coverage.provider_checkpoint_promoted_at is None + assert coverage.scope_summary.model_dump() == { + "eligible": 3, + "completed_full": 2 if loaded else 3, + "completed_changes": 1 if loaded else 0, + "unfinished": 0, + } diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 676b62b7a7..600df790ca 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -22,6 +22,7 @@ from uuid import UUID, uuid4 import canonical_capture as harness # settings must precede application imports +from calendar_lifecycle import count_calendar_request, event_checkpoint, verify_calendar_scopes from provider_sample import rest_source, wispr_source from sqlalchemy import select, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine @@ -33,6 +34,7 @@ InvalidCaptureCheckpoint, InvalidScanContinuation, KnownObjectSource, + ScopedPageSource, ) from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope from airweave.domains.entities.canonical.service import CanonicalCaptureService @@ -134,8 +136,32 @@ class BoundedCheckpointKnownSource(BoundedCheckpointSource, BoundedKnownSource): pass +class BoundedScopedSource(BoundedPageSource): + async def prepare_cycle(self, previous): + return await self.source.prepare_cycle(previous) + + async def prepare_scope(self, scope, cycle, previous, *, parent, force_full): + return await self.source.prepare_scope( + scope, cycle, previous, parent=parent, force_full=force_full + ) + + def initial_scope_continuation(self, scope, cycle, plan): + return self.source.initial_scope_continuation(scope, cycle, plan) + + +class BoundedScopedKnownSource(BoundedScopedSource, BoundedKnownSource): + pass + + def bounded_page_source(source, counters, record_limit, probe=None): """Preserve only protocols the actual source implements.""" + if isinstance(source, ScopedPageSource): + wrapper = ( + BoundedScopedKnownSource + if isinstance(source, KnownObjectSource) + else BoundedScopedSource + ) + return wrapper(source, counters, record_limit, probe) checkpoint = isinstance(source, CheckpointedPageSource) known = isinstance(source, KnownObjectSource) wrapper = ( @@ -210,12 +236,11 @@ def safe_failure_reason(error): def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attempt, sequence): """Match production cursor/scope capabilities; never manufacture a resume claim.""" if name == "google_calendar": - calendar_id = manifest["calendar_config"]["calendar_ids"][0] - assert set(saved["calendar_tokens"]) == {calendar_id} - assert saved["calendar_tokens"][calendar_id] - assert set(saved["occurrence_coverage"]) == {calendar_id} - expected_scopes = 2 if loaded else 3 - assert counters["started"] == counters["completed"] == expected_scopes + cycle = saved["canonical_cycle"] + assert cycle["phase"] == "complete" and cycle["mode"] == "mixed" + assert cycle["completed_job_id"] + assert cycle["last_full_capture"] is None and cycle["promoted_checkpoint"] is None + assert counters["started"] == counters["completed"] == 0 assert bool(counters["sync_token_requests"]) == loaded elif name == "google_drive": cycle = saved["canonical_cycle"] @@ -313,8 +338,7 @@ def count_drive_request(name, request, previous, counters): promoted = previous.get("canonical_cycle", {}).get("promoted_checkpoint") if ( promoted - and request.url.params.get("pageToken") - == promoted["checkpoint"]["value"]["page_token"] + and request.url.params.get("pageToken") == promoted["checkpoint"]["value"]["page_token"] ): counters["resumed_changes_requests"] += 1 @@ -331,6 +355,7 @@ async def child(manifest): counters = {"provider_requests": 0, "records_observed": 0, "started": 0, "completed": 0} result = {"failed": True} previous = {} + previous_calendar_checkpoint = None name = manifest["provider"] is_calendar = name == "google_calendar" counters["sync_token_requests"] = 0 @@ -376,8 +401,13 @@ async def request_hook(request): raise BudgetExceeded("provider_requests") counters["provider_requests"] += 1 count_gmail_request(name, request, previous, counters, identity_verified) - if request.url.params.get("syncToken"): - counters["sync_token_requests"] += 1 + if is_calendar: + count_calendar_request( + request, + manifest["calendar_config"]["calendar_ids"][0], + previous_calendar_checkpoint, + counters, + ) count_drive_request(name, request, previous, counters) @asynccontextmanager @@ -417,6 +447,10 @@ async def db_context(): cursor_service = SyncCursorService() async with sessions() as db: previous = await cursor_service.get_cursor_data(db, sync_id, ctx) + if is_calendar: + previous_calendar_checkpoint = await event_checkpoint( + db, organization_id, sync_id, manifest["calendar_config"]["calendar_ids"][0] + ) cursor_schema = { "gmail": GmailCursor, "google_calendar": GoogleCalendarCursor, @@ -539,6 +573,7 @@ async def db_context(): rows = list((await db.scalars(select(Entity).where(Entity.sync_id == sync_id))).all()) assert status == "completed" validate_checkpoint(name, manifest, counters, saved, previous, loaded, attempt, sequence) + await verify_calendar_scopes(sessions, organization_id, sync_id, manifest, saved, loaded) visible = [r for r in rows if r.deleted_at is None and r.source_payload is not None] digest = hashlib.sha256( json.dumps( @@ -572,7 +607,8 @@ async def db_context(): result = { "failed": False, "full_scope_completed": name != "wispr" - and not (name == "gmail" and saved["canonical_cycle"]["mode"] == "changes"), + and not (name == "gmail" and saved["canonical_cycle"]["mode"] == "changes") + and not (is_calendar and loaded), "counters_complete": True, "resumed_saved_page": resume_probe.resumed_saved_page if resume_probe else False, "saved_cursor_expired": resume_probe.saved_cursor_expired if resume_probe else False, From 6987d612f68f0a249a1dfe85b5e319033f1f5a0b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:37:46 -0700 Subject: [PATCH 108/266] Record live scoped Calendar delta and HTTP consumer proof --- .../calendar-scoped-cycles-20260930.json | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 backend/tests/live/evidence/calendar-scoped-cycles-20260930.json diff --git a/backend/tests/live/evidence/calendar-scoped-cycles-20260930.json b/backend/tests/live/evidence/calendar-scoped-cycles-20260930.json new file mode 100644 index 0000000000..51c56a543d --- /dev/null +++ b/backend/tests/live/evidence/calendar-scoped-cycles-20260930.json @@ -0,0 +1,54 @@ +{ + "commit": "826a2c1db50966509251206ff8a9e19eafdf18ab", + "almanac_commit": "c31ac0d9b0f7d9c25e052f8c5446db055d9da654", + "provider": "google_calendar", + "verification": "Real read-only Google provider; production capture pipeline; disposable local PostgreSQL; two fresh worker processes and real loopback HTTP consumer", + "first": { + "job_status": "completed", + "provider_requests": 4, + "records_observed": 188, + "stored_records": 181, + "partial_records": 0, + "observed_change_sequence": 188, + "loaded_durable_checkpoint": false + }, + "second": { + "job_status": "completed", + "provider_requests": 4, + "records_observed": 20, + "stored_records": 181, + "partial_records": 0, + "observed_change_sequence": 188, + "loaded_durable_checkpoint": true, + "sync_token_requests": 1, + "identical_payload_revision_digest": true, + "scope_summary": { + "eligible": 3, + "completed_full": 2, + "completed_changes": 1, + "unfinished": 0 + } + }, + "almanac_consumer": { + "verified": true, + "events": 19, + "pages": 4, + "native_identities_verified": true, + "coverage_verified": true, + "wrong_source_denied": true, + "uncaptured_range_denied": true + }, + "harness_exit_code": 0, + "independently_verified_cleanup": { + "remaining_private_schemas": 0, + "remaining_private_directories": 0, + "archives_removed": true + }, + "limitations": [ + "Selected primary calendar and fixed seven-day occurrence window only", + "No live induced edits, deletion, revocation or all-day/DST example", + "Loopback fixture authentication; no hosted WorkOS owner proof", + "No staging deployment" + ], + "harness_summary_note": "Generic identical_second_read and saved_cursor_reused summarize full-scope/recovery trials and are false for mixed Calendar; per-cycle native token assertions and matching revision digests establish this delta proof." +} From 33e7ed28493224cc0b35ef7c4b924e131c42a306 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:41:20 -0700 Subject: [PATCH 109/266] Report live token reuse independently of full capture evidence --- .../canonical/tests/test_resume_probe.py | 32 +++++++++++++++++++ backend/tests/live/provider_lifecycle.py | 31 +++++++++++++----- 2 files changed, 55 insertions(+), 8 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py index 8b6e6f98c4..257317f616 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py @@ -219,3 +219,35 @@ async def test_parent_trial_keeps_aggregate_budget_and_marks_killed_usage_unknow ) 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/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index 600df790ca..e2f26d1c8d 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -816,6 +816,28 @@ def trial_counter_summary(results): } +def trial_read_summary(results): + """Report observed token reuse and revision stability independently of full scopes.""" + successful = [item for item in results if not item.get("failed", True)] + digests = [item.get("payload_revision_digest") for item in successful] + reused = any( + item.get("loaded_durable_checkpoint", False) + and not item.get("saved_cursor_expired", False) + and ( + item.get("resumed_saved_page", False) + or item.get("sync_token_requests", 0) > 0 + or item.get("resumed_changes_requests", 0) > 0 + ) + for item in successful + ) + return { + "saved_cursor_reused": reused, + "identical_second_read": len(results) == len(successful) == 2 + and bool(digests[0]) + and digests[0] == digests[1], + } + + def wispr_request_limit(value: str) -> int: """Validate an explicitly reduced trial budget before setup or network access.""" try: @@ -962,7 +984,6 @@ async def resume_evidence(): async with admin.begin() as connection: await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) await admin.dispose() - completed = [item for item in results if item.get("full_scope_completed")] resume_verified = ( name == "slack" and len(results) == 3 @@ -998,14 +1019,8 @@ async def resume_evidence(): item.get("resumed_page_committed", False) for item in results ), "resumed_partial": any(item.get("resumed_partial", False) for item in results), - "saved_cursor_reused": bool( - (resume_verified or gmail_verified) - and not results[1].get("saved_cursor_expired") - ), **trial_counter_summary(results), - "identical_second_read": len(completed) == 2 - and completed[0].get("payload_revision_digest") - == completed[1].get("payload_revision_digest"), + **trial_read_summary(results), } ) ) From f47eb906a33ccf3e4583c3199616ba787949df39 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:43:41 -0700 Subject: [PATCH 110/266] Retain native Slack timestamps for canonical date filters --- .../canonical/tests/test_slack_recovery.py | 65 +++++++++++++++++++ backend/airweave/platform/sources/slack.py | 36 +++++++++- .../platform/sources/test_slack_capture.py | 46 +++++++++++++ 3 files changed, 146 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index 1e07836f46..2ee1adcf76 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -533,3 +533,68 @@ async def test_missing_queued_thread_restarts_inventory_before_absence(database, 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, _, _ = 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: + (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/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 774442ff91..c3416cfe82 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -4,7 +4,7 @@ import hashlib import re -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from functools import partial from typing import Any, AsyncGenerator, Dict, List, Optional @@ -50,6 +50,32 @@ from airweave.platform.sources.http_helpers import _parse_retry_after, raise_for_status 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) @@ -208,6 +234,8 @@ async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict: @staticmethod def _capture_message(message: dict, channel_id: str) -> CaptureRecord: """Keep the complete native response; file bytes are explicitly not captured yet.""" + 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 @@ -215,6 +243,8 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: 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: @@ -255,6 +285,10 @@ async def _validated_capture_page( 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 ) diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index bb47c9feba..c2539dd4ee 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -174,3 +174,49 @@ async def test_thread_not_found_outside_reply_fetch_does_not_invalidate_inventor ScanContinuation(), files=MagicMock(), ) + + +@pytest.mark.parametrize( + "native,created,updated", + [ + ( + {"ts": "1355517523.000005", "edited": {"ts": "1355517536.000001"}}, + "2012-12-14T20:38:43.000005+00:00", + "2012-12-14T20:38:56.000001+00:00", + ), + ( + {"ts": "1355517523.000005", "latest_reply": "1855517536.999999"}, + "2012-12-14T20:38:43.000005+00:00", + None, + ), + ({"ts": "invalid", "edited": {"ts": "NaN"}}, None, None), + ({"ts": "999999999999.123456", "edited": {"ts": True}}, None, None), + ], +) +def test_message_native_dates_preserve_precision_and_unknowns(native, created, updated): + record = SlackSource._capture_message(native, "C1") + assert record.payload == native + assert (record.source_created_at.isoformat() if record.source_created_at else None) == created + assert (record.source_updated_at.isoformat() if record.source_updated_at else None) == updated + + +@pytest.mark.asyncio +async def test_conversation_dates_use_distinct_native_units(): + connector = source() + connector._get = AsyncMock( + return_value={ + "channels": [ + {"id": "C1", "created": 1449252889, "updated": 1689965803820}, + {"id": "C2", "created": True, "updated": "1689965803820"}, + {"id": "C3"}, + ] + } + ) + result = await connector.capture_page( + CompletedScope(record_type="channel"), ScanContinuation(), files=MagicMock() + ) + first, malformed, absent = result.records + assert first.source_created_at.isoformat() == "2015-12-04T18:14:49+00:00" + assert first.source_updated_at.isoformat() == "2023-07-21T18:56:43.820000+00:00" + assert malformed.source_created_at is malformed.source_updated_at is None + assert absent.source_created_at is absent.source_updated_at is None From b070aabfc4806aa43d43b80f0bf8cc5e442fbaeb Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:49:05 -0700 Subject: [PATCH 111/266] Record existing Drive account native Docs access proof --- .../google-workspace-access-20260930.json | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 backend/tests/live/evidence/google-workspace-access-20260930.json diff --git a/backend/tests/live/evidence/google-workspace-access-20260930.json b/backend/tests/live/evidence/google-workspace-access-20260930.json new file mode 100644 index 0000000000..814ff07066 --- /dev/null +++ b/backend/tests/live/evidence/google-workspace-access-20260930.json @@ -0,0 +1,40 @@ +{ + "probe": "native_google_workspace_read_access", + "date": "2026-09-30", + "transport_revision": "f47eb90", + "temporary_probe_sha256": "1295b94da8446e7b517aeee4b7cc0457ceeb04cf6a857f5ba7194215152dd1bb", + "read_only": true, + "actual_http_calls": 5, + "call_limit": 6, + "timeout_seconds": 150, + "identity_checks": [ + "Exact Composio connected-account ID", + "Active googledrive toolkit", + "Native Drive account email matches selected account" + ], + "docs": { + "sample_source": "Drive files.list, MIME document, pageSize1", + "native_request": "GET docs.googleapis.com/v1/documents/{same_file_id}?includeTabsContent=true", + "status": 200, + "document_id_matches_drive_file": true, + "response_bytes": 47353, + "tabs_array_returned": true + }, + "sheets": { + "sample_source": "Drive files.list, MIME spreadsheet, pageSize1", + "sample_returned": false, + "native_api_tested": false + }, + "privacy": { + "provider_content_persisted": false, + "provider_identities_in_evidence": false, + "grants_changed": false, + "production_host_allowlist_changed": false, + "provider_writes": false + }, + "limitations": [ + "One document proves account capability, not complete Docs capture", + "No Sheets native API request or scope proof", + "No durable structured representations, indexing, widgets or hosted authorization verified" + ] +} From fc7ae61f4636614c9937a02c790ecbdef371ee05 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:57:42 -0700 Subject: [PATCH 112/266] Define retained Docs representations without changing legacy blob hashes --- .../domains/entities/canonical/requests.py | 20 +++ .../sources/records/workspace_manifest.py | 169 ++++++++++++++++++ .../records/test_workspace_manifest.py | 113 ++++++++++++ 3 files changed, 302 insertions(+) create mode 100644 backend/airweave/platform/sources/records/workspace_manifest.py create mode 100644 backend/tests/unit/platform/sources/records/test_workspace_manifest.py diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 972ac36cf0..1674938642 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -12,6 +12,7 @@ Field, JsonValue, field_validator, + model_serializer, model_validator, ) @@ -68,6 +69,23 @@ class BlobReference(BaseModel): default=None, description="RFC 6901 JSON Pointer into the unchanged native payload" ) + role: Literal["representation_manifest"] | None = None + + @model_validator(mode="after") + def manifest_is_not_native_location(self) -> "BlobReference": + """A representation manifest is provenance, not part of the provider payload.""" + if self.role is not None and self.source_path is not None: + raise ValueError("Representation manifests cannot have a native source path") + return self + + @model_serializer(mode="wrap") + def historical_descriptor(self, handler): + """Preserve historical hashes without dropping unrelated nullable fields.""" + value = handler(self) + if self.role is None: + value.pop("role", None) + return value + class CaptureRecord(BaseModel): """One current provider observation; sparse deletions are valid records.""" @@ -112,6 +130,8 @@ def unique_visibility_fields(cls, value): @model_validator(mode="after") def validate_removal(self) -> "CaptureRecord": """Keep provider deletion distinct from inaccessible or deselected scope.""" + if sum(blob.role == "representation_manifest" for blob in self.blobs) > 1: + raise ValueError("A record can retain at most one representation manifest") if self.kind == "delete" and self.removal_reason is None: raise ValueError("delete observations require removal_reason") if self.kind == "upsert" and self.removal_reason is not None: 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/tests/unit/platform/sources/records/test_workspace_manifest.py b/backend/tests/unit/platform/sources/records/test_workspace_manifest.py new file mode 100644 index 0000000000..132588a077 --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_workspace_manifest.py @@ -0,0 +1,113 @@ +"""Representation closure and compatibility, without changing Gmail MIME ownership.""" + +import hashlib +from datetime import datetime, timezone + +import pytest +from pydantic import ValidationError + +from airweave.domains.entities.canonical.requests import ( + BlobReference, + CaptureRecord, + RecordIdentity, +) +from airweave.domains.entities.canonical.store import capture_fingerprint +from airweave.platform.sources.records.workspace_manifest import ( + DocsState, + ExportState, + WorkspaceManifestV1, + canonical_json, + parse_manifest, + validate_document, +) + + +def descriptor(content, **extra): + digest = hashlib.sha256(content).hexdigest() + return BlobReference( + key=f"canonical/sync/blobs/sha256/{digest}", sha256=digest, size_bytes=len(content), **extra + ) + + +def fixture_manifest(): + native = descriptor(b'{"documentId":"doc","tabs":[]}') + model = WorkspaceManifestV1( + file_id="doc", + drive_version="7", + export=ExportState(status="unavailable", reason="export_size_limit"), + native=DocsState(status="complete", document_blob=native.sha256, embedded_media="retained"), + ) + content = canonical_json(model.model_dump(mode="json")) + return content, (native, descriptor(content, role="representation_manifest")) + + +def test_manifest_exact_closure_and_identity(): + content, blobs = fixture_manifest() + assert ( + parse_manifest(content, file_id="doc", drive_version="7", blobs=blobs).native.status + == "complete" + ) + for bad_blobs in (blobs[1:], blobs + (descriptor(b"unreferenced"),), blobs + (blobs[0],)): + with pytest.raises(ValueError): + parse_manifest(content, file_id="doc", drive_version="7", blobs=bad_blobs) + for file_id, version in (("foreign", "7"), ("doc", "8")): + with pytest.raises(ValueError): + parse_manifest(content, file_id=file_id, drive_version=version, blobs=blobs) + + +def test_generic_role_preserves_historical_shape_and_repeated_mime_digest(): + blob = descriptor(b"same", source_path="/payload/parts/0/body") + other = blob.model_copy(update={"source_path": "/payload/parts/1/body"}) + record = CaptureRecord( + identity=RecordIdentity(record_type="message", native_id="one"), + payload={}, + observed_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + blobs=(blob, other), + ) + assert "role" not in record.model_dump(mode="json")["blobs"][0] + assert record.model_dump(mode="json")["blobs"][0]["media_type"] is None + assert len(record.blobs) == 2 + assert ( + capture_fingerprint(record) + == "97c6dc0cc6d1376f18aa2f04d7fc66781a3bcfb729dda843dc90ec7037d4bc9d" + ) + with pytest.raises(ValidationError): + descriptor(b"manifest", role="representation_manifest", source_path="/payload") + with pytest.raises(ValidationError): + CaptureRecord.model_validate( + record.model_dump() + | { + "blobs": ( + descriptor(b"first", role="representation_manifest"), + descriptor(b"second", role="representation_manifest"), + ) + } + ) + + +def test_all_tabs_identity_and_unknown_native_fields(): + document = { + "documentId": "doc", + "futureField": {"kept": True}, + "tabs": [ + { + "tabProperties": {"tabId": "a"}, + "documentTab": {}, + "childTabs": [ + { + "tabProperties": {"tabId": "b"}, + "documentTab": {"future": 1}, + } + ], + } + ], + } + validate_document(document, file_id="doc") + assert "revisionId" not in document + with pytest.raises(ValueError): + validate_document(document, file_id="foreign") + with pytest.raises(ValueError): + validate_document({"documentId": "doc", "body": {}}, file_id="doc") + document["tabs"].append(document["tabs"][0]) + with pytest.raises(ValueError): + validate_document(document, file_id="doc") From 89aadd088aede1171e457b410f81d8bab473c0e8 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 11:58:31 -0700 Subject: [PATCH 113/266] feat: read and project retained native Google Docs --- .../entities/canonical/projection_mappers.py | 27 +- .../entities/canonical/workspace_docs.py | 165 ++++++++++ .../domains/entities/test_workspace_docs.py | 293 ++++++++++++++++++ 3 files changed, 476 insertions(+), 9 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/workspace_docs.py create mode 100644 backend/tests/unit/domains/entities/test_workspace_docs.py diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index e371b33dd2..0224e00f54 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -65,15 +65,24 @@ async def _drive( breadcrumbs=[], ), ) - 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: - raise ProjectionMappingError("Drive captured bytes have no known file format") - suffix = suffix.lower() + if any(blob.role == "representation_manifest" for blob in record.blobs): + from airweave.domains.entities.canonical.workspace_docs import read_document + + document = await read_document(record, storage) + content = document.text().encode("utf-8") + media_type = "text/plain" + suffix = ".txt" + 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: + raise ProjectionMappingError("Drive captured bytes have no known file format") + suffix = suffix.lower() local_path = await write_blob(content, directory, suffix=suffix) metadata = { **data, 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/tests/unit/domains/entities/test_workspace_docs.py b/backend/tests/unit/domains/entities/test_workspace_docs.py new file mode 100644 index 0000000000..642effcb1e --- /dev/null +++ b/backend/tests/unit/domains/entities/test_workspace_docs.py @@ -0,0 +1,293 @@ +"""Native Docs representation preserves structure and yields one deliberate search view.""" + +import hashlib +import json +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.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity +from airweave.domains.entities.canonical.workspace_docs import DocumentReadError, read_document + + +def paragraph(text): + return {"paragraph": {"elements": [{"textRun": {"content": text}}]}} + + +def document(): + return { + "documentId": "doc", + "title": "Native title", + "futureDocumentField": {"keep": True}, + "tabs": [ + { + "tabProperties": {"tabId": "first", "title": "First"}, + "documentTab": { + "body": { + "content": [ + paragraph("Body"), + { + "table": { + "tableRows": [ + { + "tableCells": [{"content": [paragraph("Cell")]}], + } + ] + } + }, + ] + }, + "headers": {"h": {"content": [paragraph("Header")]}}, + "footers": {"f": {"content": [paragraph("Footer")]}}, + "footnotes": {"n": {"content": [paragraph("Footnote")]}}, + }, + "childTabs": [ + { + "tabProperties": {"tabId": "child", "title": "Nested"}, + "documentTab": {"body": {"content": [paragraph("Child content")]}}, + "unknownTabProperty": 12, + } + ], + } + ], + } + + +def capture(native=None, *, schema_version=1): + native = document() if native is None else native + sync = uuid4() + contents = {} + + def blob(value, media_type="application/json", role=None): + data = value if isinstance(value, bytes) else json.dumps(value).encode() + digest = hashlib.sha256(data).hexdigest() + key = f"canonical/{sync}/blobs/sha256/{digest}" + contents[key] = data + return BlobReference( + key=key, sha256=digest, size_bytes=len(data), media_type=media_type, role=role + ) + + native_blob = blob(native) + export = blob(b"NOT THE SEARCH BODY", "application/octet-stream") + manifest = blob( + { + "schema_version": schema_version, + "file_id": "doc", + "drive_version": "42", + "export": {"status": "retained", "blob": export.sha256, "reason": None}, + "native": { + "kind": "docs", + "status": "complete", + "document_blob": native_blob.sha256, + "include_tabs_content": True, + "suggestions_view": "DEFAULT_FOR_CURRENT_ACCESS", + "comments_view": "COMMENTS_VIEW_MODE_OMITTED", + "embedded_media": "not_retained", + "media": [], + "missing": [], + }, + }, + role="representation_manifest", + ) + record = SourceRecord( + id=uuid4(), + sync_id=sync, + identity=RecordIdentity(record_type="file", native_id="doc"), + revision=1, + payload={ + "id": "doc", + "name": "Document", + "version": "42", + "mimeType": "application/vnd.google-apps.document", + }, + payload_schema_version=1, + capture_hash="hash", + content_hash=manifest.sha256, + completeness="complete", + observed_at=datetime.now(timezone.utc), + source_created_at=None, + source_updated_at=None, + deleted_at=None, + removal_reason=None, + blobs=(export, native_blob, manifest), + indexed_revision=None, + indexed_pipeline_version=None, + ) + storage = AsyncMock() + + async def read(key, *, max_bytes): + assert len(contents[key]) <= max_bytes + return contents[key] + + storage.read_file.side_effect = read + return record, storage, contents + + +@pytest.mark.asyncio +async def test_recursive_tabs_preserve_original_structure_without_edit_revision(): + record, storage, _ = capture() + result = await read_document(record, storage) + assert result.document == document() + assert "revisionId" not in result.document + assert result.tab("child") == document()["tabs"][0]["childTabs"][0] + with pytest.raises(DocumentReadError, match="not present"): + result.tab("missing") + text = result.text() + for expected in ("Body", "Cell", "Header", "Footer", "Footnote", "Child content"): + assert text.count(expected) == 1 + + +@pytest.mark.asyncio +async def test_projection_reads_native_once_and_never_indexes_export(): + record, storage, _ = capture() + async with map_record(record, "google_drive", storage) as entities: + assert len(entities) == 1 + entity = entities[0] + assert entity.file_id == "doc" + assert entity.mime_type == "text/plain" + text = Path(entity.local_path).read_text() + assert "Child content" in text and "NOT THE SEARCH BODY" not in text + assert all(call.args[0] != record.blobs[0].key for call in storage.read_file.call_args_list) + + +@pytest.mark.asyncio +async def test_unknown_text_shape_keeps_original_read_but_fails_projection(): + native = document() + native["tabs"][0]["documentTab"]["body"]["content"].append( + { + "paragraph": {"elements": [{"futureText": {"content": "Do not silently lose"}}]}, + } + ) + record, storage, _ = capture(native) + assert (await read_document(record, storage)).document == native + with pytest.raises(DocumentReadError, match="Unsupported"): + async with map_record(record, "google_drive", storage): + pass + + +@pytest.mark.asyncio +async def test_unavailable_record_does_not_read_storage(): + record, storage, _ = capture() + with pytest.raises(DocumentReadError, match="unavailable"): + await read_document(record.model_copy(update={"content_access": "unavailable"}), storage) + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_native_bytes_are_verified_not_trusted_from_manifest(): + record, storage, contents = capture() + contents[record.blobs[1].key] = b"corrupt" + with pytest.raises(ValueError, match="size and digest"): + await read_document(record, storage) + + +@pytest.mark.asyncio +async def test_manifest_cannot_resolve_a_blob_outside_its_record(): + record, storage, _ = capture() + record = record.model_copy(update={"blobs": (record.blobs[0], record.blobs[2])}) + with pytest.raises(ValueError, match="account"): + await read_document(record, storage) + assert storage.read_file.await_count == 1 + + +@pytest.mark.asyncio +async def test_native_document_identity_must_match_drive_original(): + native = document() + native["documentId"] = "other-file" + record, storage, _ = capture(native) + with pytest.raises(ValueError, match="another file"): + await read_document(record, storage) + + +@pytest.mark.asyncio +async def test_duplicate_tab_ids_are_not_resolved_by_first_match(): + native = document() + native["tabs"][0]["childTabs"][0]["tabProperties"]["tabId"] = "first" + record, storage, _ = capture(native) + with pytest.raises(ValueError, match="duplicated"): + await read_document(record, storage) + + +@pytest.mark.asyncio +async def test_drive_version_mismatch_does_not_read_native_response(): + record, storage, _ = capture() + record = record.model_copy(update={"payload": {**record.payload, "version": "43"}}) + with pytest.raises(ValueError, match="version"): + await read_document(record, storage) + assert storage.read_file.await_count == 1 + + +@pytest.mark.asyncio +async def test_smart_chip_display_text_is_used_without_following_link(): + native = document() + native["tabs"][0]["documentTab"]["body"]["content"] = [ + { + "paragraph": { + "elements": [ + { + "person": { + "personProperties": {"name": "Alice", "email": "private@example.com"} + } + }, + {"textRun": {"content": " and "}}, + {"person": {"personProperties": {"email": "bob@example.com"}}}, + {"textRun": {"content": " read "}}, + { + "richLink": { + "richLinkProperties": { + "title": "Roadmap", + "uri": "https://untrusted.example/never-fetch", + } + } + }, + ] + } + } + ] + record, storage, _ = capture(native) + result = await read_document(record, storage) + assert "Alice and bob@example.com read Roadmap" in result.text() + assert "untrusted.example" not in result.text() + assert "private@example.com" not in result.text() + assert storage.read_file.await_count == 2 + + +@pytest.mark.asyncio +async def test_equation_is_retained_but_not_invented_as_search_text(): + native = document() + native["tabs"][0]["documentTab"]["body"]["content"] = [ + {"paragraph": {"elements": [{"equation": {}}]}}, + ] + record, storage, _ = capture(native) + result = await read_document(record, storage) + assert result.document == native + with pytest.raises(DocumentReadError, match="equation text"): + result.text() + + +@pytest.mark.asyncio +async def test_unsupported_manifest_never_falls_back_to_export(): + record, storage, _ = capture(schema_version=2) + with pytest.raises(ValueError): + async with map_record(record, "google_drive", storage): + pass + assert storage.read_file.await_count == 1 + assert storage.read_file.call_args.args[0] == record.blobs[2].key + + +@pytest.mark.asyncio +async def test_incomplete_native_fields_are_not_published_as_complete_text(): + record, storage, _ = capture() + result = await read_document(record, storage) + native = result.manifest.native.model_copy(update={"status": "partial"}) + partial = result.model_copy( + update={"manifest": result.manifest.model_copy(update={"native": native})} + ) + assert partial.tab("child")["tabProperties"]["tabId"] == "child" + with pytest.raises(DocumentReadError, match="incomplete"): + partial.text() From 6927ca1f4b744c56f9029a0cca3a175c899f3b45 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:03:36 -0700 Subject: [PATCH 114/266] feat(records): authorize exact stored native document reads --- backend/airweave/api/v1/endpoints/records.py | 23 +++ .../domains/entities/canonical/query.py | 53 +++++ .../entities/canonical/query_models.py | 17 +- .../canonical/tests/test_document_http.py | 186 ++++++++++++++++++ 4 files changed, 278 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_document_http.py diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index b9f463b5f1..b56a13445e 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -16,6 +16,7 @@ from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query import CanonicalQueryService from airweave.domains.entities.canonical.query_models import ( + DocumentRead, MailThreadPage, RecordChangePage, RecordFilters, @@ -47,6 +48,8 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> "blob_not_found": 404, "stale_record_revision": 409, "blob_unavailable": 503, + "document_unavailable": 409, + "document_incomplete": 409, "calendar_changed_restart": 409, "calendar_read_incomplete": 409, "calendar_range_not_captured": 409, @@ -103,6 +106,26 @@ async def read_record( return await service.read(db, ctx.organization.id, sync_id, record_id) +@router.get("/{sync_id}/records/{record_id}/document", response_model=DocumentRead) +async def read_stored_document( + sync_id: UUID, + record_id: UUID, + response: Response, + revision: int = Query(ge=1), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), + container: Container = Depends(deps.get_container), +) -> DocumentRead: + """Read exact retained native Docs content; never fetch or fall back to a provider.""" + result = await service.document( + db, ctx.organization.id, sync_id, record_id, revision, container.storage_backend + ) + response.headers["Cache-Control"] = "private, no-store" + response.headers["X-Content-Type-Options"] = "nosniff" + return result + + @router.get("/{sync_id}/mail/threads/{thread_id}", response_model=MailThreadPage) async def mail_thread( sync_id: UUID, diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py index 645e3a6b6a..72e1eace76 100644 --- a/backend/airweave/domains/entities/canonical/query.py +++ b/backend/airweave/domains/entities/canonical/query.py @@ -10,6 +10,7 @@ from airweave.domains.entities.canonical.coverage import capture_coverage from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import ( + DocumentRead, MailThreadCursor, MailThreadPage, RecordChangePage, @@ -19,6 +20,7 @@ ) from airweave.domains.entities.canonical.query_store import CanonicalQueryStore from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.workspace_docs import read_document from airweave.domains.storage.exceptions import StorageException from airweave.domains.storage.protocols import StorageBackend @@ -53,6 +55,18 @@ class BlobUnavailable(CanonicalStoreError): code = "blob_unavailable" +class DocumentUnavailable(CanonicalStoreError): + """This record has no valid retained native document representation.""" + + code = "document_unavailable" + + +class DocumentIncomplete(CanonicalStoreError): + """Retained native document fields do not cover the requested document.""" + + code = "document_incomplete" + + class CanonicalQueryService: """Own read semantics; cursor signatures never replace tenant authorization.""" @@ -215,6 +229,45 @@ async def mail_thread( thread_id=thread_id, messages=messages, next_cursor=next_cursor, has_more=more ) + async def document( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + record_id: UUID, + revision: int, + storage: StorageBackend, + ) -> DocumentRead: + """Read verified stored Docs bytes under current record and ancestor authority.""" + record = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(record, revision) + try: + captured = await read_document(record, storage) + except (StorageException, BlobIntegrityError) as exc: + raise BlobUnavailable("Committed document bytes are unavailable; retry later") from exc + except ValueError as exc: + raise DocumentUnavailable( + "No valid stored Google document is available; refresh this source and reread." + ) from exc + db.expire_all() + current = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(current, revision) + if record.blobs != current.blobs: + raise StaleRecordRevision("Document changed; reread before requesting its content") + if captured.manifest.native.status != "complete": + raise DocumentIncomplete( + "Stored document content is incomplete; refresh this source and reread." + ) + return DocumentRead( + id=current.id, + identity=current.identity, + revision=current.revision, + observed_at=current.observed_at, + completeness=current.completeness, + manifest=captured.manifest, + document=captured.document, + ) + async def blob( self, db: AsyncSession, diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index 2d64cae6c2..204e30b753 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -4,10 +4,25 @@ from typing import Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict, Field +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue from airweave.domains.entities.canonical.coverage import CaptureCoverage from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.platform.sources.records.workspace_manifest import WorkspaceManifestV1 + + +class DocumentRead(BaseModel): + """One authorized stored revision's native document and representation coverage.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + id: UUID + identity: RecordIdentity + revision: int = Field(ge=1) + observed_at: AwareDatetime + completeness: Literal["complete", "partial", "metadata_only"] + manifest: WorkspaceManifestV1 + document: dict[str, JsonValue] class RecordFilters(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/tests/test_document_http.py b/backend/airweave/domains/entities/canonical/tests/test_document_http.py new file mode 100644 index 0000000000..4206b5159c --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_document_http.py @@ -0,0 +1,186 @@ +"""Stored Docs HTTP reads use real SQL visibility and verified synthetic blob bytes.""" + +import hashlib +import json +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 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, RecordIdentity +from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError +from airweave.domains.entities.canonical.tests.helpers import capture, observation +from airweave.domains.storage.exceptions import StorageNotFoundError + + +def document_capture(sync_id, variant): + contents = {} + + def blob(value, role=None): + content = json.dumps(value).encode() + digest = hashlib.sha256(content).hexdigest() + key = f"canonical/{sync_id}/blobs/sha256/{digest}" + contents[key] = content + return BlobReference( + key=key, + sha256=digest, + size_bytes=len(content), + media_type="application/json", + role=role, + ) + + document = { + "documentId": "wrong-file" if variant == "wrong_file" else "doc", + "title": "Retained document", + "futureField": {"preserve": [1, True]}, + "tabs": [{"tabProperties": {"tabId": "tab"}, "documentTab": {"body": {"content": []}}}], + } + native = blob(document) + manifest = blob( + { + "file_id": "doc", + "drive_version": "42", + "export": {"status": "unavailable", "reason": "unsupported"}, + "native": { + "status": "unavailable" + if variant == "native_unavailable" + else ("partial" if variant == "native_partial" else "complete"), + "document_blob": None if variant == "native_unavailable" else native.sha256, + "embedded_media": "not_retained", + "missing": [{"reason": "capture_budget"}] + if variant in {"native_partial", "native_unavailable"} + else [], + }, + }, + "representation_manifest", + ) + root = observation("folder", identity=RecordIdentity(record_type="folder", native_id="folder")) + record = observation( + "doc", + identity=RecordIdentity(record_type="file", native_id="doc"), + parent=root.identity, + payload={"id": "doc", "version": "42", "mimeType": "application/vnd.google-apps.document"}, + blobs=(native,) if variant == "missing_manifest" else (native, manifest), + completeness="partial", + ) + return root, record, contents, document + + +def document_storage(database, service, fence, root, document_record, contents, variant): + mutated = False + + async def read_bytes(key, *, max_bytes): + nonlocal mutated + if variant == "missing_bytes": + raise StorageNotFoundError("private-storage-key") + if variant == "corrupt_bytes": + return b"corrupt private content" + if not mutated: + mutated = True + update = None + if variant == "changed_during_read": + update = document_record.model_copy( + update={"payload": {**document_record.payload, "version": "43"}} + ) + elif variant == "deleted_during_read": + update = document_record.model_copy( + update={"kind": "delete", "removal_reason": "provider_deleted"} + ) + elif variant == "ancestor_revoked_during_read": + update = root.model_copy( + update={"kind": "delete", "removal_reason": "access_revoked"} + ) + if update is not None: + await capture(database, service, fence, update) + assert len(contents[key]) <= max_bytes + return contents[key] + + storage = AsyncMock() + storage.read_file.side_effect = read_bytes + return storage + + +@pytest.mark.parametrize( + "variant,status,code", + [ + ("available", 200, None), + ("missing_manifest", 409, "document_unavailable"), + ("wrong_file", 409, "document_unavailable"), + ("native_unavailable", 409, "document_unavailable"), + ("native_partial", 409, "document_incomplete"), + ("missing_bytes", 503, "blob_unavailable"), + ("corrupt_bytes", 503, "blob_unavailable"), + ("changed_during_read", 409, "stale_record_revision"), + ("deleted_during_read", 404, "record_not_found"), + ("ancestor_revoked_during_read", 404, "record_not_found"), + ("revoked_before_read", 404, "record_not_found"), + ("other_tenant", 404, "record_not_found"), + ("stale_revision", 409, "stale_record_revision"), + ("missing_revision", 422, None), + ], +) +async def test_document_read_authorization_integrity_and_native_contract( + database, source, variant, status, code +): + service, fence = source + root, document_record, contents, document = document_capture(fence.sync_id, variant) + saved = await capture(database, service, fence, root, document_record) + record_id = saved.changes[-1].record.id + if variant == "revoked_before_read": + await capture( + database, + service, + fence, + root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}), + ) + storage = document_storage(database, service, fence, root, document_record, contents, variant) + 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 + + owner = uuid4() if variant == "other_tenant" else fence.organization_id + app.dependency_overrides[deps.get_context] = lambda: SimpleNamespace( + organization=SimpleNamespace(id=owner) + ) + app.dependency_overrides[get_db] = session + app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(storage_backend=storage) + app.dependency_overrides[deps.get_canonical_query_service] = lambda: CanonicalQueryService( + CanonicalRecordStore(), CanonicalQueryStore(), "fixture" + ) + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.get( + f"/sync/{fence.sync_id}/records/{record_id}/document", + params={} + if variant == "missing_revision" + else {"revision": 2 if variant == "stale_revision" else 1}, + ) + assert response.status_code == status, response.text + if status == 200: + result = response.json() + assert result["document"] == document + assert result["id"] == str(record_id) + assert result["identity"]["native_id"] == "doc" + assert result["revision"] == 1 and result["observed_at"] + assert result["completeness"] == "partial" + assert result["manifest"]["native"]["status"] == "complete" + assert response.headers["cache-control"] == "private, no-store" + assert "key" not in result and "blobs" not in result + elif code: + assert response.json()["error"]["code"] == code + assert "private" not in response.text + assert "documentId" not in response.text + assert "canonical/" not in response.text + if variant in {"other_tenant", "revoked_before_read", "stale_revision", "missing_revision"}: + storage.read_file.assert_not_awaited() From 2fc0868a4ecac67c9c4dc65d285a8573e1e07d01 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:05:07 -0700 Subject: [PATCH 115/266] Capture native Docs and export representations under Drive ownership --- .../auth_provider/providers/composio.py | 2 +- .../canonical/tests/test_workspace_capture.py | 156 ++++++++++++++ .../airweave/platform/sources/google_drive.py | 2 +- .../sources/records/google_docs_content.py | 136 +++++++++++++ .../sources/records/google_drive_content.py | 52 ++++- .../http_client/test_composio_transport.py | 24 +++ .../records/test_google_docs_content.py | 190 ++++++++++++++++++ .../records/test_google_drive_content.py | 8 +- 8 files changed, 553 insertions(+), 17 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py create mode 100644 backend/airweave/platform/sources/records/google_docs_content.py create mode 100644 backend/tests/unit/platform/sources/records/test_google_docs_content.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 3409abe9d0..d73c86735e 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -114,7 +114,7 @@ async def get_auth_result( hosts = { "gmail": frozenset({"gmail.googleapis.com"}), "google_calendar": frozenset({"www.googleapis.com"}), - "google_drive": frozenset({"www.googleapis.com"}), + "google_drive": frozenset({"www.googleapis.com", "docs.googleapis.com"}), "slack": frozenset({"slack.com"}), "linear": frozenset({"api.linear.app", "uploads.linear.app"}), "github": frozenset({"api.github.com"}), 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..1837936e7b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py @@ -0,0 +1,156 @@ +"""Actual Drive pipeline publishes Docs references atomically; HTTP is synthetic.""" + +import hashlib +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 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 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 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/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py index 9394fdec08..25a892f8bf 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -102,7 +102,7 @@ def capture_cycle_configuration(self) -> CycleConfiguration: if self.include_patterns: raise ValueError("Drive path selection is not yet supported by canonical capture") return CycleConfiguration( - fingerprint=hashlib.sha256(b"drive-all-accessible-v1").hexdigest(), + fingerprint=hashlib.sha256(b"drive-all-accessible-docs-native-v2").hexdigest(), parents={"file": (None,)}, known_object_validation=("file",), ) 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_content.py b/backend/airweave/platform/sources/records/google_drive_content.py index 338756488b..5d0feaa041 100644 --- a/backend/airweave/platform/sources/records/google_drive_content.py +++ b/backend/airweave/platform/sources/records/google_drive_content.py @@ -5,13 +5,15 @@ import httpx from airweave.core.logging import ContextualLogger -from airweave.domains.entities.canonical.requests import CaptureRecord +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.workspace_manifest import DOCS_MIME, ExportState async def capture_file_content( @@ -44,28 +46,56 @@ async def capture_file_content( else: media_type = mime download_url = url + "?alt=media&supportsAllDrives=true" - try: - blob = await files.capture_canonical_url( - url=download_url, client=client, auth=auth, logger=logger, media_type=media_type + blob, export_state = await _download_representation( + download_url, media_type, mime, files=files, client=client, auth=auth, logger=logger + ) + if mime == DOCS_MIME: + record = await capture_document_parts( + record, + export=export_state, + export_blob=blob, + files=files, + client=client, + auth=auth, ) - except httpx.HTTPStatusError as error: - if not _export_limit_reached(error, mime): - raise - logger.info("Drive export exceeds provider size limit; retained metadata only") - return record - except FileSkippedException: - # Declared/streamed size limits are a visible metadata-only record, not absent data. + 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 mime == DOCS_MIME: + return record return record.model_copy( update={"blobs": (blob,), "content_hash": blob.sha256, "completeness": "complete"} ) +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."): diff --git a/backend/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py index 67cd97e0c5..3be38ffaf5 100644 --- a/backend/tests/unit/platform/http_client/test_composio_transport.py +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -308,3 +308,27 @@ async def aclose(self): import asyncio await asyncio.get_running_loop().shutdown_asyncgens() + + +@pytest.mark.asyncio +async def test_drive_binding_allows_native_docs_only_with_same_verified_toolkit(monkeypatch): + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "key"}, config={"account_id": "ca_drive"} + ) + monkeypatch.setattr( + provider, + "_get_with_auth", + AsyncMock( + return_value={ + "id": "ca_drive", + "toolkit": {"slug": "googledrive"}, + "status": "ACTIVE", + } + ), + ) + result = await provider.get_auth_result("google_drive", ["access_token"]) + assert result.managed_auth.connected_account_id == "ca_drive" + assert result.managed_auth.allowed_hosts == {"www.googleapis.com", "docs.googleapis.com"} + assert "sheets.googleapis.com" not in result.managed_auth.allowed_hosts diff --git a/backend/tests/unit/platform/sources/records/test_google_docs_content.py b/backend/tests/unit/platform/sources/records/test_google_docs_content.py new file mode 100644 index 0000000000..ffb8653a77 --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_google_docs_content.py @@ -0,0 +1,190 @@ +"""Real bounded HTTP acquisition and immutable filesystem storage; no provider calls.""" + +import json +from unittest.mock import AsyncMock, Mock +from uuid import uuid4 + +import httpx +import pytest + +from airweave.adapters.storage.filesystem import FilesystemBackend +from airweave.domains.sources.exceptions import ( + SourceAuthError, + SourceRateLimitError, + SourceServerError, +) +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.storage.file_service import FileService +from airweave.platform.sources.records import google_docs_content +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.workspace_manifest import parse_manifest + +DOCUMENT = { + "documentId": "doc", + "title": "All tabs", + "tabs": [ + { + "tabProperties": {"tabId": "one"}, + "documentTab": {"body": {"content": []}}, + "childTabs": [{"tabProperties": {"tabId": "two"}, "documentTab": {"unknown": "kept"}}], + } + ], +} + + +@pytest.fixture +def files(tmp_path, monkeypatch): + monkeypatch.setattr( + "airweave.domains.storage.file_service.paths.temp_sync_dir", + lambda _: str(tmp_path / "temp"), + ) + return FileService(uuid4(), FilesystemBackend(tmp_path / "storage"), sync_id=uuid4()) + + +async def acquire(files, handler, *, version="7"): + record = file_record( + {"id": "doc", "mimeType": "application/vnd.google-apps.document", "version": "7"} + ) + get = AsyncMock(return_value={"version": version}) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + result = await capture_file_content( + record, + files=files, + get=get, + client=client, + auth=StaticTokenProvider("test"), + logger=Mock(), + ) + return result, get + + +def response(request): + if request.url.host == "docs.googleapis.com": + assert request.url.params["includeTabsContent"] == "true" + assert request.url.params["commentsViewMode"] == "COMMENTS_VIEW_MODE_OMITTED" + assert request.url.params["suggestionsViewMode"] == "DEFAULT_FOR_CURRENT_ACCESS" + assert "fields" not in request.url.params + return httpx.Response(200, json=DOCUMENT) + return httpx.Response(200, content=b"synthetic-docx") + + +async def manifest_for(record, files): + marked = next(blob for blob in record.blobs if blob.role) + content = await files.storage.read_file(marked.key) + return parse_manifest(content, file_id="doc", drive_version="7", blobs=record.blobs) + + +async def test_all_tabs_native_and_export_are_published_together_and_deterministic(files): + first, get = await acquire(files, response) + second, _ = await acquire(files, response) + assert first.blobs == second.blobs and first.content_hash == second.content_hash + assert len(first.blobs) == 3 and first.completeness == "complete" + assert first.payload == { + "id": "doc", + "mimeType": "application/vnd.google-apps.document", + "version": "7", + } + manifest = await manifest_for(first, files) + native = next(blob for blob in first.blobs if blob.sha256 == manifest.native.document_blob) + assert json.loads(await files.storage.read_file(native.key)) == DOCUMENT + assert get.await_count == 1 + + +async def test_changed_drive_version_rejects_staged_representations(files): + with pytest.raises(ValueError, match="changed during"): + await acquire(files, response, version="8") + # Uncommitted immutable bytes can remain; no source blob GC is claimed. + assert len(await files.storage.list_files()) == 3 + + +async def test_native_success_survives_explicit_export_size_limit(files): + def limited(request): + if request.url.host == "docs.googleapis.com": + return response(request) + return httpx.Response( + 403, json={"error": {"errors": [{"reason": "exportSizeLimitExceeded"}]}} + ) + + record, _ = await acquire(files, limited) + manifest = await manifest_for(record, files) + assert record.completeness == "complete" + assert manifest.export.reason == "export_size_limit" and manifest.native.status == "complete" + + +async def test_embedded_objects_are_explicitly_partial_not_downloaded(files): + def media(request): + if request.url.host == "docs.googleapis.com": + return httpx.Response( + 200, + json={ + **DOCUMENT, + "inlineObjects": { + "image": { + "embeddedObject": { + "imageProperties": {"contentUri": "https://foreign.invalid/image"} + } + } + }, + }, + ) + return response(request) + + record, _ = await acquire(files, media) + manifest = await manifest_for(record, files) + assert record.completeness == "partial" + assert manifest.native.status == "complete" and manifest.native.embedded_media == "not_retained" + assert manifest.native.missing[0].source_path == "/inlineObjects/image/embeddedObject" + + +@pytest.mark.parametrize( + "status,error", [(401, SourceAuthError), (429, SourceRateLimitError), (500, SourceServerError)] +) +async def test_oversized_error_is_never_successful_native_gap(files, monkeypatch, status, error): + monkeypatch.setattr(google_docs_content, "MAX_DOCS_BYTES", 20) + + def failure(request): + if request.url.host == "docs.googleapis.com": + return httpx.Response(status, content=b"x" * 100, headers={"Retry-After": "300"}) + return response(request) + + with pytest.raises(error): + await acquire(files, failure) + + +async def test_successful_oversize_records_explicit_native_gap(files, monkeypatch): + monkeypatch.setattr(google_docs_content, "MAX_DOCS_BYTES", 20) + record, _ = await acquire(files, response) + manifest = await manifest_for(record, files) + assert record.completeness == "partial" + assert ( + manifest.native.document_blob is None + and manifest.native.missing[0].reason == "read_size_limit" + ) + + +async def test_foreign_document_and_storage_failure_do_not_publish(files): + def foreign(request): + if request.url.host == "docs.googleapis.com": + return httpx.Response(200, json={**DOCUMENT, "documentId": "foreign"}) + return response(request) + + with pytest.raises(ValueError, match="another file"): + await acquire(files, foreign) + files.store_canonical_blob = AsyncMock(side_effect=ConnectionError("storage unavailable")) + with pytest.raises(ConnectionError): + await acquire(files, response) + + +async def test_identical_part_bytes_share_one_workspace_descriptor_without_losing_uses(files): + from airweave.platform.sources.records.workspace_manifest import canonical_json + + def identical(request): + return httpx.Response(200, content=canonical_json(DOCUMENT)) + + record, _ = await acquire(files, identical) + manifest = await manifest_for(record, files) + assert len(record.blobs) == 2 + assert manifest.export.blob == manifest.native.document_blob + part = next(blob for blob in record.blobs if blob.role is None) + assert part.media_type == "application/octet-stream" diff --git a/backend/tests/unit/platform/sources/records/test_google_drive_content.py b/backend/tests/unit/platform/sources/records/test_google_drive_content.py index a03eebbb9f..41a6eaa910 100644 --- a/backend/tests/unit/platform/sources/records/test_google_drive_content.py +++ b/backend/tests/unit/platform/sources/records/test_google_drive_content.py @@ -62,16 +62,16 @@ async def test_storage_failure_is_not_silently_a_partial_success(inputs): await capture_file_content(record, **inputs) -async def test_native_document_uses_existing_export_format(inputs): +async def test_native_spreadsheet_keeps_existing_export_only_format(inputs): record = file_record( - {"id": "doc", "mimeType": "application/vnd.google-apps.document", "version": "7"} + {"id": "doc", "mimeType": "application/vnd.google-apps.spreadsheet", "version": "7"} ) await capture_file_content(record, **inputs) call = inputs["files"].capture_canonical_url.call_args.kwargs assert "/export?" in call["url"] assert ( call["media_type"] - == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" ) @@ -89,7 +89,7 @@ async def test_provider_export_limit_is_narrow_and_preserves_native_payload(inpu ) inputs["files"].capture_canonical_url.side_effect = failure record = file_record( - {"id": "doc", "mimeType": "application/vnd.google-apps.document", "version": "7"} + {"id": "doc", "mimeType": "application/vnd.google-apps.spreadsheet", "version": "7"} ) if reason != "exportSizeLimitExceeded": with pytest.raises(httpx.HTTPStatusError): From 85603b0da58afa989e0244c92baeb8923c4e82c6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:08:31 -0700 Subject: [PATCH 116/266] test: qualify stored Docs capture through fresh HTTP reader --- .../tests/test_workspace_document_probe.py | 87 +++++++++++++++++++ backend/tests/live/canonical_capture.py | 16 ++++ backend/tests/live/provider_sample.py | 5 +- backend/tests/live/workspace_document.py | 79 +++++++++++++++++ backend/tests/live/workspace_document_read.py | 71 +++++++++++++++ 5 files changed, 256 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py create mode 100644 backend/tests/live/workspace_document.py create mode 100644 backend/tests/live/workspace_document_read.py 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..2728491d5b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py @@ -0,0 +1,87 @@ +"""Exercise the opt-in Docs sample with fake provider HTTP and real SQL/HTTP/process reads.""" + +import json + +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':{'emailAddress':'synthetic@example.test'}}) + 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()) + await verify_rest_identity(name,connector,email) + 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"] diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index e1be77ea1e..f6dae31521 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -84,6 +84,11 @@ async def fetch_records(name, account, expected_email, key, fence, root): from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.entities.canonical.source import ContainerScopedSource + if name == "google_drive" and os.environ.get("LIVE_WORKSPACE_DOCUMENT") == "1": + from workspace_document import fetch_workspace_document + + return await fetch_workspace_document(account, expected_email, key, fence, root) + storage = FilesystemBackend(root / "blobs") files = FileService(fence.job_id, storage, sync_id=fence.sync_id) files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 @@ -235,7 +240,16 @@ async def verify(name, account, email, key, sessions, engine, root): sync_id=sync_id, root=root, ) + document_reader = None + if name == "google_drive" and os.environ.get("LIVE_WORKSPACE_DOCUMENT") == "1": + from workspace_document import verify_workspace_document_reader + + document_reader = await verify_workspace_document_reader( + records=listed, storage=storage, query=query, sessions=sessions, + organization_id=organization_id, sync_id=sync_id, root=root, + ) return { + "document_reader": document_reader, "almanac_reader": consumer, "provider": name, "identity_verification": identity_verification, @@ -269,6 +283,8 @@ async def main(): raise ValueError( "LIVE_PROVIDERS must select gmail, google_drive, google_calendar, slack or wispr" ) + if os.environ.get("LIVE_WORKSPACE_DOCUMENT") == "1" and selected != ["google_drive"]: + raise ValueError("Workspace document trial requires only google_drive") inputs = [ (name, os.environ[variable], os.environ["LIVE_EXPECTED_EMAIL"] if name != "wispr" else "") for name, variable in ( diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index e10db47fdf..3c62c680fa 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -72,8 +72,9 @@ async def rest_source( max_file_bytes=10 * 1024 * 1024, ): host = {"gmail": "gmail.googleapis.com", "slack": "slack.com"}.get(name, "www.googleapis.com") - auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts={host}) - transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts={host}) + hosts = {host, "docs.googleapis.com"} if name == "google_drive" else {host} + auth = ManagedAuthProvider(api_key=key, connected_account_id=account, allowed_hosts=hosts) + transport = ComposioTransport(api_key=key, connected_account_id=account, allowed_hosts=hosts) transport.MAX_BINARY_BYTES = max_file_bytes hooks = {"request": [request_hook]} if request_hook is not None else {} if response_hook is not None: diff --git a/backend/tests/live/workspace_document.py b/backend/tests/live/workspace_document.py new file mode 100644 index 0000000000..8bbf9194ad --- /dev/null +++ b/backend/tests/live/workspace_document.py @@ -0,0 +1,79 @@ +"""One explicit native Docs hydration sample; never a Drive enumeration or checkpoint.""" + +import asyncio +import hashlib +import sys +from pathlib import Path + +from provider_sample import rest_source + +from airweave.adapters.storage.filesystem import FilesystemBackend +from airweave.domains.storage.file_service import FileService +from airweave.platform.sources.records.google_drive import BASE, file_record +from airweave.platform.sources.records.workspace_manifest import canonical_json + + +async def fetch_workspace_document(account, email, key, fence, root): + storage = FilesystemBackend(root / "blobs") + files = FileService(fence.job_id, storage, sync_id=fence.sync_id) + files.MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 + requests = 0 + + async def count_request(request): + nonlocal requests + if request.method != "GET" or requests >= 12: + raise ValueError("Workspace read-only provider request budget exceeded") + requests += 1 + + async with ( + asyncio.timeout(150), + rest_source("google_drive", account, email, key, fence, request_hook=count_request) as ( + source, + identity_verification, + ), + ): + listing = await source._get( + BASE + "/files", + params={ + "q": "trashed=false and mimeType='application/vnd.google-apps.document'", + "pageSize": "1", + "fields": "files(*)", + }, + ) + items = listing.get("files", []) + if len(items) != 1: + raise ValueError("The selected account returned no document sample") + captured = await source._capture_body(file_record(items[0]), files) + if not any(blob.role == "representation_manifest" for blob in captured.blobs): + raise ValueError("The document sample has no structured representation") + # Fixed metadata only; never print native IDs, names or contents. + print('{"document_sample_provider_requests": ' + str(requests) + "}", flush=True) + return [captured], storage, identity_verification, 1 + + +async def verify_workspace_document_reader( + *, records, storage, query, sessions, organization_id, sync_id, root +): + from almanac_handoff import run_consumer + + if len(records) != 1: + raise ValueError("Workspace trial requires exactly one retained document") + record = records[0] + async with sessions() as db: + result = await query.document( + db, organization_id, sync_id, record.id, record.revision, storage + ) + expected = hashlib.sha256(canonical_json(result.document)).hexdigest() + backend = Path(__file__).resolve().parents[2] + return await run_consumer( + sessions=sessions, + organization_id=organization_id, + sync_id=sync_id, + root=root, + storage=storage, + query=query, + consumer=Path(__file__).with_name("workspace_document_read.py"), + repository=backend.parent, + python=Path(sys.executable), + extra={"record_id": str(record.id), "revision": record.revision, "expected": expected}, + ) diff --git a/backend/tests/live/workspace_document_read.py b/backend/tests/live/workspace_document_read.py new file mode 100644 index 0000000000..66922b1a86 --- /dev/null +++ b/backend/tests/live/workspace_document_read.py @@ -0,0 +1,71 @@ +"""Private subprocess read of the real stored-document HTTP route; no provider keys.""" + +import asyncio +import hashlib +import json +import logging +import sys +from pathlib import Path +from uuid import uuid4 + +import httpx + +logging.disable(logging.CRITICAL) + + +async def main(): + path = Path(sys.argv[1]).resolve() + if path.stat().st_mode & 0o077 or path.parent.stat().st_mode & 0o077: + raise ValueError("Private handoff permissions required") + config = json.loads(path.read_text()) + endpoint = f"/sync/{config['sync_id']}/records/{config['record_id']}/document" + async with httpx.AsyncClient( + base_url=config["origin"], + headers={"x-api-key": config["api_key"]}, + timeout=30, + trust_env=False, + follow_redirects=False, + ) as client: + response = await client.get(endpoint, params={"revision": config["revision"]}) + response.raise_for_status() + result = response.json() + assert result["id"] == config["record_id"] and result["revision"] == config["revision"] + encoded = json.dumps( + result["document"], + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode() + assert hashlib.sha256(encoded).hexdigest() == config["expected"] + assert result["manifest"]["file_id"] == result["document"]["documentId"] + assert result["manifest"]["native"]["status"] == "complete" + assert response.headers["cache-control"] == "private, no-store" + stale = await client.get(endpoint, params={"revision": config["revision"] + 1}) + assert stale.status_code == 409 + foreign = await client.get( + f"/sync/{uuid4()}/records/{config['record_id']}/document", + params={"revision": config["revision"]}, + ) + assert foreign.status_code == 404 + print( + json.dumps( + { + "verified": True, + "native_document_digest_verified": True, + "identity_verified": True, + "stale_revision_denied": True, + "wrong_source_denied": True, + "top_level_tabs": len(result["document"]["tabs"]), + "scope": "stored_document_http_route_with_fixture_auth", + } + ) + ) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except Exception: + print(json.dumps({"verified": False, "stage": "exact_record_read"})) + raise SystemExit(1) from None From 9b44828c06eb480f1274ead174deb0720d589e68 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:10:30 -0700 Subject: [PATCH 117/266] test: record bounded live Docs storage and HTTP verification --- .../evidence/docs-stored-http-20260930.json | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 backend/tests/live/evidence/docs-stored-http-20260930.json diff --git a/backend/tests/live/evidence/docs-stored-http-20260930.json b/backend/tests/live/evidence/docs-stored-http-20260930.json new file mode 100644 index 0000000000..3732417501 --- /dev/null +++ b/backend/tests/live/evidence/docs-stored-http-20260930.json @@ -0,0 +1,40 @@ +{ + "date": "2026-09-30", + "commit": "85603b0", + "provider": "google_drive", + "provider_requests": 5, + "identity_verification": "provider_email", + "records": 1, + "journal_changes": 1, + "partial_records": 0, + "blobs": 3, + "blob_bytes": 1053473, + "blob_sha_verified": true, + "replay_new_changes": 0, + "fresh_connections_readback": true, + "document_reader": { + "verified": true, + "native_document_digest_verified": true, + "identity_verified": true, + "stale_revision_denied": true, + "wrong_source_denied": true, + "top_level_tabs": 4, + "scope": "stored_document_http_route_with_fixture_auth" + }, + "full_scope_completed": false, + "checkpoint_saved": false, + "almanac_reader": null, + "harness_exit_code": 0, + "cleanup": { + "remaining_private_schemas": 0, + "remaining_private_blob_directories": 0, + "immutable_archive_removed": true + }, + "limitations": [ + "One selected document, not full Drive enumeration or incremental-cycle proof", + "Local HTTP route with fixture authentication, not hosted ownership verification", + "No Almanac document CLI consumer tested", + "No induced provider mutation or deletion", + "Sheets not tested" + ] +} From 69d02962072f24b2ecce9986401a1de08ac2d2c2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:24:56 -0700 Subject: [PATCH 118/266] test: connect document capture probe to the Almanac CLI consumer --- .../tests/test_workspace_document_probe.py | 31 +++++++++++++++++-- backend/tests/live/workspace_document.py | 16 ++++++++-- 2 files changed, 42 insertions(+), 5 deletions(-) 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 index 2728491d5b..edf7b6940a 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py @@ -1,12 +1,14 @@ """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''' +SCRIPT = r""" import asyncio,json,os,sys from contextlib import asynccontextmanager from pathlib import Path @@ -66,7 +68,7 @@ async def main(): finally: await engine.dispose() asyncio.run(main()) -''' +""" async def test_workspace_sample_survives_sql_reopen_and_real_http_reader(database, tmp_path): @@ -85,3 +87,28 @@ async def test_workspace_sample_survives_sql_reopen_and_real_http_reader(databas 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/tests/live/workspace_document.py b/backend/tests/live/workspace_document.py index 8bbf9194ad..2fe8b156d2 100644 --- a/backend/tests/live/workspace_document.py +++ b/backend/tests/live/workspace_document.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import os import sys from pathlib import Path @@ -65,6 +66,15 @@ async def verify_workspace_document_reader( ) expected = hashlib.sha256(canonical_json(result.document)).hexdigest() backend = Path(__file__).resolve().parents[2] + consumer = Path(__file__).with_name("workspace_document_read.py") + repository = backend.parent + python = Path(sys.executable) + if os.environ.get("LIVE_ALMANAC_PYTHON"): + python = Path(os.environ["LIVE_ALMANAC_PYTHON"]).absolute() + repository = Path(os.environ["LIVE_ALMANAC_ROOT"]).resolve() + consumer = repository / "backend/tests/live/read_owned_document.py" + if not python.is_file() or not consumer.is_file(): + raise ValueError("Explicit Almanac Python and document consumer required") return await run_consumer( sessions=sessions, organization_id=organization_id, @@ -72,8 +82,8 @@ async def verify_workspace_document_reader( root=root, storage=storage, query=query, - consumer=Path(__file__).with_name("workspace_document_read.py"), - repository=backend.parent, - python=Path(sys.executable), + consumer=consumer, + repository=repository, + python=python, extra={"record_id": str(record.id), "revision": record.revision, "expected": expected}, ) From 8927e481c5f0044d3d7d988e941d1aab0c379072 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:26:56 -0700 Subject: [PATCH 119/266] test: record real Docs capture through Almanac CLI --- .../evidence/docs-almanac-cli-20260930.json | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 backend/tests/live/evidence/docs-almanac-cli-20260930.json diff --git a/backend/tests/live/evidence/docs-almanac-cli-20260930.json b/backend/tests/live/evidence/docs-almanac-cli-20260930.json new file mode 100644 index 0000000000..036ebbeb48 --- /dev/null +++ b/backend/tests/live/evidence/docs-almanac-cli-20260930.json @@ -0,0 +1,36 @@ +{ + "date": "2026-09-30", + "airweave_commit": "69d0296", + "almanac_commit": "8ec369087", + "provider_requests": 5, + "records": 1, + "journal_changes": 1, + "blobs": 3, + "blob_bytes": 1053474, + "partial_records": 0, + "blob_sha_verified": true, + "replay_new_changes": 0, + "fresh_connections_readback": true, + "native_document_digest_verified": true, + "top_level_tabs": 4, + "identity_verified": true, + "stale_revision_denied": true, + "wrong_account_denied": true, + "revoked_binding_denied": true, + "scope": "actual_cli_and_almanac_fixture_route_to_real_store_http", + "full_scope_completed": false, + "checkpoint_saved": false, + "harness_exit_code": 0, + "cleanup": { + "private_schemas": 0, + "private_blob_directories": 0, + "immutable_archives_removed": true + }, + "limitations": [ + "One selected document, not full source traversal or checkpoint", + "Almanac route/auth/binding use fixtures; source-store uses real loopback HTTP", + "No hosted ownership proof", + "No induced provider edits/deletes", + "No Sheets capture proof" + ] +} From d6c8d3e6502e811b6d8d490624f4721b8fc0199b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:30:24 -0700 Subject: [PATCH 120/266] fix: reject ambiguous Wispr text continuation markers --- backend/airweave/platform/sources/wispr.py | 16 ++++++++----- .../platform/sources/test_wispr_capture.py | 23 +++++++++++++++++++ 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 264f2145ff..7464c801e2 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -192,13 +192,17 @@ async def validate(self) -> None: def _next_offset(text: JsonValue, field: str) -> int | None: if not isinstance(text, str): raise ValueError("Wispr meeting text is not a string") - match = re.search( - rf"\(\.\.\.truncated, \d+ chars remaining; continue with " - rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", - text, - re.MULTILINE, + matches = list( + re.finditer( + rf"\(\.\.\.truncated, \d+ chars remaining; continue with " + rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", + text, + re.MULTILINE, + ) ) - return int(match.group(1)) if match else None + if len(matches) > 1: + raise ValueError("Wispr returned ambiguous continuation markers") + return int(matches[0].group(1)) if matches else None async def _meeting(self, identity: str) -> dict[str, JsonValue]: offsets = {"content": 0, "transcript": 0} diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index e304827f1b..b3b57b1f4f 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -263,3 +263,26 @@ async def test_short_listing_pages_are_not_limited_to_six_cursors(monkeypatch): ) continuation = page.continuation assert page.final + + +@pytest.mark.asyncio +async def test_ambiguous_continuation_markers_do_not_choose_a_range(monkeypatch): + connector = await source() + transcript = ( + "Quoted example:\n(...truncated, 10 chars remaining; " + "continue with view_transcript.start_char=10...)\n" + "Actual meeting text follows.\n(...truncated, 20 chars remaining; " + "continue with view_transcript.start_char=40000...)\nProvider guidance follows." + ) + execute = AsyncMock( + return_value={ + "id": "m", + "modified_at": "same", + "content": "notes", + "transcript": transcript, + } + ) + monkeypatch.setattr(connector, "_execute", execute) + with pytest.raises(ValueError, match="ambiguous continuation"): + await connector._meeting("m") + assert execute.await_count == 1 From 1e5d1c6e9e4d914335a2892842a4cb5d6c0727b5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:34:05 -0700 Subject: [PATCH 121/266] fix: preserve safe Wispr rate signals without invented retry timing --- .../canonical/tests/test_wispr_diagnostics.py | 11 ++++++++-- .../canonical/tests/test_wispr_recovery.py | 18 +++++++++++----- backend/airweave/platform/sources/wispr.py | 11 ++++++++++ .../platform/sources/test_wispr_capture.py | 21 +++++++++++++++++++ 4 files changed, 54 insertions(+), 7 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py index d9b1fb669e..26cb650819 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_diagnostics.py @@ -13,6 +13,7 @@ async def test_safe_validation_failure_metadata_without_payload_leak(): import httpx import provider_sample from wispr_diagnostics import WisprDiagnostics +from airweave.domains.sources.exceptions import SourceServerError async def main(): original_client = httpx.AsyncClient @@ -42,8 +43,14 @@ async def request(req): diagnostics.request(req) SimpleNamespace(organization_id=uuid4()), request_hook=request, response_hook=diagnostics.response, envelope_hook=diagnostics.envelope): raise AssertionError("Validation error swallowed") - except ValueError as exc: - assert str(exc) == "Wispr tool execution failed; capture is incomplete" + 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 diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py index 15ad646ca1..c02eebabf5 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py @@ -9,6 +9,7 @@ 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 @@ -16,7 +17,7 @@ from airweave.platform.sources.wispr import WisprSource -async def connector(rows, body_calls, *, fail_at=None): +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" @@ -31,7 +32,7 @@ async def execute(slug, arguments): native = arguments["meeting_id"] body_calls.append(native) if len(body_calls) == fail_at: - raise ValueError("Synthetic opaque body failure") + raise failure or ValueError("Synthetic opaque body failure") return { "id": native, "content": "notes", @@ -50,15 +51,22 @@ def driver(service, database, fence, source): ) -async def test_failed_body_resumes_without_repeating_completed_siblings(database, source): +@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) - with pytest.raises(ValueError, match="Synthetic opaque"): + 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] diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 7464c801e2..0f4867f324 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -20,6 +20,7 @@ RecordIdentity, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation +from airweave.domains.sources.exceptions import SourceServerError from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider from airweave.domains.storage.file_service import FileService from airweave.platform.decorators import source @@ -178,6 +179,16 @@ async def _execute(self, slug: str, arguments: dict) -> dict[str, JsonValue]: }, ) ) + 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 diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index b3b57b1f4f..c248608d8d 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -4,6 +4,8 @@ import pytest +from airweave.domains.sources.exceptions import SourceServerError +from airweave.domains.sources.exceptions.classifier import classify_error from airweave.domains.sources.token_providers.protocol import ManagedToolAuthProvider from airweave.platform.sources.wispr import WisprSource @@ -286,3 +288,22 @@ async def test_ambiguous_continuation_markers_do_not_choose_a_range(monkeypatch) with pytest.raises(ValueError, match="ambiguous continuation"): await connector._meeting("m") assert execute.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error", ["Private meeting: rate limit exceeded", "Too many requests: private detail"] +) +async def test_rate_signal_stops_without_retry_or_invented_status(monkeypatch, error): + connector = await source() + post = AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": error}]) + monkeypatch.setattr(connector, "_post", post) + with pytest.raises(SourceServerError, match="rate-limit signal") as failure: + await connector._meeting("m") + assert post.await_count == 2 + assert failure.value.source_short_name == "wispr" + assert failure.value.status_code is None + assert "private" not in str(failure.value).lower() + assert "retry delay are unknown" in str(failure.value) + # Do not turn unstructured tool text into credential failure or a timed retry. + assert classify_error(failure.value).category is None From 9f34ade81e5dad103ef3bf3bbf94c4cb2553806a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:55:47 -0700 Subject: [PATCH 122/266] Capture native Attio records and memberships with durable note pagination --- .../canonical/tests/test_attio_recovery.py | 138 ++++ backend/airweave/platform/configs/config.py | 4 +- backend/airweave/platform/sources/attio.py | 691 ++++++------------ .../airweave/platform/sources/attio_models.py | 57 ++ .../platform/sources/test_attio_capture.py | 230 ++++++ 5 files changed, 652 insertions(+), 468 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_attio_recovery.py create mode 100644 backend/airweave/platform/sources/attio_models.py create mode 100644 backend/tests/unit/platform/sources/test_attio_capture.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_attio_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_attio_recovery.py new file mode 100644 index 0000000000..6c4c843767 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_attio_recovery.py @@ -0,0 +1,138 @@ +"""Actual Attio HTTP adapter and PostgreSQL page recovery with synthetic provider data.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID, uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import CompleteCycle +from airweave.domains.sources.exceptions import SourceError +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 AttioConfig +from airweave.platform.sources.attio import AttioSource + +WORKSPACE, OBJECT, RECORD = [str(UUID(int=i)) for i in range(1, 4)] +OBJECT_DATA = {"id": {"workspace_id": WORKSPACE, "object_id": OBJECT}, "api_slug": "people"} +RECORD_DATA = { + "id": {"workspace_id": WORKSPACE, "object_id": OBJECT, "record_id": RECORD}, + "values": {}, +} +NOTES = [ + { + "id": {"workspace_id": WORKSPACE, "note_id": str(UUID(int=i + 100))}, + "parent_object": "people", + "parent_record_id": RECORD, + "content_markdown": "Original note", + "unknown": {"retained": True}, + } + for i in range(51) +] + + +async def connector(calls, *, fail=False): + def result(path, data, status=200): + return httpx.Response(status, json=data, request=httpx.Request("GET", path)) + + async def get(path, **kwargs): + calls.append((path, kwargs.get("params"))) + if path.endswith("/self"): + return result(path, {"active": True, "workspace_id": WORKSPACE}) + if path.endswith("/objects"): + return result(path, {"data": [OBJECT_DATA]}) + if path.endswith("/lists"): + return result(path, {"data": []}) + if path.endswith("/objects/" + OBJECT): + return result(path, {"data": OBJECT_DATA}) + assert path.endswith("/notes") + offset = kwargs["params"]["offset"] + if fail and offset == 50: + return result(path, {"code": "merge_in_progress"}, 404) + return result(path, {"data": NOTES[offset : offset + 50]}) + + async def post(path, **kwargs): + calls.append((path, kwargs["json"])) + assert path.endswith("/records/query") + return result(path, {"data": [RECORD_DATA]}) + + client = AsyncMock() + client.get.side_effect = get + client.post.side_effect = post + return await AttioSource.create( + auth=ManagedAuthProvider( + api_key="fixture", + connected_account_id="fixture", + allowed_hosts=frozenset({"api.attio.com"}), + ), + logger=MagicMock(), + http_client=client, + config=AttioConfig(workspace_id=WORKSPACE), + ) + + +def driver(service, database, fence, connector): + return CanonicalScanDriver( + service, database, fence, connector, AsyncMock(), AsyncMock(), MagicMock() + ) + + +async def test_note_page_failure_recovers_exact_offset_without_replaying_fifty_notes( + database, source +): + service, fence = source + calls = [] + first = await connector(calls, fail=True) + with pytest.raises(SourceError): + 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 == "note") + ) + ).all() + ) + assert len(notes) == 50 + assert all(row.deleted_at is None and row.record_revision == 1 for row in notes) + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "note")) + assert scan.phase == "collecting" and scan.continuation["offset"] == 50 + before = await service.read_cycle(db, fence) + assert before.phase == "active" + 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, + ) + retry_calls = [] + resumed = await connector(retry_calls) + after = await driver(service, database, newer, resumed).run() + async with database() as db: + after = await service.complete_cycle(db, CompleteCycle(fence=newer, expected=after.version)) + assert after.phase == "complete" and after.version.cycle_id == before.version.cycle_id + assert after.last_full_capture.discovery == "incomplete" + assert [params["offset"] for path, params in retry_calls if path.endswith("/notes")] == [50] + # Discovery-only ancestors are refreshed on resume; the committed note page is not. + assert any(path.endswith("/records/query") for path, _ in retry_calls) + assert all( + after.configuration.policy(kind) == "discovery_only" + for kind in resumed.canonical_record_types + ) + async with database() as db: + notes = list( + ( + await db.scalars( + select(Entity).where(Entity.entity_definition_short_name == "note") + ) + ).all() + ) + assert len(notes) == 51 + assert all(row.deleted_at is None and row.record_revision == 1 for row in notes) + assert all(row.source_payload["unknown"] == {"retained": True} for row in notes) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 22cf473231..c7f0edb315 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -35,9 +35,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): diff --git a/backend/airweave/platform/sources/attio.py b/backend/airweave/platform/sources/attio.py index 1d915a09d3..f0f4e29e48 100644 --- a/backend/airweave/platform/sources/attio.py +++ b/backend/airweave/platform/sources/attio.py @@ -1,46 +1,54 @@ -"""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.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 @@ -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/tests/unit/platform/sources/test_attio_capture.py b/backend/tests/unit/platform/sources/test_attio_capture.py new file mode 100644 index 0000000000..a9a19a43c1 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_attio_capture.py @@ -0,0 +1,230 @@ +"""Native Attio capture boundaries; credential-free HTTP fixtures.""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID + +import httpx +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import 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 ManagedAuthProvider +from airweave.platform.configs.config import AttioConfig +from airweave.platform.sources.attio import AttioSource + +WORKSPACE, OBJECT, RECORD, LIST = [str(UUID(int=i)) for i in range(1, 5)] + + +def response(data, status=200): + return httpx.Response( + status, json=data, request=httpx.Request("GET", "https://api.attio.com/v2/notes") + ) + + +def native(kind, value, **fields): + return {"id": {"workspace_id": WORKSPACE, kind + "_id": value}, **fields} + + +def parent(kind, value, container=None): + return SourceRecord( + id=UUID(int=90), + sync_id=UUID(int=91), + identity=RecordIdentity(record_type=kind, native_id=value, container_id=container), + revision=1, + payload={}, + payload_schema_version=1, + capture_hash="0" * 64, + 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, + ) + + +async def source(*pages, principal=None): + client = AsyncMock() + client.get.side_effect = [ + response( + principal + if principal is not None + else { + "active": True, + "workspace_id": WORKSPACE, + "token_level": "user", + "authorized_by_workspace_member_id": str(UUID(int=7)), + } + ), + *pages, + ] + result = await AttioSource.create( + auth=ManagedAuthProvider( + api_key="test-only", + connected_account_id="ca_test", + allowed_hosts=frozenset({"api.attio.com"}), + ), + logger=MagicMock(), + http_client=client, + config=AttioConfig(workspace_id=WORKSPACE), + ) + return result, client + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "principal", + [{"active": True}, {"active": True, "workspace_id": str(UUID(int=999))}], +) +async def test_identity_requires_active_matching_workspace(principal): + with pytest.raises(ValueError): + await source(principal=principal) + + +@pytest.mark.asyncio +async def test_native_entries_are_distinct_memberships(): + capture, client = await source() + for number in (10, 11): + list_id = str(UUID(int=number)) + owner = parent("list", list_id) + entry = native( + "entry", + str(UUID(int=number + 100)), + parent_record_id=RECORD, + parent_object="people", + entry_values={"stage": [{"status": "new"}]}, + ) + entry["id"]["list_id"] = list_id + client.post.return_value = response({"data": [entry]}) + page = await capture.capture_page( + capture.child_scope(owner, "entry"), ScanContinuation(), parent=owner, files=MagicMock() + ) + assert page.records[0].payload == entry + assert page.records[0].identity.native_id != RECORD + assert page.records[0].identity.container_id == list_id + assert page.records[0].parent == owner.identity + assert client.post.call_args.kwargs["headers"] == {} + assert all( + capture.capture_cycle_configuration.policy(k) == "discovery_only" + for k in capture.canonical_record_types + ) + + +@pytest.mark.asyncio +async def test_more_than_fifty_notes_resume_and_preserve_native_content(): + obj = response({"data": native("object", OBJECT, api_slug="people")}) + notes = [ + native( + "note", + str(UUID(int=i + 100)), + parent_object="people", + parent_record_id=RECORD, + content_plaintext="body", + content_markdown="**body**", + future={"kept": True}, + ) + for i in range(51) + ] + capture, client = await source( + obj, response({"data": notes[:50]}), obj, response({"data": notes[50:]}) + ) + owner = parent("record", RECORD, OBJECT) + scope = capture.child_scope(owner, "note") + page = await capture.capture_page(scope, ScanContinuation(), files=MagicMock(), parent=owner) + assert not page.final + assert len(page.records) == 50 + # A reconstructed source uses only persisted continuation, not in-memory iteration state. + resumed, resumed_client = await source(obj, response({"data": notes[50:]})) + last = await resumed.capture_page(scope, page.continuation, files=MagicMock(), parent=owner) + assert last.final and last.records[0].payload == notes[-1] + assert resumed_client.get.call_args.kwargs["params"]["offset"] == 50 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code,status", [("merge_in_progress", 404), ("not_found", 404), ("forbidden", 403)] +) +async def test_note_failure_does_not_finalize_or_hide_error(code, status): + capture, client = await source( + response({"data": native("object", OBJECT, api_slug="people")}), + response({"code": code}, status), + ) + owner = parent("record", RECORD, OBJECT) + with pytest.raises(SourceError): + await capture.capture_page( + capture.child_scope(owner, "note"), ScanContinuation(), files=MagicMock(), parent=owner + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fault", ["workspace", "object", "duplicate"]) +async def test_invalid_record_page_rejected_atomically(fault): + capture, client = await source() + raw = native("record", RECORD) + raw["id"]["object_id"] = OBJECT + if fault in {"workspace", "object"}: + raw["id"][fault + "_id"] = str(UUID(int=999)) + client.post.return_value = response({"data": [raw, raw] if fault == "duplicate" else [raw]}) + owner = parent("object", OBJECT) + with pytest.raises(ValueError): + await capture.capture_page( + capture.child_scope(owner, "record"), + ScanContinuation(), + files=MagicMock(), + parent=owner, + ) + + +@pytest.mark.asyncio +async def test_previous_page_overlap_fails_and_absence_never_inferred(): + capture, client = await source() + owner = parent("object", OBJECT) + raw = native("record", RECORD) + raw["id"]["object_id"] = OBJECT + client.post.return_value = response({"data": [raw]}) + with pytest.raises(ValueError, match="repeated identities"): + await capture.capture_page( + capture.child_scope(owner, "record"), + ScanContinuation(value={"offset": 500, "previous_ids": [RECORD]}), + files=MagicMock(), + parent=owner, + ) + with pytest.raises(ValueError, match="cannot confirm absence"): + await capture.confirm_absent(owner) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [429, 500]) +async def test_transient_note_failure_propagates_after_retry_budget(status): + from types import MethodType + + from tenacity import stop_after_attempt + + capture, client = await source( + response({"data": native("object", OBJECT, api_slug="people")}), + response({"code": "transient"}, status), + ) + # Exercise production translation with one attempt, without waiting on fixture failures. + capture._request = MethodType( + AttioSource._request.retry_with(stop=stop_after_attempt(1)), capture + ) + owner = parent("record", RECORD, OBJECT) + with pytest.raises(SourceError): + await capture.capture_page( + capture.child_scope(owner, "note"), ScanContinuation(), files=MagicMock(), parent=owner + ) + + +@pytest.mark.asyncio +async def test_inactive_http200_is_reconnect_error(): + with pytest.raises(SourceAuthError) as error: + await source(principal={"active": False}) + assert error.value.status_code == 200 + assert "reconnect" in str(error.value) From 11bd8cc15af023bc7051f361f73610c7a12058f3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 12:56:07 -0700 Subject: [PATCH 123/266] feat: project retained Attio originals with managed account binding --- Makefile | 2 +- .../auth_provider/providers/composio.py | 1 + .../entities/canonical/attio_projection.py | 162 ++++++++++++++++ .../entities/canonical/projection_mappers.py | 4 + .../airweave/platform/entities/__init__.py | 14 +- backend/airweave/platform/entities/attio.py | 177 ++---------------- .../domains/entities/test_attio_projection.py | 171 +++++++++++++++++ .../http_client/test_composio_transport.py | 25 +++ 8 files changed, 383 insertions(+), 173 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/attio_projection.py create mode 100644 backend/tests/unit/domains/entities/test_attio_projection.py diff --git a/Makefile b/Makefile index 72c5a051a4..3a2ec7a9fa 100644 --- a/Makefile +++ b/Makefile @@ -35,7 +35,7 @@ test-store: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/syncs/tests test-capture: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities/test_github_projection.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities/test_github_projection.py tests/unit/domains/entities/test_attio_projection.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index d73c86735e..0ee9a96402 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -119,6 +119,7 @@ async def get_auth_result( "linear": frozenset({"api.linear.app", "uploads.linear.app"}), "github": frozenset({"api.github.com"}), "notion": frozenset({"api.notion.com"}), + "attio": frozenset({"api.attio.com"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( diff --git a/backend/airweave/domains/entities/canonical/attio_projection.py b/backend/airweave/domains/entities/canonical/attio_projection.py new file mode 100644 index 0000000000..eafb9f9843 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/attio_projection.py @@ -0,0 +1,162 @@ +"""Offline search views of Attio originals; membership never owns CRM record content.""" + +import json +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, HttpUrl, JsonValue + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import ProjectionMappingError +from airweave.platform.entities.attio import AttioOriginalEntity + + +class _Native(BaseModel): + model_config = ConfigDict(extra="ignore") + + +class _Identity(_Native): + workspace_id: UUID + object_id: UUID | None = None + record_id: UUID | None = None + list_id: UUID | None = None + entry_id: UUID | None = None + note_id: UUID | None = None + + +class _Original(_Native): + id: _Identity + + +class _Object(_Original): + singular_noun: str + plural_noun: str + api_slug: str + + +class _List(_Original): + name: str + + +class _Record(_Original): + web_url: HttpUrl | None = None + values: dict[str, JsonValue] + + +class _Entry(_Original): + parent_record_id: UUID + parent_object: str + entry_values: dict[str, JsonValue] + + +class _Note(_Original): + title: str + parent_record_id: UUID + parent_object: str + content_plaintext: str + content_markdown: str + + +class _Name(_Native): + full_name: str | None = None + value: str | None = None + + +def _attributes(values: dict[str, JsonValue]) -> str: + # Arbitrary custom attributes are data, not guessed standard properties. Preserve + # their names, values and references without fetching or joining other originals. + return json.dumps(values, ensure_ascii=False, sort_keys=True, indent=2) + + +def _parent(record: SourceRecord, kind: str, native_id: UUID | None) -> None: + if ( + native_id is None + or record.identity.container_id != str(native_id) + or record.parent is None + or record.parent.record_type != kind + or record.parent.native_id != str(native_id) + or (kind in {"object", "list"} and record.parent.container_id is not None) + ): + raise ProjectionMappingError("Attio parent differs from captured identity") + + +def _url(url: HttpUrl | None) -> str | None: + if url is None: + return None + if ( + url.scheme != "https" + or url.host != "app.attio.com" + or url.port not in (None, 443) + or url.username + or url.password + ): + raise ProjectionMappingError("Attio URL is outside the native application host") + return str(url) + + +def _record_title(original: _Record, native_id: UUID) -> str: + title = f"Record {native_id}" + names = original.values.get("name") + if isinstance(names, list) and names: + name = _Name.model_validate(names[0]) + title = name.full_name or name.value or title + return title + + +def map_attio(record: SourceRecord) -> tuple[AttioOriginalEntity, ...]: + """Render each native identity separately, without inventing URLs or related content.""" + native = _Original.model_validate(record.payload) + kind = record.identity.record_type + native_url = None + ids = native.id + native_id = { + "object": ids.object_id, + "record": ids.record_id, + "list": ids.list_id, + "entry": ids.entry_id, + "note": ids.note_id, + }.get(kind) + if native_id is None or str(native_id) != record.identity.native_id: + raise ProjectionMappingError("Attio native identity differs from captured identity") + if kind in {"object", "list"}: + if record.parent is not None or record.identity.container_id is not None: + raise ProjectionMappingError("Attio root has an unexpected capture parent") + if kind == "object": + obj = _Object.model_validate(record.payload) + title, text = obj.plural_noun, f"{obj.singular_noun}\n{obj.api_slug}" + else: + title, text = _List.model_validate(record.payload).name, "" + coverage = "Container metadata only; records and memberships are separate originals." + elif kind == "record": + _parent(record, "object", ids.object_id) + original = _Record.model_validate(record.payload) + native_url = _url(original.web_url) + title = _record_title(original, native_id) + text = _attributes(original.values) + coverage = "Retained attributes only; notes, memberships and file contents are separate." + elif kind == "entry": + _parent(record, "list", ids.list_id) + entry = _Entry.model_validate(record.payload) + title = f"List entry {native_id}" + text = f"Referenced {entry.parent_object} record: {entry.parent_record_id}\n" + text += _attributes(entry.entry_values) + coverage = "List membership and its attributes only; referenced CRM record is separate." + elif kind == "note": + note = _Note.model_validate(record.payload) + _parent(record, "record", note.parent_record_id) + title, text = note.title, note.content_plaintext + coverage = "Retained note text only; embedded media bytes are not included." + else: + raise ProjectionMappingError("Unsupported Attio original kind") + return ( + AttioOriginalEntity( + native_id=str(native_id), + original_kind=kind, + title=title or kind.capitalize(), + text=text, + content_coverage=coverage, + native_url=native_url, + breadcrumbs=[], + created_at=record.source_created_at, + updated_at=record.source_updated_at, + ), + ) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 0224e00f54..3adfd1a947 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -244,6 +244,10 @@ async def map_record( from airweave.domains.entities.canonical.linear_projection import map_linear entities = await map_linear(record, storage, directory) + elif source_name == "attio": + from airweave.domains.entities.canonical.attio_projection import map_attio + + entities = map_attio(record) elif source_name == "notion": from airweave.domains.entities.canonical.notion_projection import map_notion diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index 811f37f5fd..dfe2f926cc 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, @@ -388,12 +383,7 @@ AsanaTaskEntity, AsanaWorkspaceEntity, ], - "attio": [ - AttioListEntity, - AttioNoteEntity, - AttioObjectEntity, - AttioRecordEntity, - ], + "attio": [AttioOriginalEntity], "bitbucket": [ BitbucketCodeFileEntity, BitbucketDirectoryEntity, 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/tests/unit/domains/entities/test_attio_projection.py b/backend/tests/unit/domains/entities/test_attio_projection.py new file mode 100644 index 0000000000..0a1b2b0ae0 --- /dev/null +++ b/backend/tests/unit/domains/entities/test_attio_projection.py @@ -0,0 +1,171 @@ +"""Attio search preserves original identity rather than joining CRM and membership.""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import ( + ProjectionMappingError, + map_record, +) +from airweave.domains.entities.canonical.requests import RecordIdentity + +WORKSPACE, OBJECT, RECORD, LIST, ENTRY, NOTE = (str(uuid4()) for _ in range(6)) + + +def original(kind, native_id, payload, parent=None): + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity( + record_type=kind, native_id=native_id, container_id=parent.native_id if parent else None + ), + parent=parent, + revision=1, + payload=payload, + payload_schema_version=1, + 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, + ) + + +def entry(): + return original( + "entry", + ENTRY, + { + "id": {"workspace_id": WORKSPACE, "list_id": LIST, "entry_id": ENTRY}, + "parent_record_id": RECORD, + "parent_object": "people", + "entry_values": {"status": [{"status": {"title": "Interested"}}]}, + "unknown_native": "must survive", + }, + RecordIdentity(record_type="list", native_id=LIST), + ) + + +@pytest.mark.asyncio +async def test_membership_search_keeps_entry_identity_and_never_fetches_record(): + captured = entry() + before = captured.model_dump() + storage = AsyncMock() + async with map_record(captured, "attio", storage) as entities: + assert len(entities) == 1 + result = entities[0] + assert result.native_id == ENTRY + assert result.original_kind == "entry" + assert RECORD in result.text and "Interested" in result.text + assert result.web_url == "" + assert "separate" in result.content_coverage + assert captured.model_dump() == before + assert storage.mock_calls == [] + + +@pytest.mark.asyncio +async def test_note_uses_current_plaintext_field_and_its_record_parent(): + captured = original( + "note", + NOTE, + { + "id": {"workspace_id": WORKSPACE, "note_id": NOTE}, + "title": "Call", + "parent_record_id": RECORD, + "parent_object": "people", + "content_plaintext": "Discussed fundraising", + "content_markdown": "**Discussed** fundraising", + }, + RecordIdentity(record_type="record", native_id=RECORD, container_id=OBJECT), + ) + async with map_record(captured, "attio", AsyncMock()) as entities: + assert entities[0].text == "Discussed fundraising" + assert entities[0].title == "Call" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("corruption", ["entry_id", "list_id", "parent", "unavailable"]) +async def test_mismatched_or_unavailable_membership_never_projects(corruption): + captured = entry() + if corruption in {"entry_id", "list_id"}: + payload = captured.model_dump()["payload"] + payload["id"][corruption] = str(uuid4()) + captured = captured.model_copy(update={"payload": payload}) + elif corruption == "parent": + captured = captured.model_copy(update={"parent": None}) + else: + captured = captured.model_copy(update={"content_access": "unavailable"}) + with pytest.raises(ProjectionMappingError): + async with map_record(captured, "attio", AsyncMock()): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", ["https://app.attio.com/workspace/person/record", "https://evil.example/record"] +) +async def test_record_attributes_and_native_url(url): + captured = original( + "record", + RECORD, + { + "id": {"workspace_id": WORKSPACE, "object_id": OBJECT, "record_id": RECORD}, + "values": {"name": [{"full_name": "Sam Example"}], "custom_score": [{"value": 12}]}, + "web_url": url, + }, + RecordIdentity(record_type="object", native_id=OBJECT), + ) + if "evil" in url: + with pytest.raises(ProjectionMappingError): + async with map_record(captured, "attio", AsyncMock()): + pass + else: + async with map_record(captured, "attio", AsyncMock()) as entities: + assert entities[0].title == "Sam Example" + assert "custom_score" in entities[0].text + assert entities[0].web_url == url + + +@pytest.mark.asyncio +async def test_two_memberships_for_one_crm_record_remain_distinct(): + first = entry() + second_entry, second_list = str(uuid4()), str(uuid4()) + payload = first.model_dump()["payload"] + payload["id"].update(entry_id=second_entry, list_id=second_list) + second = original( + "entry", + second_entry, + payload, + RecordIdentity(record_type="list", native_id=second_list), + ) + async with map_record(first, "attio", AsyncMock()) as one: + async with map_record(second, "attio", AsyncMock()) as two: + assert one[0].native_id != two[0].native_id + assert RECORD in one[0].text and RECORD in two[0].text + assert one[0].original_kind == two[0].original_kind == "entry" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["object", "list"]) +async def test_container_projection_does_not_invent_children(kind): + native_id = OBJECT if kind == "object" else LIST + payload = {"id": {"workspace_id": WORKSPACE, kind + "_id": native_id}} + if kind == "object": + payload.update(singular_noun="Person", plural_noun="People", api_slug="people") + else: + payload.update(name="Fundraising") + captured = original(kind, native_id, payload) + async with map_record(captured, "attio", AsyncMock()) as entities: + assert len(entities) == 1 + assert "metadata only" in entities[0].content_coverage + assert entities[0].web_url == "" diff --git a/backend/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py index 3be38ffaf5..7d92552965 100644 --- a/backend/tests/unit/platform/http_client/test_composio_transport.py +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -332,3 +332,28 @@ async def test_drive_binding_allows_native_docs_only_with_same_verified_toolkit( assert result.managed_auth.connected_account_id == "ca_drive" assert result.managed_auth.allowed_hosts == {"www.googleapis.com", "docs.googleapis.com"} assert "sheets.googleapis.com" not in result.managed_auth.allowed_hosts + + +@pytest.mark.asyncio +async def test_attio_binding_uses_only_native_api_host_and_verified_account(monkeypatch): + from airweave.domains.auth_provider.exceptions import AuthProviderConfigError + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "key"}, config={"account_id": "ca_attio"} + ) + fetch = AsyncMock( + return_value={ + "id": "ca_attio", + "toolkit": {"slug": "attio"}, + "status": "ACTIVE", + } + ) + monkeypatch.setattr(provider, "_get_with_auth", fetch) + result = await provider.get_auth_result("attio", ["access_token"]) + assert result.credentials is None + assert result.managed_auth.connected_account_id == "ca_attio" + assert result.managed_auth.allowed_hosts == {"api.attio.com"} + fetch.return_value = {"id": "ca_other", "toolkit": {"slug": "attio"}, "status": "ACTIVE"} + with pytest.raises(AuthProviderConfigError, match="identity"): + await provider.get_auth_result("attio", ["access_token"]) From d6f74e23df4ea0e4e8cf363bf49f63d43d7e5b8f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:10:46 -0700 Subject: [PATCH 124/266] Capture Wispr scratchpad notes with durable inventory and body recovery --- .../canonical/tests/test_wispr_recovery.py | 132 ++++++++++++++ backend/airweave/platform/sources/wispr.py | 163 ++++++++++++------ .../platform/sources/test_wispr_capture.py | 126 +++++++++++++- 3 files changed, 364 insertions(+), 57 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py index c02eebabf5..2717365cb6 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py @@ -27,6 +27,8 @@ async def connector(rows, body_calls, *, fail_at=None, failure=None): ) 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"] @@ -144,3 +146,133 @@ async def test_old_parentless_body_requires_explicit_migration(database, source) ) 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, + ) + 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 + assert calls == [ + ("n2", 0), + ("n2", 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 diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 0f4867f324..34e8e08dcd 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -1,4 +1,4 @@ -"""Capture Wispr's available meeting representations through account-bound sessions.""" +"""Capture Wispr meeting and scratchpad representations through account-bound sessions.""" from __future__ import annotations @@ -6,6 +6,7 @@ import json import re from datetime import datetime, timezone +from typing import Literal from urllib.parse import quote from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, StrictBool @@ -53,12 +54,22 @@ class MeetingPage(BaseModel): 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 _MeetingVersion(BaseModel): +class _Version(BaseModel): model_config = ConfigDict(extra="ignore") modified_at: AwareDatetime | None = None @@ -101,8 +112,11 @@ class WisprSource(BaseSource): Captures remain explicitly partial, even after every exposed text range is fetched. """ - canonical_record_types = ("meeting_listing", "meeting") - canonical_container_parents = {"meeting": "meeting_listing"} + 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: @@ -127,10 +141,9 @@ async def create( fingerprint = hashlib.sha256( json.dumps( { - "version": 2, + "version": 3, "account": auth.connected_account_id, - "inventory": "meeting_listing", - "body": "meeting", + "kinds": cls.canonical_record_types, "policy": "discovery_only", "listing_limit": 200, }, @@ -141,7 +154,10 @@ async def create( fingerprint=fingerprint, record_types=cls.canonical_record_types, container_parents=cls.canonical_container_parents, - completion_policies={"meeting_listing": "discovery_only"}, + completion_policies={ + "meeting_listing": "discovery_only", + "scratchpad_listing": "discovery_only", + }, ) return instance @@ -194,15 +210,18 @@ async def _execute(self, slug: str, arguments: dict) -> dict[str, JsonValue]: return result.data async def validate(self) -> None: - """Verify the bound session can actually read meetings.""" + """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 meeting text is not a string") + raise ValueError("Wispr requested text is not a string") matches = list( re.finditer( rf"\(\.\.\.truncated, \d+ chars remaining; continue with " @@ -215,24 +234,38 @@ def _next_offset(text: JsonValue, field: str) -> int | None: raise ValueError("Wispr returned ambiguous continuation markers") return int(matches[0].group(1)) if matches else None - async def _meeting(self, identity: str) -> dict[str, JsonValue]: - offsets = {"content": 0, "transcript": 0} + 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] = [] for _ in range(100): arguments = { - "meeting_id": identity, + id_field: identity, **{ "view_" + field: {"char_limit": 40000, "start_char": offset} for field, offset in offsets.items() }, } - response = await self._execute("WISPR_FLOW_MCP_GET_MEETING", arguments) + response = await self._execute(slug, arguments) if response.get("id") != identity: - raise ValueError("Wispr returned a different meeting 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 meeting changed during paginated capture") + raise ValueError("Wispr body changed during paginated capture") responses.append({"requested_ranges": arguments, "response": response}) following = {} for field, offset in offsets.items(): @@ -245,14 +278,14 @@ async def _meeting(self, identity: str) -> dict[str, JsonValue]: return {"responses": responses} # Request only continuing ranges; complete ranges already remain in responses. offsets = following - raise ValueError("Wispr meeting exceeds the bounded capture page limit") + 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 parent.identity.record_type != "meeting_listing" or record_type != "meeting": + 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="meeting", parent=parent.identity) + 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.""" @@ -271,12 +304,18 @@ async def capture_page( 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. """ - if scope.record_type == "meeting_listing" and scope.parent is None and parent is None: - return await self._listing_page(continuation) + kind = scope.record_type if ( - scope.record_type != "meeting" + 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 parent.identity.record_type != "meeting_listing" + 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 @@ -285,13 +324,13 @@ async def capture_page( 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._meeting(identity) + payload = await self._body(kind, identity) payload["listing"] = parent.payload - modified = _MeetingVersion.model_validate(payload["responses"][0]["response"]).modified_at + modified = _Version.model_validate(payload["responses"][0]["response"]).modified_at return CapturePage( records=( CaptureRecord( - identity=RecordIdentity(record_type="meeting", native_id=identity), + identity=RecordIdentity(record_type=kind, native_id=identity), parent=parent.identity, payload=payload, completeness="partial", @@ -303,7 +342,9 @@ async def capture_page( final=True, ) - async def _listing_page(self, continuation: ScanContinuation) -> CapturePage: + 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: @@ -312,18 +353,52 @@ async def _listing_page(self, continuation: ScanContinuation) -> CapturePage: arguments["since"] = progress.window.since.isoformat() if progress.window.until: arguments["until"] = progress.window.until.isoformat() - page = MeetingPage.model_validate( - await self._execute("WISPR_FLOW_MCP_SEARCH_MEETINGS", arguments) + 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) ) - count = progress.count + len(page.meetings) + 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 - starts = [_MeetingStart.model_validate(row).start for row in page.meetings] - dates_complete = progress.dates_complete and all(value is not None for value in starts) - known_starts = [value for value in starts if value is not 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_starts) if known_starts else None - latest = max(known_starts) if known_starts else 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: @@ -354,23 +429,7 @@ async def _listing_page(self, continuation: ScanContinuation) -> CapturePage: else: next_progress = _Progress() final = True - identities = [_Listing.model_validate(row).id for row in page.meetings] - if len(identities) != len(set(identities)): - raise ValueError("Wispr listing page repeats a meeting identity") - records = tuple( - CaptureRecord( - identity=RecordIdentity(record_type="meeting_listing", native_id=identity), - payload=row, - completeness="metadata_only", - observed_at=datetime.now(timezone.utc), - ) - for identity, row in zip(identities, page.meetings, strict=True) - ) - return CapturePage( - records=records, - continuation=ScanContinuation(value=next_progress.model_dump(mode="json")), - final=final, - ) + return next_progress, final @staticmethod def _partition( diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index c248608d8d..4ba6a619c6 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -27,6 +27,7 @@ async def test_session_is_bound_and_never_auto_connects(monkeypatch): side_effect=[ {"session_id": "session"}, {"data": {"meetings": [], "has_more": False}, "error": None}, + {"data": {"notes": [], "has_more": False}, "error": None}, ] ) monkeypatch.setattr(connector, "_post", post) @@ -56,7 +57,7 @@ async def test_preserves_native_ranges_and_continues_exact_offset(monkeypatch): ] ) monkeypatch.setattr(connector, "_execute", execute) - result = await connector._meeting("m") + result = await connector._body("meeting", "m") assert len(result["responses"]) == 2 assert result["responses"][0]["response"]["native_extra"] is True assert execute.call_args.args[1]["view_transcript"]["start_char"] == 3 @@ -81,7 +82,7 @@ async def test_changed_meeting_fails_instead_of_mixing_versions(monkeypatch): ) monkeypatch.setattr(connector, "_execute", execute) with pytest.raises(ValueError, match="changed"): - await connector._meeting("m") + await connector._body("meeting", "m") @pytest.mark.asyncio @@ -93,7 +94,7 @@ async def test_explicit_tool_failure_never_becomes_partial_success(monkeypatch): AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": "failed"}]), ) with pytest.raises(ValueError, match="tool execution failed; capture is incomplete"): - await connector._meeting("m") + await connector._body("meeting", "m") @pytest.mark.asyncio @@ -286,7 +287,7 @@ async def test_ambiguous_continuation_markers_do_not_choose_a_range(monkeypatch) ) monkeypatch.setattr(connector, "_execute", execute) with pytest.raises(ValueError, match="ambiguous continuation"): - await connector._meeting("m") + await connector._body("meeting", "m") assert execute.await_count == 1 @@ -299,7 +300,7 @@ async def test_rate_signal_stops_without_retry_or_invented_status(monkeypatch, e post = AsyncMock(side_effect=[{"session_id": "session"}, {"data": {}, "error": error}]) monkeypatch.setattr(connector, "_post", post) with pytest.raises(SourceServerError, match="rate-limit signal") as failure: - await connector._meeting("m") + await connector._body("meeting", "m") assert post.await_count == 2 assert failure.value.source_short_name == "wispr" assert failure.value.status_code is None @@ -307,3 +308,118 @@ async def test_rate_signal_stops_without_retry_or_invented_status(monkeypatch, e assert "retry delay are unknown" in str(failure.value) # Do not turn unstructured tool text into credential failure or a timed retry. assert classify_error(failure.value).category is None + + +@pytest.mark.asyncio +async def test_scratchpad_ranges_preserve_normalized_text_and_native_fields(monkeypatch): + connector = await source() + execute = AsyncMock( + side_effect=[ + { + "id": "n", + "title": "Note", + "modified_at": "2026-09-30T00:00:00Z", + "content": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_content.start_char=3...)" + ), + "future_native": {"retained": True}, + }, + {"id": "n", "title": "Note", "modified_at": "2026-09-30T00:00:00Z", "content": "def"}, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + result = await connector._body("scratchpad_note", "n") + assert len(result["responses"]) == 2 + assert result["responses"][0]["response"]["future_native"] == {"retained": True} + assert execute.call_args_list[0].args == ( + "WISPR_FLOW_MCP_GET_SCRATCHPAD_NOTE", + {"note_id": "n", "view_content": {"char_limit": 40000, "start_char": 0}}, + ) + assert execute.call_args.args[1]["view_content"]["start_char"] == 3 + assert all("view_transcript" not in call.args[1] for call in execute.call_args_list) + + +@pytest.mark.asyncio +async def test_scratchpad_partitions_by_modified_time_and_keeps_native_listing(monkeypatch): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + + connector = await source() + rows = [ + {"id": "n1", "modified_at": "2026-01-01T00:00:00Z", "content_excerpt": "original"}, + {"id": "n2", "modified_at": "2026-01-03T00:00:00Z", "content_excerpt": "next"}, + ] + execute = AsyncMock( + side_effect=[ + {"notes": rows, "has_more": True, "truncated": True}, + {"notes": rows[:1], "has_more": False}, + {"notes": rows[1:], "has_more": False}, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + scope = CompletedScope(record_type="scratchpad_listing") + first = await connector.capture_page(scope, ScanContinuation(), files=MagicMock()) + assert first.records[0].payload == rows[0] and not first.final + assert first.records[0].identity.record_type == "scratchpad_listing" + second = await connector.capture_page(scope, first.continuation, files=MagicMock()) + third = await connector.capture_page(scope, second.continuation, files=MagicMock()) + assert third.final + assert execute.call_args_list[1].args[1]["until"] == "2026-01-02T00:00:00+00:00" + assert execute.call_args_list[2].args[1]["since"] == "2026-01-02T00:00:00+00:00" + assert connector.capture_cycle_configuration.policy("scratchpad_listing") == "discovery_only" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second", + [ + {"id": "wrong", "modified_at": "same", "content": "def"}, + {"id": "n", "modified_at": "changed", "content": "def"}, + {"id": "n", "modified_at": "same", "content": None}, + ], +) +async def test_scratchpad_mismatched_or_changed_ranges_fail(monkeypatch, second): + connector = await source() + execute = AsyncMock( + side_effect=[ + { + "id": "n", + "modified_at": "same", + "content": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_content.start_char=3...)" + ), + }, + second, + ] + ) + monkeypatch.setattr(connector, "_execute", execute) + with pytest.raises(ValueError): + await connector._body("scratchpad_note", "n") + + +@pytest.mark.asyncio +async def test_scratchpad_cap_never_substitutes_meeting_start_for_modified_time(monkeypatch): + from airweave.domains.entities.canonical.requests import CompletedScope + from airweave.domains.entities.canonical.scan_models import ScanContinuation + + connector = await source() + monkeypatch.setattr( + connector, + "_execute", + AsyncMock( + return_value={ + "notes": [ + {"id": "n1", "start": "2026-01-01T00:00:00Z"}, + {"id": "n2", "start": "2026-01-03T00:00:00Z"}, + ], + "has_more": True, + "truncated": True, + } + ), + ) + with pytest.raises(ValueError, match="cannot be partitioned safely"): + await connector.capture_page( + CompletedScope(record_type="scratchpad_listing"), ScanContinuation(), files=MagicMock() + ) From eded5614f568b82a5d67e0002421db216ce0526c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:13:16 -0700 Subject: [PATCH 125/266] Project retained Wispr scratchpad notes into search --- .../entities/canonical/projection_mappers.py | 39 +++++++++---- .../airweave/platform/entities/__init__.py | 4 +- backend/airweave/platform/entities/wispr.py | 16 ++++++ .../entities/test_projection_mappers.py | 56 +++++++++++++++++++ 4 files changed, 101 insertions(+), 14 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 3adfd1a947..31fc9ec669 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -22,7 +22,7 @@ ) from airweave.platform.entities.google_drive import GoogleDriveFileEntity, GoogleDriveFolderEntity from airweave.platform.entities.slack import SlackChannelEntity, SlackMessageEntity -from airweave.platform.entities.wispr import WisprMeetingEntity +from airweave.platform.entities.wispr import WisprMeetingEntity, WisprNoteEntity class ProjectionMappingError(ValueError): @@ -163,11 +163,16 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: if not isinstance(window, dict) or window.get("start_char") != expected or complete: raise ProjectionMappingError("Wispr text ranges overlap or have a gap") text = _string(response.get(field)) - marker = re.search( - rf"(?m)^\(\.\.\.truncated, \d+ chars remaining; continue with " - rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", - text, + markers = list( + re.finditer( + rf"(?m)^\(\.\.\.truncated, \d+ chars remaining; continue with " + rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", + text, + ) ) + if len(markers) > 1: + raise ProjectionMappingError("Wispr returned ambiguous continuation markers") + marker = markers[0] if markers else None if marker: following = int(marker.group(1)) if following <= expected: @@ -183,9 +188,9 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: - if record.identity.record_type == "meeting_listing": + if record.identity.record_type in {"meeting_listing", "scratchpad_listing"}: return () - if record.identity.record_type != "meeting": + if record.identity.record_type not in {"meeting", "scratchpad_note"}: raise ProjectionMappingError("Unsupported Wispr record type") responses = record.payload.get("responses") if not isinstance(responses, list) or not responses or not isinstance(responses[0], dict): @@ -193,6 +198,18 @@ def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: first = responses[0].get("response") if not isinstance(first, dict): raise ProjectionMappingError("Wispr meeting response is invalid") + if record.identity.record_type == "scratchpad_note": + if first.get("id") != record.identity.native_id: + raise ProjectionMappingError("Wispr scratchpad identity differs from retained identity") + return ( + WisprNoteEntity( + note_id=record.identity.native_id, + title=_string(first.get("title"), default="Note"), + content=_wispr_text(responses, "content"), + modified_at=_datetime(first.get("modified_at")), + breadcrumbs=[], + ), + ) return ( WisprMeetingEntity( meeting_id=record.identity.native_id, @@ -210,11 +227,9 @@ def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: def excluded_from_search(record: SourceRecord, source_name: str) -> bool: """Retained calendar exclusions intentionally publish no searchable meeting.""" - return ( - source_name == "google_calendar" - and (record.identity.record_type == "event_occurrence" or ( - record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload) - )) + return source_name == "google_calendar" and ( + record.identity.record_type == "event_occurrence" + or (record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload)) ) diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index dfe2f926cc..e16ae894f1 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -334,7 +334,7 @@ TrelloMemberEntity, ) from .web import WebFileEntity -from .wispr import WisprMeetingEntity +from .wispr import WisprMeetingEntity, WisprNoteEntity from .word import WordDocumentEntity from .zendesk import ( ZendeskAttachmentEntity, @@ -704,7 +704,7 @@ "web": [ WebFileEntity, ], - "wispr": [WisprMeetingEntity], + "wispr": [WisprMeetingEntity, WisprNoteEntity], "word": [ WordDocumentEntity, ], diff --git a/backend/airweave/platform/entities/wispr.py b/backend/airweave/platform/entities/wispr.py index cbb03236d6..f2ab126b9c 100644 --- a/backend/airweave/platform/entities/wispr.py +++ b/backend/airweave/platform/entities/wispr.py @@ -30,3 +30,19 @@ class WisprMeetingEntity(BaseEntity): 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/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 405f32df0c..1a17550ddf 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -214,3 +214,59 @@ async def test_owned_drive_spreadsheet_uses_real_converter(): result = await converter.convert_batch([path]) assert "Owned capture" in result[path] and "Verified" in result[path] assert not Path(path).exists() + + +@pytest.mark.asyncio +async def test_wispr_scratchpad_retains_all_text_ranges_without_meeting_fields(): + payload = { + "responses": [ + { + "requested_ranges": {"view_content": {"start_char": 0}}, + "response": { + "id": "note", + "title": "Ideas", + "content": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_content.start_char=3...)" + ), + "modified_at": "2026-09-30T00:00:00Z", + }, + }, + { + "requested_ranges": {"view_content": {"start_char": 3}}, + "response": { + "id": "note", + "title": "Ideas", + "content": "def", + "modified_at": "2026-09-30T00:00:00Z", + }, + }, + ] + } + original = record("scratchpad_note", payload, native_id="note") + before = original.model_dump() + storage = AsyncMock() + async with map_record(original, "wispr", storage) as entities: + assert entities[0].note_id == "note" + assert entities[0].content == "abc\ndef" + assert entities[0].web_url == "" + assert original.model_dump() == before + assert storage.mock_calls == [] + + +@pytest.mark.asyncio +async def test_wispr_projection_rejects_ambiguous_stored_continuation(): + marker = "(...truncated, 3 chars remaining; continue with view_content.start_char=3...)" + payload = { + "responses": [ + { + "requested_ranges": {"view_content": {"start_char": 0}}, + "response": {"id": "note", "content": marker + "\n" + marker}, + } + ] + } + with pytest.raises(ProjectionMappingError, match="ambiguous"): + async with map_record( + record("scratchpad_note", payload, native_id="note"), "wispr", AsyncMock() + ): + pass From 01cd82b07aef13a154034b5ba57ed198d964269b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:15:47 -0700 Subject: [PATCH 126/266] Run all owned source projection contracts in capture verification --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 3a2ec7a9fa..474600ccd9 100644 --- a/Makefile +++ b/Makefile @@ -35,7 +35,7 @@ test-store: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/syncs/tests test-capture: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities/test_github_projection.py tests/unit/domains/entities/test_attio_projection.py tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search From ec3338fa77ab91903d2c18c2aeb548a1de891f6f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:17:29 -0700 Subject: [PATCH 127/266] fix(gmail): retry structured native quota failures without losing evidence --- backend/airweave/platform/sources/gmail.py | 17 +- .../airweave/platform/sources/gmail_errors.py | 75 ++++++++ .../sources/tests/test_gmail_errors.py | 170 ++++++++++++++++++ 3 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/platform/sources/gmail_errors.py create mode 100644 backend/airweave/platform/sources/tests/test_gmail_errors.py diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index 50e861bea3..ab2a48048e 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -62,6 +62,7 @@ ) 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 @@ -70,6 +71,8 @@ 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. @@ -94,6 +97,14 @@ def _should_retry_gmail_request(exception: Exception) -> bool: 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) + + @source( name="Gmail", short_name="gmail", @@ -350,7 +361,7 @@ async def _refresh_and_get_headers(self) -> Dict[str, str]: @retry( stop=stop_after_attempt(5), retry=retry_if_exception(_should_retry_gmail_request), - wait=wait_rate_limit_with_backoff, + wait=_wait_gmail_request, reraise=True, ) async def _get(self, url: str, params: Optional[dict] = None) -> dict: @@ -370,6 +381,7 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: if response.status_code == 429: self.logger.warning("Gmail rate limit reached") + raise_gmail_throttle(response) raise_for_status( response, source_short_name=self.short_name, @@ -383,7 +395,7 @@ async def _get(self, url: str, params: Optional[dict] = None) -> dict: @retry( stop=stop_after_attempt(5), retry=retry_if_exception(_should_retry_gmail_request), - wait=wait_rate_limit_with_backoff, + wait=_wait_gmail_request, reraise=True, ) async def _get_capture_json(self, url: str, params: Optional[dict] = None) -> dict: @@ -425,6 +437,7 @@ async def _get_mime_json( 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, 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/tests/test_gmail_errors.py b/backend/airweave/platform/sources/tests/test_gmail_errors.py new file mode 100644 index 0000000000..403aa7b3c6 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_gmail_errors.py @@ -0,0 +1,170 @@ +"""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", ["_get", "_get_capture_json"]) +@pytest.mark.parametrize("reason", ["rateLimitExceeded", "userRateLimitExceeded"]) +@pytest.mark.parametrize("delay", [None, "3", "180", "invalid", "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("method", ["_get", "_get_capture_json"]) +@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", + ], +) +async def test_permission_unknown_and_malformed_errors_do_not_retry(method, body): + calls = 0 + + def handler(request): + nonlocal calls + calls += 1 + return 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 getattr(source, method)("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 From d5a906bccc1bf42a579a24cbb3ab9da95c14c00e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:17:47 -0700 Subject: [PATCH 128/266] test(gmail): cover malformed quota response bytes --- .../airweave/platform/sources/tests/test_gmail_errors.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/backend/airweave/platform/sources/tests/test_gmail_errors.py b/backend/airweave/platform/sources/tests/test_gmail_errors.py index 403aa7b3c6..bfcfe4d65f 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_errors.py +++ b/backend/airweave/platform/sources/tests/test_gmail_errors.py @@ -75,6 +75,7 @@ def handler(request): } }, "not json", + b"{broken", ], ) async def test_permission_unknown_and_malformed_errors_do_not_retry(method, body): @@ -83,7 +84,11 @@ async def test_permission_unknown_and_malformed_errors_do_not_retry(method, body def handler(request): nonlocal calls calls += 1 - return httpx.Response(403, json=body) + 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( From 20ad9a5265d78d620ad2291b87e302ad8ab2e950 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:18:14 -0700 Subject: [PATCH 129/266] Qualify Wispr resume probes for mixed meeting and scratchpad capture --- .../tests/test_wispr_resume_probe.py | 70 +++++++++++++++++- backend/tests/live/wispr_resume.py | 73 +++++++++++++++---- 2 files changed, 126 insertions(+), 17 deletions(-) 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 index ba04eda148..1d8b9aa6d2 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_resume_probe.py @@ -39,6 +39,19 @@ async def main(): 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): @@ -76,11 +89,11 @@ async def test_wispr_process_restart_skips_completed_bodies(database, source, tm 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"] == 4 + 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"] == 2 + 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 @@ -156,3 +169,56 @@ async def test_wispr_pipeline_rejects_actual_topology_mismatch(database, source, ) 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/tests/live/wispr_resume.py b/backend/tests/live/wispr_resume.py index d87bf11dc3..8a271cdb51 100644 --- a/backend/tests/live/wispr_resume.py +++ b/backend/tests/live/wispr_resume.py @@ -3,14 +3,17 @@ import hashlib import json import os +from typing import Literal from uuid import UUID -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, ValidationError from sqlalchemy import select +from airweave.domains.entities.canonical.cycle_models import CaptureCycle from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity from airweave.models.sync import Sync +from airweave.models.sync_cursor import SyncCursor from airweave.models.sync_job import SyncJob @@ -18,11 +21,14 @@ class CompletedBody(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) scan_id: UUID revision: int - native_digest: str + record_type: Literal["meeting", "scratchpad_note"] + native_digest: str = Field(pattern=r"^[a-f0-9]{64}$") class WisprTarget(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) + schema_version: Literal[3] + configuration_digest: str = Field(pattern=r"^[a-f0-9]{64}$") cycle_id: UUID bodies: tuple[CompletedBody, ...] @@ -40,18 +46,41 @@ def __init__( self.organization_id, self.sync_id = organization_id, sync_id self.job_id, self.attempt_id = job_id, attempt_id self.target_path, self.mode, self.counters = target_path, mode, counters - self.target = ( - WisprTarget.model_validate_json(target_path.read_text()) if mode == "resume" else None - ) + try: + self.target = ( + WisprTarget.model_validate_json(target_path.read_text()) + if mode == "resume" + else None + ) + except ValidationError: + raise ValueError( + "Incompatible Wispr resume target; start a new private v3 trial" + ) from None async def before_page(self, scope, continuation): - if scope.record_type != "meeting": + if scope.record_type not in {"meeting", "scratchpad_note"}: return async with self.sessions() as db: sync, job = await db.get(Sync, self.sync_id), await db.get(SyncJob, self.job_id) assert sync.organization_id == job.organization_id == self.organization_id assert sync.writer_job_id == job.id and sync.writer_attempt_id == self.attempt_id assert job.status == "running" + cursor = await db.scalar( + select(SyncCursor).where( + SyncCursor.organization_id == self.organization_id, + SyncCursor.sync_id == self.sync_id, + ) + ) + if cursor is None: + raise ValueError("Wispr resume requires a persisted capture cycle") + cycle = CaptureCycle.model_validate(cursor.cursor_data.get("canonical_cycle")) + if cycle.configuration.parents != { + "meeting_listing": (None,), + "meeting": ("meeting_listing",), + "scratchpad_listing": (None,), + "scratchpad_note": ("scratchpad_listing",), + }: + raise ValueError("Wispr resume requires the v3 mixed-resource capture topology") rows = ( await db.execute( select(CaptureScan, Entity.native_id) @@ -59,7 +88,8 @@ async def before_page(self, scope, continuation): .where( CaptureScan.organization_id == self.organization_id, CaptureScan.sync_id == self.sync_id, - CaptureScan.record_type == "meeting", + CaptureScan.record_type.in_(("meeting", "scratchpad_note")), + CaptureScan.cycle_id == cycle.version.cycle_id, CaptureScan.phase == "complete", ) ) @@ -68,6 +98,7 @@ async def before_page(self, scope, continuation): CompletedBody( scan_id=row.id, revision=row.revision, + record_type=row.record_type, native_digest=hashlib.sha256(native.encode()).hexdigest(), ) for row, native in rows @@ -75,22 +106,30 @@ async def before_page(self, scope, continuation): cycles = {row.cycle_id for row, _ in rows} if self.target is not None: assert cycles == {self.target.cycle_id} + assert cycle.configuration.digest() == self.target.configuration_digest, ( + "Capture configuration changed" + ) assert all(body in bodies for body in self.target.bodies), "Completed sibling changed" digest = hashlib.sha256(scope.parent.native_id.encode()).hexdigest() - assert digest not in {body.native_digest for body in self.target.bodies}, ( - "Completed body repeated" - ) + assert (scope.record_type, digest) not in { + (body.record_type, body.native_digest) for body in self.target.bodies + }, "Completed body repeated" if len(bodies) == len(self.target.bodies) + 1: - self._stop(True) + self._stop(True, bodies) elif len(bodies) == 3: assert len(cycles) == 1 self.target_path.write_text( - WisprTarget(cycle_id=next(iter(cycles)), bodies=bodies).model_dump_json() + WisprTarget( + schema_version=3, + configuration_digest=cycle.configuration.digest(), + cycle_id=next(iter(cycles)), + bodies=bodies, + ).model_dump_json() ) os.chmod(self.target_path, 0o600) - self._stop(False) + self._stop(False, bodies) - def _stop(self, verified): + def _stop(self, verified, bodies): print( json.dumps( { @@ -99,7 +138,11 @@ def _stop(self, verified): "intentional_interruption": not verified, "wispr_recovery_verified": verified, "full_scope_completed": False, - "completed_bodies": 4 if verified else 3, + "completed_bodies": len(bodies), + "completed_body_kinds": { + kind: sum(body.record_type == kind for body in bodies) + for kind in ("meeting", "scratchpad_note") + }, **self.counters, } ), From d4f988540fccca57c8d31cc95a245136e6e5ac7a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:23:31 -0700 Subject: [PATCH 130/266] Keep Gmail error tests focused on distinct failure boundaries --- .../sources/tests/test_gmail_errors.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/backend/airweave/platform/sources/tests/test_gmail_errors.py b/backend/airweave/platform/sources/tests/test_gmail_errors.py index bfcfe4d65f..e65709c193 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_errors.py +++ b/backend/airweave/platform/sources/tests/test_gmail_errors.py @@ -18,9 +18,16 @@ def envelope(reason, *, domain="usageLimits", code=403): @pytest.mark.asyncio -@pytest.mark.parametrize("method", ["_get", "_get_capture_json"]) -@pytest.mark.parametrize("reason", ["rateLimitExceeded", "userRateLimitExceeded"]) -@pytest.mark.parametrize("delay", [None, "3", "180", "invalid", "NaN"]) +@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 @@ -54,7 +61,6 @@ def handler(request): @pytest.mark.asyncio -@pytest.mark.parametrize("method", ["_get", "_get_capture_json"]) @pytest.mark.parametrize( "body", [ @@ -78,7 +84,7 @@ def handler(request): b"{broken", ], ) -async def test_permission_unknown_and_malformed_errors_do_not_retry(method, body): +async def test_permission_unknown_and_malformed_errors_do_not_retry(body): calls = 0 def handler(request): @@ -98,7 +104,7 @@ def handler(request): config=GmailConfig(), ) with pytest.raises(SourceEntityForbiddenError): - await getattr(source, method)("https://gmail.googleapis.com/gmail/v1/users/me/profile") + await source._get_capture_json("https://gmail.googleapis.com/gmail/v1/users/me/profile") assert calls == 1 From ddb5ea5900907eca22ba16139c7282d314f69164 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:27:44 -0700 Subject: [PATCH 131/266] Index retained Gmail body when attachment capture is partial --- .../entities/canonical/gmail_projection.py | 15 ++++++-- .../canonical/tests/test_gmail_projection.py | 36 +++++++++++++++++-- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py index 580b5db5e9..b3367eae15 100644 --- a/backend/airweave/domains/entities/canonical/gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -77,7 +77,13 @@ def __init__(self, record: SourceRecord, storage: StorageBackend, directory: Pat self.attachments: list[GmailAttachmentEntity] = [] async def attachment(self, part: dict, path: str) -> None: - """Retain each MIME attachment identity, including duplicate filenames.""" + """Project retained attachment bytes; keep missing bytes explicit on the original.""" + if ( + self.record.completeness == "partial" + and part.get("body", {}).get("attachmentId") + and not any(blob.source_path == path + "/body" for blob in self.record.blobs) + ): + return content = await _body(part, path, self.record, self.storage) filename = part.get("filename") or "attachment" mime = part.get("mimeType", "application/octet-stream") @@ -131,8 +137,11 @@ async def render(self, part: dict, path: str) -> str: async def map(self) -> tuple[BaseEntity, ...]: """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 != "complete": - raise ValueError("Gmail projection requires a complete active source record") + 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") diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py index 5110f672ab..46bdec7632 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -156,9 +156,39 @@ async def test_corrupt_blob_bytes_fail_before_projection(tmp_path): @pytest.mark.asyncio -async def test_partial_capture_cannot_be_marked_successfully_projected(tmp_path): - with pytest.raises(ValueError, match="complete active"): - await map_gmail(record(part(), completeness="partial"), AsyncMock(), tmp_path) +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) + 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 From 4a63b6b2e57771d00077fffb059de9a41ce43a6b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:33:32 -0700 Subject: [PATCH 132/266] Require pinned native mailbox identity before owned Gmail capture --- .../canonical/tests/test_gmail_recovery.py | 49 ++++++++++--- .../tests/test_gmail_resume_probe.py | 7 +- backend/airweave/platform/configs/config.py | 9 +++ backend/airweave/platform/sources/gmail.py | 39 ++++++++++- .../sources/tests/test_gmail_identity.py | 70 +++++++++++++++++++ backend/tests/live/provider_sample.py | 7 +- 6 files changed, 162 insertions(+), 19 deletions(-) create mode 100644 backend/airweave/platform/sources/tests/test_gmail_identity.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py index c78695840b..899fe92df4 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_recovery.py @@ -1,5 +1,6 @@ """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 @@ -27,6 +28,8 @@ def __init__(self, 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 @@ -44,7 +47,11 @@ async def setup(database, source, native, *, attempt=1, query=None): logger=MagicMock(), http_client=client, config=GmailConfig( - gmail_query=query, included_labels=[], excluded_labels=[], excluded_categories=[] + expected_mailbox="synthetic@example.test", + gmail_query=query, + included_labels=[], + excluded_labels=[], + excluded_categories=[], ), ) ctx, _, runtime, bus = components(database, source) @@ -98,7 +105,7 @@ async def test_full_resume_interleaved_history_known_omission_and_next_delta(dat ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): with pytest.raises(ConnectionError, match="interruption"): await run(pipeline, ctx, runtime) async with database() as db: @@ -106,7 +113,7 @@ async def test_full_resume_interleaved_history_known_omission_and_next_delta(dat 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 client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert len([url for url, _ in native.calls if url.endswith("/profile")]) == 1 assert [ @@ -145,7 +152,7 @@ async def test_full_resume_interleaved_history_known_omission_and_next_delta(dat ] ) pipeline, ctx, runtime, client = await setup(database, (service, next_fence), delta) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert delta.calls[0][1]["startHistoryId"] == "120" assert not delta.responses @@ -169,7 +176,7 @@ async def test_filtered_mutating_view_omission_is_not_provider_deletion(database [("/messages", {"messages": [{"id": "a"}]}), ("/messages/a", message("a", labelIds=[]))] ) pipeline, ctx, runtime, client = await setup(database, source, native, query="in:inbox") - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert native.calls[0][1]["q"] == "in:inbox" async with database() as db: @@ -193,7 +200,7 @@ async def test_history_expiry_restarts_full_without_using_partial_sightings(data ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) async with database() as db: cycle = await source[0].read_cycle(db, pipeline._writer()) @@ -221,17 +228,17 @@ async def test_history_offset_resume_and_lost_final_ack_do_not_rehydrate_committ ) native = NativeHTTP(responses) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + 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 client: + 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 client: + 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 @@ -268,7 +275,7 @@ async def test_new_job_resumes_active_cycle_after_exhausted_server_error( ) service, fence = source pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): with pytest.raises(SourceServerError): await run(pipeline, ctx, runtime) async with database() as db: @@ -300,7 +307,7 @@ async def test_new_job_resumes_active_cycle_after_exhausted_server_error( attempt_number=1, ) pipeline, ctx, runtime, client = await setup(database, (service, next_fence), native) - async with client: + 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 @@ -311,3 +318,23 @@ async def test_new_job_resumes_active_cycle_after_exhausted_server_error( 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 index 48058cd6e2..57c88dd83f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_resume_probe.py @@ -47,12 +47,15 @@ async def main(): '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(), @@ -183,7 +186,6 @@ async def test_actual_lifecycle_child_context_reaches_committed_gmail_page( import conftest import httpx import provider_lifecycle as lifecycle -from provider_sample import verify_rest_identity 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 @@ -194,7 +196,6 @@ async def test_actual_lifecycle_child_context_reaches_committed_gmail_page( async def synthetic_source(name, account, expected_email, key, fence, **options): assert name == 'gmail' and options['gmail_unfiltered'] is True native = NativeHTTP([ - ('/profile', {'emailAddress': expected_email, 'historyId': 'identity-only'}), ('/profile', {'historyId': 'capture-boundary'}), ('/messages', {'messages': [{'id': 'synthetic-a'}], 'nextPageToken': 'next'}), ('/messages/synthetic-a', message('synthetic-a')), @@ -203,8 +204,8 @@ async def synthetic_source(name, account, expected_email, key, fence, **options) 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=[])) - await verify_rest_identity(name, connector, expected_email) yield connector, 'provider_email' lifecycle.rest_source = synthetic_source diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index c7f0edb315..76ceaa7aab 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -240,6 +240,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", diff --git a/backend/airweave/platform/sources/gmail.py b/backend/airweave/platform/sources/gmail.py index ab2a48048e..8acb1765ff 100644 --- a/backend/airweave/platform/sources/gmail.py +++ b/backend/airweave/platform/sources/gmail.py @@ -14,6 +14,7 @@ from typing import Any, AsyncGenerator, Dict, List, Optional, Set 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 @@ -105,6 +106,13 @@ def _wait_gmail_request(retry_state) -> float: 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", @@ -134,15 +142,23 @@ class GmailSource(BaseSource): 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": 1, + "capture_version": 2, "query": query, "filters": self.capture_config.model_dump(mode="json"), }, @@ -156,6 +172,7 @@ def capture_cycle_configuration(self) -> CycleConfiguration: ) def _canonical_capture(self, files: FileService | None = None) -> GmailCapture: + self._require_mailbox() return GmailCapture( self._get_capture_json, self._build_gmail_query(), @@ -171,6 +188,7 @@ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: 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( @@ -233,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 @@ -1074,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/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/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index 3c62c680fa..5f42d0a792 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -87,6 +87,7 @@ async def rest_source( "gmail": ( GmailSource, GmailConfig( + expected_mailbox=expected_email, gmail_query=None, included_labels=[], excluded_labels=[], @@ -94,7 +95,7 @@ async def rest_source( after_date=None, ) if gmail_unfiltered - else GmailConfig(gmail_query=gmail_query), + else GmailConfig(gmail_query=gmail_query, expected_mailbox=expected_email), ), "google_drive": (GoogleDriveSource, GoogleDriveConfig()), "google_calendar": (GoogleCalendarSource, calendar_config or GoogleCalendarConfig()), @@ -103,7 +104,9 @@ async def rest_source( source = await source_type.create( auth=auth, logger=logging.getLogger("probe"), http_client=wrapped, config=config ) - await verify_rest_identity(name, source, expected_email) + # Gmail create attests this selected mailbox using its native profile. + if name != "gmail": + await verify_rest_identity(name, source, expected_email) yield source, "provider_email" From 466fa926d5244c031d66ecaec904a05e7e6ca260 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:41:50 -0700 Subject: [PATCH 133/266] Allow retained XLSX files through the search text pipeline --- .../domains/sync_pipeline/file_types.py | 2 +- .../entities/test_projection_mappers.py | 18 +++++++++++++----- 2 files changed, 14 insertions(+), 6 deletions(-) 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/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 1a17550ddf..e8bae1ccf9 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -1,7 +1,8 @@ """Search projections preserve canonical data and cannot fetch missing source content.""" from datetime import datetime, timezone -from unittest.mock import AsyncMock +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 import pytest @@ -184,6 +185,7 @@ async def test_owned_drive_spreadsheet_uses_real_converter(): from airweave.domains.converters.registry import ConverterRegistry from airweave.domains.entities.canonical.requests import BlobReference + from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder workbook = openpyxl.Workbook() workbook.active.append(["Project", "Status"]) @@ -209,10 +211,16 @@ async def test_owned_drive_spreadsheet_uses_real_converter(): async with map_record(item, "google_drive", storage) as entities: path = entities[0].local_path assert Path(path).suffix == ".xlsx" - converter = ConverterRegistry().for_extension(".xlsx") - assert converter is not None - result = await converter.convert_batch([path]) - assert "Owned capture" in result[path] and "Verified" in result[path] + tracker = AsyncMock() + result = await TextualRepresentationBuilder(ConverterRegistry()).build_for_batch( + list(entities), + SimpleNamespace(source_short_name="google_drive", logger=MagicMock()), + SimpleNamespace(entity_tracker=tracker), + ) + assert len(result) == 1 + assert "Owned capture" in result[0].textual_representation + assert "Verified" in result[0].textual_representation + tracker.record_skipped.assert_not_awaited() assert not Path(path).exists() From 28f8fe96b80b50d8470c9577b3d9a62447e8992d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:47:33 -0700 Subject: [PATCH 134/266] Gate owned Calendar capture on pinned primary calendar identity --- .../tests/test_calendar_lifecycle_probe.py | 8 +-- .../canonical/tests/test_calendar_recovery.py | 27 +++++++- .../platform/sources/google_calendar.py | 63 +++++++++++++---- .../sources/records/calendar_pages.py | 2 +- .../sources/tests/test_calendar_identity.py | 67 +++++++++++++++++++ 5 files changed, 147 insertions(+), 20 deletions(-) create mode 100644 backend/airweave/platform/sources/tests/test_calendar_identity.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py index a3773319c9..b6ed148461 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py @@ -15,7 +15,6 @@ import conftest import httpx import provider_lifecycle as lifecycle -from provider_sample import verify_rest_identity from airweave.domains.entities.canonical.tests.test_calendar_recovery import NativeHTTP, event from airweave.domains.sources.token_providers.static import StaticTokenProvider from airweave.platform.sources.google_calendar import GoogleCalendarSource @@ -26,17 +25,16 @@ async def synthetic_source(name, account, expected_email, key, fence, **options) path = '/calendars/' + quote(expected_email, safe='') + '/events' first = os.environ['TEST_STAGE'] == 'first' native = NativeHTTP([ - ('/users/me/calendarList/primary', {'id': expected_email, 'primary': True}), ('/users/me/calendarList', {'items': [{'id': expected_email, 'timeZone': 'UTC'}]}), (path, {'items': [event('a')] if first else [], 'nextSyncToken': 'full' if first else 'delta'}), (path, {'items': [event('a')]}), - ]) + ], principal=expected_email) async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle), event_hooks={'request': [options['request_hook']]}) as client: connector = await GoogleCalendarSource.create(auth=StaticTokenProvider('synthetic'), - logger=MagicMock(), http_client=client, config=options['calendar_config']) - await verify_rest_identity(name, connector, expected_email) + logger=MagicMock(), http_client=client, config=options['calendar_config'].model_copy( + update={'expected_primary_calendar_id': expected_email})) yield connector, 'provider_email' assert not native.replies diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py index 0cb932a9ca..c811dd0476 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_recovery.py @@ -34,10 +34,13 @@ def event(native_id): class NativeHTTP: - def __init__(self, replies): + def __init__(self, replies, principal="fixture-primary"): + self.principal = principal self.replies, self.calls = list(replies), [] async def handle(self, request): + if request.url.path.endswith("/calendars/primary"): + return httpx.Response(200, json={"id": self.principal}) self.calls.append((request.url.path, dict(request.url.params))) path, value = self.replies.pop(0) assert str(request.url).split("?")[0] == BASE + path @@ -49,6 +52,7 @@ async def handle(self, request): async def setup(database, source, native, attempt=1, config=CONFIG): service, fence = source + config = config.model_copy(update={"expected_primary_calendar_id": "fixture-primary"}) client = httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) connector = await GoogleCalendarSource.create( auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=client, config=config @@ -414,3 +418,24 @@ async def stop_after_role_commit(): parent = await db.scalar(select(Entity).where(Entity.native_id == "cal")) assert parent.visibility_epoch == old_epoch + 1 await client.aclose() + + +async def test_wrong_primary_identity_never_writes_capture(database, source): + from airweave.models.sync_cursor import SyncCursor + + native = NativeHTTP([], principal="wrong-primary") + async with database() as db: + before = [(row.id, row.cursor_data) for row in (await db.scalars(select(SyncCursor))).all()] + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) as client: + with pytest.raises(ValueError, match="does not match"): + await GoogleCalendarSource.create( + auth=StaticTokenProvider("same-broker-account"), + logger=MagicMock(), + http_client=client, + config=GoogleCalendarConfig(expected_primary_calendar_id="trusted-primary"), + ) + assert native.calls == [] + 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/platform/sources/google_calendar.py b/backend/airweave/platform/sources/google_calendar.py index 52737608e4..eb8cc4f86f 100644 --- a/backend/airweave/platform/sources/google_calendar.py +++ b/backend/airweave/platform/sources/google_calendar.py @@ -17,6 +17,7 @@ from typing import Any, AsyncGenerator, Dict, List, Optional 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 @@ -80,6 +81,13 @@ def _retry_capture(exception: BaseException) -> bool: 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", @@ -106,17 +114,32 @@ 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 CalendarPages(self._get_capture_json, self.calendar_config).configuration + return self._calendar_pages().configuration async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: """Resolve the expansion window once per persisted cycle.""" - return CalendarPages(self._get_capture_json, self.calendar_config).prepare() + return self._calendar_pages().prepare() async def prepare_scope( self, @@ -128,7 +151,7 @@ async def prepare_scope( force_full: bool, ) -> ScopePlan: """Choose native delta only from compatible completed exact-scope evidence.""" - return CalendarPages(self._get_capture_json, self.calendar_config).scope_plan( + return self._calendar_pages().scope_plan( scope, cycle, previous, parent=parent, force_full=force_full ) @@ -136,6 +159,7 @@ 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: @@ -153,17 +177,13 @@ async def capture_page( parent: SourceRecord | None = None, ) -> CapturePage: """Originals and pagination commit together through the production page driver.""" - return await CalendarPages(self._get_capture_json, self.calendar_config).page( - scope, continuation - ) + 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 CalendarPages(self._get_capture_json, self.calendar_config).confirm_member_absent( - record.identity.native_id - ) + await self._calendar_pages().confirm_member_absent(record.identity.native_id) # ----------------------- # Construction / Config @@ -187,6 +207,8 @@ async def create( 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 @@ -493,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/records/calendar_pages.py b/backend/airweave/platform/sources/records/calendar_pages.py index 7ad52ead89..18ddefb587 100644 --- a/backend/airweave/platform/sources/records/calendar_pages.py +++ b/backend/airweave/platform/sources/records/calendar_pages.py @@ -63,7 +63,7 @@ def __init__(self, get: GetJSON, config: GoogleCalendarConfig): @property def configuration(self) -> CycleConfiguration: """Fingerprint native selection and capture semantics, not moving wall-clock bounds.""" - raw = {"version": 2, "config": self.config.model_dump(mode="json")} + raw = {"version": 3, "config": self.config.model_dump(mode="json")} return CycleConfiguration( fingerprint=hashlib.sha256(json.dumps(raw, sort_keys=True).encode()).hexdigest(), parents={ 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(), + ) From ab8d0431c6db685503c386bdc8af165f591c810d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:48:51 -0700 Subject: [PATCH 135/266] Label Calendar probe evidence by native primary identity --- .../entities/canonical/tests/test_calendar_lifecycle_probe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py index b6ed148461..b2deb4c628 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_lifecycle_probe.py @@ -35,7 +35,7 @@ async def synthetic_source(name, account, expected_email, key, fence, **options) connector = await GoogleCalendarSource.create(auth=StaticTokenProvider('synthetic'), logger=MagicMock(), http_client=client, config=options['calendar_config'].model_copy( update={'expected_primary_calendar_id': expected_email})) - yield connector, 'provider_email' + yield connector, 'provider_primary_calendar_id' assert not native.replies lifecycle.rest_source = synthetic_source From 2b00e3d40500ffc31ecb7f04862ec23f30327f87 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:49:11 -0700 Subject: [PATCH 136/266] Require native Drive principal attestation for owned capture --- .../tests/test_drive_lifecycle_probe.py | 6 +- .../canonical/tests/test_drive_recovery.py | 20 ++-- .../tests/test_workspace_document_probe.py | 5 +- .../airweave/platform/sources/google_drive.py | 66 ++++++++++++- .../sources/tests/test_drive_identity.py | 95 +++++++++++++++++++ 5 files changed, 173 insertions(+), 19 deletions(-) create mode 100644 backend/airweave/platform/sources/tests/test_drive_identity.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py index 72ff306dd2..75d14429e6 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py @@ -23,7 +23,7 @@ @asynccontextmanager async def synthetic_source(name, account, expected_email, key, fence, **options): assert name == 'google_drive' - responses = [('/about', {'user': {'emailAddress': expected_email}})] + responses = [] if os.environ['TEST_STAGE'] == 'first': responses += [('/changes/startPageToken', {'startPageToken': 'before'}), ('/files', {'files': [{'id': 'folder-a'}]}), @@ -35,8 +35,8 @@ async def synthetic_source(name, account, expected_email, key, fence, **options) async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle), event_hooks={'request': [options['request_hook']]}) as client: connector = await GoogleDriveSource.create(auth=StaticTokenProvider('synthetic'), - logger=MagicMock(), http_client=client, config=GoogleDriveConfig()) - await verify_rest_identity(name, connector, expected_email) + logger=MagicMock(), http_client=client, + config=GoogleDriveConfig(expected_permission_id="principal-a")) yield connector, 'provider_email' assert not native.replies diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py index 29dedd19d6..0fce356bd9 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_recovery.py @@ -1,5 +1,6 @@ """Real canonical SQL/pipeline and Drive source; native HTTP is synthetic.""" +from contextlib import aclosing from unittest.mock import MagicMock from uuid import uuid4 @@ -24,6 +25,9 @@ def __init__(self, replies): self.calls = [] async def handle(self, request): + if request.url.path.endswith("/about"): + assert request.url.params["fields"] == "user(permissionId,me)" + return httpx.Response(200, json={"user": {"permissionId": "principal-a", "me": True}}) self.calls.append((request.url.path, dict(request.url.params))) path, value = self.replies.pop(0) assert str(request.url).split("?")[0] == BASE + path @@ -44,7 +48,7 @@ async def setup(database, source, native, attempt=1): auth=StaticTokenProvider("synthetic"), logger=MagicMock(), http_client=client, - config=GoogleDriveConfig(), + config=GoogleDriveConfig(expected_permission_id="principal-a"), ) ctx, _, runtime, bus = components(database, source) pipeline = CanonicalCapturePipeline( @@ -91,7 +95,7 @@ async def test_resume_pending_file_and_changes_then_lost_final_ack(database, sou ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): with pytest.raises(ConnectionError): await run(pipeline, ctx, runtime) async with database() as db: @@ -99,11 +103,11 @@ async def test_resume_pending_file_and_changes_then_lost_final_ack(database, sou assert cycle.promoted_checkpoint is None assert len((await db.scalars(select(Entity))).all()) == 1 pipeline, ctx, runtime, client = await setup(database, source, native, 2) - async with client: + async with aclosing(client): await pipeline.run_scans(ctx, runtime, no_limits) # Crash after final durable page, before publication. No provider call on retry. pipeline, ctx, runtime, client = await setup(database, source, native, 3) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert not native.replies assert sum(path.endswith("/files") for path, _ in native.calls) == 1 @@ -146,7 +150,7 @@ async def test_resume_pending_file_and_changes_then_lost_final_ack(database, sou ] ) pipeline, ctx, runtime, client = await setup(database, (service, next_fence), delta) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert delta.calls[0][1]["pageToken"] == "after" assert not delta.replies @@ -174,7 +178,7 @@ async def test_drive_membership_restart_preserves_records_until_exact_validation ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert not native.replies async with database() as db: @@ -214,7 +218,7 @@ async def test_rejected_inventory_token_restarts_sweep_not_boundary(database, so ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) assert not native.replies assert sum(path.endswith("/startPageToken") for path, _ in native.calls) == 1 @@ -247,7 +251,7 @@ async def test_provider_failure_cannot_promote_or_remove_prior_capture(database, ] ) pipeline, ctx, runtime, client = await setup(database, source, native) - async with client: + async with aclosing(client): with pytest.raises(ValueError if failure == "incomplete" else SourceEntityForbiddenError): await run(pipeline, ctx, runtime) async with database() as db: 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 index edf7b6940a..ea9442c099 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_document_probe.py @@ -34,7 +34,7 @@ def response(request): assert request.method=='GET' seen.append(request.url.path) if request.url.path.endswith('/about'): - return httpx.Response(200,json={'user':{'emailAddress':'synthetic@example.test'}}) + 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'): @@ -49,8 +49,7 @@ 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()) - await verify_rest_identity(name,connector,email) + logger=MagicMock(),http_client=client,config=GoogleDriveConfig(expected_permission_id="principal-a")) yield connector,'provider_email' async def main(): diff --git a/backend/airweave/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py index 25a892f8bf..7a6b9d5792 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -17,10 +17,12 @@ 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 @@ -65,6 +67,21 @@ 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", @@ -95,14 +112,34 @@ class GoogleDriveSource(BaseSource): canonical_record_types = ("file",) canonical_container_parents = {} + expected_permission_id: str | None = None + _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(b"drive-all-accessible-docs-native-v2").hexdigest(), + fingerprint=hashlib.sha256( + json.dumps( + { + "capture_version": 3, + "permission_id": self.expected_permission_id, + "scope": "drive-all-accessible-docs-native", + }, + sort_keys=True, + ).encode() + ).hexdigest(), parents={"file": (None,)}, known_object_validation=("file",), ) @@ -113,6 +150,9 @@ async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: 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: @@ -134,6 +174,7 @@ async def capture_page( 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") @@ -144,6 +185,7 @@ async def hydrate(record: CaptureRecord) -> CaptureRecord: 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( @@ -170,11 +212,14 @@ 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.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 @@ -183,11 +228,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), 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)) From 918e4593d4a23001189ba43a7637202b99a8320f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:49:21 -0700 Subject: [PATCH 137/266] Require explicit trusted Google identities in owned capture configuration --- backend/airweave/platform/configs/config.py | 14 ++++++++++ backend/tests/live/provider_sample.py | 31 +++++++++++++++++---- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 76ceaa7aab..d5aa8088fa 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -330,6 +330,13 @@ def bounded(self): class GoogleCalendarConfig(SourceConfig): """Rolling expanded capture, or an explicit window through source configuration.""" + 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") @@ -375,6 +382,13 @@ 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.", + ) + include_patterns: list[str] = Field( default=[], title="Include Patterns", diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index 5f42d0a792..e89afd24cb 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -83,6 +83,15 @@ async def rest_source( wrapped = AirweaveHttpClient( client, fence.organization_id, name, feature_flag_enabled=False ) + if name == "google_calendar": + calendar_config = calendar_config or GoogleCalendarConfig() + if calendar_config.expected_primary_calendar_id is None: + calendar_config = GoogleCalendarConfig.model_validate( + { + **calendar_config.model_dump(), + "expected_primary_calendar_id": os.environ["LIVE_CALENDAR_PRIMARY_ID"], + } + ) source_type, config = { "gmail": ( GmailSource, @@ -97,17 +106,29 @@ async def rest_source( if gmail_unfiltered else GmailConfig(gmail_query=gmail_query, expected_mailbox=expected_email), ), - "google_drive": (GoogleDriveSource, GoogleDriveConfig()), - "google_calendar": (GoogleCalendarSource, calendar_config or GoogleCalendarConfig()), + "google_drive": ( + GoogleDriveSource, + GoogleDriveConfig(expected_permission_id=os.environ["LIVE_DRIVE_PERMISSION_ID"]) + if name == "google_drive" + else GoogleDriveConfig(), + ), + "google_calendar": ( + GoogleCalendarSource, + calendar_config or GoogleCalendarConfig(), + ), "slack": (SlackSource, SlackConfig()), }[name] source = await source_type.create( auth=auth, logger=logging.getLogger("probe"), http_client=wrapped, config=config ) - # Gmail create attests this selected mailbox using its native profile. - if name != "gmail": + # Google sources attest their explicitly pinned native identity during create. + if name not in {"gmail", "google_calendar", "google_drive"}: await verify_rest_identity(name, source, expected_email) - yield source, "provider_email" + identity_kind = { + "google_drive": "provider_permission_id", + "google_calendar": "provider_primary_calendar_id", + }.get(name, "provider_email") + yield source, identity_kind @asynccontextmanager From b54561d063468dbcb513a841b457ecea94f90dee Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 13:59:17 -0700 Subject: [PATCH 138/266] Reject limited Slack history before exhaustive reconciliation --- .../tests/test_slack_limited_recovery.py | 39 +++++++++++++++++++ backend/airweave/platform/sources/slack.py | 27 ++++++++++++- .../sources/tests/test_slack_history_limit.py | 25 ++++++++++++ 3 files changed, 89 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py create mode 100644 backend/airweave/platform/sources/tests/test_slack_history_limit.py 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..9e3acd4086 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py @@ -0,0 +1,39 @@ +"""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, _, _ = runner(database, source, [ROOT, {"messages": [{"ts": "1", "text": "old"}]}]) + await run(baseline) + _, before, _ = await saved(database) + source = await next_job(database, source) + refresh, _, _ = 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/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index c3416cfe82..f6318c4b06 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -8,7 +8,15 @@ from functools import partial from typing import Any, AsyncGenerator, Dict, List, Optional -from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from pydantic import ( + BaseModel, + ConfigDict, + Field, + JsonValue, + StrictBool, + TypeAdapter, + ValidationError, +) from tenacity import retry, retry_if_exception_type, stop_after_attempt from airweave.core.logging import ContextualLogger @@ -28,7 +36,7 @@ RecordIdentity, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation -from airweave.domains.sources.exceptions import SourceAuthError, SourceRateLimitError +from airweave.domains.sources.exceptions import SourceAuthError, SourceError, SourceRateLimitError from airweave.domains.sources.token_providers.protocol import ( SourceAuthProvider, authorization_headers, @@ -124,6 +132,13 @@ def __init__(self, code: str): super().__init__(f"Slack request failed: {code}") +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.""" @@ -373,6 +388,14 @@ async def _capture_page_response( # 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() 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..daaed326f0 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_slack_history_limit.py @@ -0,0 +1,25 @@ +"""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.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._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(), + ) From b661c73ef26d233f8159b936d9befc21dfd00fc1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:01:35 -0700 Subject: [PATCH 139/266] Update native Docs recovery fixture for Drive identity attestation --- .../entities/canonical/tests/test_workspace_capture.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py b/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py index 1837936e7b..906c8c85f8 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py +++ b/backend/airweave/domains/entities/canonical/tests/test_workspace_capture.py @@ -1,6 +1,7 @@ """Actual Drive pipeline publishes Docs references atomically; HTTP is synthetic.""" import hashlib +from contextlib import aclosing from unittest.mock import AsyncMock import httpx @@ -25,6 +26,11 @@ def __init__(self): 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") @@ -84,7 +90,7 @@ async def test_recapture_original_uuid_atomically_adds_manifest_and_withdraws_al native = NativeDocs() pipeline, ctx, runtime, client = await setup(database, source, native) pipeline.files = files - async with client: + async with aclosing(client): with pytest.raises(ConnectionError, match="interrupted"): await run(pipeline, ctx, runtime) async with database() as db: @@ -96,7 +102,7 @@ async def test_recapture_original_uuid_atomically_adds_manifest_and_withdraws_al native.fail = False pipeline, ctx, runtime, client = await setup(database, source, native, attempt=2) pipeline.files = files - async with client: + async with aclosing(client): await run(pipeline, ctx, runtime) async with database() as db: completed = await service.read_cycle(db, pipeline._writer()) From 77a096a29bea3be95823a994d5bc76ef415f1c50 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:08:38 -0700 Subject: [PATCH 140/266] Attest Slack workspace and user before owned capture --- .../canonical/tests/test_resume_probe.py | 2 +- .../tests/test_slack_identity_writes.py | 37 +++++++++ .../tests/test_slack_limited_recovery.py | 6 +- .../canonical/tests/test_slack_membership.py | 4 + .../sync_pipeline/tests/test_factory.py | 6 ++ backend/airweave/platform/sources/slack.py | 77 +++++++++++++++++-- .../sources/tests/test_slack_history_limit.py | 4 + .../sources/tests/test_slack_identity.py | 66 ++++++++++++++++ .../platform/sources/test_slack_capture.py | 29 ++++--- 9 files changed, 209 insertions(+), 22 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_identity_writes.py create mode 100644 backend/airweave/platform/sources/tests/test_slack_identity.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py index 257317f616..a52865796e 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_resume_probe.py @@ -61,7 +61,7 @@ async def main(): responses = [ROOT, SlackApiError("not_in_channel")] if mode == "expired": responses = [ROOT, SlackApiError("invalid_cursor"), HISTORY, REPLIES] - instance, connector, pipeline = runner(sessions, (service, fence), responses, + 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 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 index 9e3acd4086..b9e4d93826 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_limited_recovery.py @@ -10,11 +10,13 @@ async def test_limited_final_history_preserves_prior_originals_and_committed_progress( database, source ): - baseline, _, _ = runner(database, source, [ROOT, {"messages": [{"ts": "1", "text": "old"}]}]) + 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, _, _ = runner( + refresh, _, _ = await runner( database, source, [ diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py index 98e34c7af1..09ed529aba 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py @@ -10,6 +10,7 @@ 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 @@ -39,6 +40,9 @@ async def test_first_completed_membership_hides_prior_partial_channel( connector = SlackSource( auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock() ) + connector.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U1") + 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") diff --git a/backend/airweave/domains/sync_pipeline/tests/test_factory.py b/backend/airweave/domains/sync_pipeline/tests/test_factory.py index 2b0bd6d395..4e396d275a 100644 --- a/backend/airweave/domains/sync_pipeline/tests/test_factory.py +++ b/backend/airweave/domains/sync_pipeline/tests/test_factory.py @@ -182,11 +182,17 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): mock_source.generate_entities = MagicMock(return_value=AsyncMock()) if canonical == "page": 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() elif canonical: mock_source.canonical_record_types = ("event",) mock_source.generate_observations = MagicMock(return_value=AsyncMock()) diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index f6318c4b06..b03ae0c9c0 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import json import re from datetime import datetime, timedelta, timezone from functools import partial @@ -16,6 +17,7 @@ StrictBool, TypeAdapter, ValidationError, + model_validator, ) from tenacity import retry, retry_if_exception_type, stop_after_attempt @@ -132,6 +134,22 @@ def __init__(self, code: str): super().__init__(f"Slack request failed: {code}") +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.""" @@ -186,11 +204,36 @@ class SlackSource(BaseSource): canonical_record_types = ("channel", "message") canonical_container_parents = {"message": "channel"} - capture_cycle_configuration = CycleConfiguration.from_source( - fingerprint=hashlib.sha256(b"slack:whole-conversations:history-and-replies:v1").hexdigest(), - record_types=canonical_record_types, - container_parents=canonical_container_parents, - ) + slack_config: SlackConfig | None = None + _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() + fingerprint = hashlib.sha256( + json.dumps( + {"version": 2, "team_id": principal.team_id, "user_id": principal.user_id}, + sort_keys=True, + ).encode() + ).hexdigest() + return CycleConfiguration.from_source( + fingerprint=fingerprint, + record_types=self.canonical_record_types, + container_parents=self.canonical_container_parents, + ) @classmethod async def create( @@ -202,7 +245,11 @@ async def create( 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 @@ -277,6 +324,7 @@ async def capture_page( parent: SourceRecord | None = None, ) -> CapturePage: """Fetch one page; record and nested reply progress are committed by the pipeline.""" + self._require_principal() try: return await self._validated_capture_page(scope, continuation) except ValidationError: @@ -410,6 +458,7 @@ async def _capture_page_response( async def confirm_absent(self, record: SourceRecord) -> None: """Require provider confirmation before hiding an omitted prior conversation.""" + self._require_principal() native_id = record.identity.native_id try: await self._get("https://slack.com/api/conversations.info", {"channel": native_id}) @@ -579,5 +628,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/tests/test_slack_history_limit.py b/backend/airweave/platform/sources/tests/test_slack_history_limit.py index daaed326f0..71f0e1b931 100644 --- a/backend/airweave/platform/sources/tests/test_slack_history_limit.py +++ b/backend/airweave/platform/sources/tests/test_slack_history_limit.py @@ -7,6 +7,7 @@ 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 @@ -16,6 +17,9 @@ 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( 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/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index c2539dd4ee..4e900a7c97 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -1,6 +1,6 @@ """Slack native page contracts; durable recovery is tested separately against PostgreSQL.""" -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,18 +8,25 @@ from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity 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 SlackApiError, SlackSource -def source(): - return SlackSource( +async def source(): + connector = SlackSource( auth=StaticTokenProvider("token"), logger=MagicMock(), http_client=MagicMock() ) + connector.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U1") + with patch.object( + connector, "_get", AsyncMock(return_value={"ok": True, "team_id": "T1", "user_id": "U1"}) + ): + await connector.validate() + return connector @pytest.mark.asyncio async def test_page_commits_history_before_threads_and_preserves_native_fields(): - connector = source() + connector = await source() connector._get = AsyncMock( side_effect=[ {"channels": [{"id": "C1", "unknown_native_field": "kept"}]}, @@ -59,7 +66,7 @@ async def test_page_commits_history_before_threads_and_preserves_native_fields() ], ) async def test_incomplete_pagination_fails(pages): - connector = source() + connector = await source() connector._get = AsyncMock(side_effect=pages) progress = ScanContinuation() with pytest.raises(ValueError): @@ -74,7 +81,7 @@ async def test_incomplete_pagination_fails(pages): @pytest.mark.asyncio async def test_prior_channel_loss_requires_explicit_provider_confirmation(): - connector = source() + connector = await source() connector._get = AsyncMock(side_effect=SlackApiError("channel_not_found")) await connector.confirm_absent( SourceRecord.model_construct(identity=RecordIdentity(record_type="channel", native_id="C1")) @@ -95,7 +102,7 @@ async def test_rate_limit_honors_entire_wait_and_preserves_proxy_type(kind): from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError - connector = source() + connector = await source() request = httpx.Request("GET", "https://slack.com/api/conversations.list") first = ( AuthProviderRateLimitError(provider_name="composio", retry_after=240) @@ -121,7 +128,7 @@ async def test_rate_limit_honors_entire_wait_and_preserves_proxy_type(kind): async def test_proxy_rate_limit_stops_after_five_attempts(): from airweave.domains.auth_provider.exceptions import AuthProviderRateLimitError - connector = source() + connector = await source() connector.http_client.get = AsyncMock( side_effect=AuthProviderRateLimitError(provider_name="composio", retry_after=240) ) @@ -150,7 +157,7 @@ async def test_queued_thread_failure_preserves_error_meaning(code, expected): ScopeAccessLost, ) - connector = source() + connector = await source() connector._get = AsyncMock(side_effect=SlackApiError(code)) error = {"restart": InvalidScanContinuation, "access": ScopeAccessLost, "error": SlackApiError}[ expected @@ -166,7 +173,7 @@ async def test_queued_thread_failure_preserves_error_meaning(code, expected): @pytest.mark.asyncio async def test_thread_not_found_outside_reply_fetch_does_not_invalidate_inventory(): - connector = source() + connector = await source() connector._get = AsyncMock(side_effect=SlackApiError("thread_not_found")) with pytest.raises(SlackApiError): await connector.capture_page( @@ -202,7 +209,7 @@ def test_message_native_dates_preserve_precision_and_unknowns(native, created, u @pytest.mark.asyncio async def test_conversation_dates_use_distinct_native_units(): - connector = source() + connector = await source() connector._get = AsyncMock( return_value={ "channels": [ From 337a1bed9ea72e6aaf2134eca07490adca3943ed Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:08:52 -0700 Subject: [PATCH 141/266] Add durable owned provisioning intent and capture generation schema --- backend/airweave/models/__init__.py | 2 + backend/airweave/models/owned_provisioning.py | 43 +++++++++++ backend/airweave/models/sync.py | 5 ++ backend/airweave/models/sync_job.py | 3 +- backend/airweave/schemas/sync.py | 2 + backend/airweave/schemas/sync_job.py | 1 + .../versions/0010_owned_provisioning.py | 71 +++++++++++++++++++ 7 files changed, 126 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/models/owned_provisioning.py create mode 100644 backend/alembic/versions/0010_owned_provisioning.py diff --git a/backend/airweave/models/__init__.py b/backend/airweave/models/__init__.py index fb137f2072..c95bf93808 100644 --- a/backend/airweave/models/__init__.py +++ b/backend/airweave/models/__init__.py @@ -18,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 @@ -33,6 +34,7 @@ from .vector_db_deployment_metadata import VectorDbDeploymentMetadata __all__ = [ + "OwnedProvisioning", "AccessControlMembership", "APIKey", "AuthProvider", diff --git a/backend/airweave/models/owned_provisioning.py b/backend/airweave/models/owned_provisioning.py new file mode 100644 index 0000000000..80229008be --- /dev/null +++ b/backend/airweave/models/owned_provisioning.py @@ -0,0 +1,43 @@ +"""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 + ) + 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/sync.py b/backend/airweave/models/sync.py index 918c98b67c..4948495376 100644 --- a/backend/airweave/models/sync.py +++ b/backend/airweave/models/sync.py @@ -41,6 +41,11 @@ 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) diff --git a/backend/airweave/models/sync_job.py b/backend/airweave/models/sync_job.py index 8c0cdc6f64..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) 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/alembic/versions/0010_owned_provisioning.py b/backend/alembic/versions/0010_owned_provisioning.py new file mode 100644 index 0000000000..e296ab62bb --- /dev/null +++ b/backend/alembic/versions/0010_owned_provisioning.py @@ -0,0 +1,71 @@ +"""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("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', '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") From 2f0479a9d27b07f37177b69b87996146b40bd0b5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:10:45 -0700 Subject: [PATCH 142/266] Publish truthful per-part extraction coverage with partial searchable content --- backend/airweave/api/v1/endpoints/records.py | 13 +- .../entities/canonical/extraction_models.py | 64 ++++ .../entities/canonical/gmail_projection.py | 81 ++-- .../entities/canonical/projection_inputs.py | 26 ++ .../entities/canonical/projection_mappers.py | 108 ++++-- .../entities/canonical/projection_store.py | 78 +++- .../domains/entities/canonical/projector.py | 79 ++-- .../entities/canonical/query_models.py | 7 + .../tests/test_calendar_exclusions.py | 8 +- .../tests/test_extraction_coverage.py | 345 ++++++++++++++++++ .../canonical/tests/test_gmail_projection.py | 5 + .../canonical/tests/test_linear_projection.py | 3 + .../canonical/tests/test_notion_projection.py | 1 + .../canonical/tests/test_owned_search.py | 70 +++- backend/airweave/domains/search/owned.py | 50 ++- .../airweave/domains/search/owned_models.py | 5 + .../airweave/models/projection_generation.py | 3 + .../versions/0009_extraction_coverage.py | 23 ++ .../domains/entities/test_attio_projection.py | 6 + .../entities/test_github_projection.py | 5 + .../entities/test_projection_mappers.py | 6 + .../domains/entities/test_workspace_docs.py | 1 + 22 files changed, 890 insertions(+), 97 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/extraction_models.py create mode 100644 backend/airweave/domains/entities/canonical/projection_inputs.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py create mode 100644 backend/alembic/versions/0009_extraction_coverage.py diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index b56a13445e..163b240c77 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -13,10 +13,11 @@ from airweave.core.container import Container from airweave.db.session import get_db from airweave.domains.entities.canonical.calendar_query import CalendarRangeNotCaptured -from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_store import current_extraction from airweave.domains.entities.canonical.query import CanonicalQueryService from airweave.domains.entities.canonical.query_models import ( DocumentRead, + IndexedRecordRead, MailThreadPage, RecordChangePage, RecordFilters, @@ -94,16 +95,20 @@ async def record_changes( return await service.changes(db, ctx.organization.id, sync_id, cursor=cursor, limit=limit) -@router.get("/{sync_id}/records/{record_id}", response_model=SourceRecord) +@router.get("/{sync_id}/records/{record_id}", response_model=IndexedRecordRead) async def read_record( sync_id: UUID, record_id: UUID, db: AsyncSession = Depends(get_db), ctx: ApiContext = Depends(deps.get_context), service: CanonicalQueryService = Depends(deps.get_canonical_query_service), -) -> SourceRecord: +) -> IndexedRecordRead: """Read current committed provider state, including tombstones and completeness.""" - return await service.read(db, ctx.organization.id, sync_id, record_id) + record = await service.read(db, ctx.organization.id, sync_id, record_id) + extraction = await current_extraction( + db, ctx.organization.id, sync_id, record_id, record.revision + ) + return IndexedRecordRead(**record.model_dump(), extraction=extraction) @router.get("/{sync_id}/records/{record_id}/document", response_model=DocumentRead) diff --git a/backend/airweave/domains/entities/canonical/extraction_models.py b/backend/airweave/domains/entities/canonical/extraction_models.py new file mode 100644 index 0000000000..5c070129e1 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/extraction_models.py @@ -0,0 +1,64 @@ +"""Generation-owned extraction evidence, separate from captured original completeness.""" + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator + + +class ExtractionPart(BaseModel): + """Stable source-local identity and format; no private content or download URLs.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + part_index: int = Field(ge=0) + key: str = Field(min_length=1, max_length=2048) + kind: Literal["body", "file", "record"] + media_type: str | None = Field(default=None, max_length=256) + extension: str | None = Field(default=None, max_length=32) + + +class ExtractionOutcome(ExtractionPart): + """Indexed means every required chunk for this part was successfully published.""" + + outcome: Literal["indexed", "unsupported", "unavailable_original"] + reason: Literal["unsupported_format", "original_not_captured"] | None = None + + @model_validator(mode="after") + def reason_matches(self): + """Reject contradictory or free-text explanations.""" + expected = { + "indexed": None, + "unsupported": "unsupported_format", + "unavailable_original": "original_not_captured", + } + if self.reason != expected[self.outcome]: + raise ValueError("Extraction outcome requires its exact bounded reason") + return self + + +class ExtractionCoverage(BaseModel): + """Complete accounting of expected parts in one immutable publication generation.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + parts: tuple[ExtractionOutcome, ...] + + @model_validator(mode="after") + def unique_parts(self): + """Each stable expected part appears once, with no ordinal gaps.""" + if len({p.key for p in self.parts}) != len(self.parts): + raise ValueError("Extraction part keys must be unique") + if tuple(p.part_index for p in self.parts) != tuple(range(len(self.parts))): + raise ValueError("Extraction must account for every expected part exactly once") + return self + + @computed_field + @property + def status(self) -> Literal["complete", "partial", "unavailable", "excluded"]: + """Empty coverage is reserved for intentional exclusions/tombstones.""" + if not self.parts: + return "excluded" + indexed = sum(p.outcome == "indexed" for p in self.parts) + return "complete" if indexed == len(self.parts) else "partial" if indexed else "unavailable" + + def persisted(self) -> dict: + """Computed status remains queryable but validates as derived, not trusted input.""" + return self.model_dump(mode="json", exclude={"status"}) diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py index b3367eae15..1c9002fb6b 100644 --- a/backend/airweave/domains/entities/canonical/gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -10,9 +10,11 @@ from pathlib import Path from airweave.domains.entities.canonical.blob_materializer import read_blob, 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.storage.protocols import StorageBackend -from airweave.platform.entities._base import BaseEntity, Breadcrumb +from airweave.platform.entities._base import Breadcrumb from airweave.platform.entities.gmail import GmailAttachmentEntity, GmailMessageEntity @@ -74,45 +76,50 @@ def __init__(self, record: SourceRecord, storage: StorageBackend, directory: Pat self.record = record self.storage = storage self.directory = directory - self.attachments: list[GmailAttachmentEntity] = [] + self.attachments: list[ProjectionInput] = [] async def attachment(self, part: dict, path: str) -> None: """Project retained attachment bytes; keep missing bytes explicit on the original.""" + filename = part.get("filename") or "attachment" + mime = part.get("mimeType", "application/octet-stream") + suffix = Path(filename).suffix.lower() or mimetypes.guess_extension(mime) or ".bin" + descriptor = ExtractionPart( + part_index=len(self.attachments) + 1, + key=path, + kind="file", + media_type=mime, + extension=suffix, + ) if ( self.record.completeness == "partial" and part.get("body", {}).get("attachmentId") and not any(blob.source_path == path + "/body" for blob in self.record.blobs) ): + self.attachments.append(ProjectionInput(part=descriptor, entity=None)) return content = await _body(part, path, self.record, self.storage) - filename = part.get("filename") or "attachment" - mime = part.get("mimeType", "application/octet-stream") - suffix = Path(filename).suffix.lower() or mimetypes.guess_extension(mime) or ".bin" local = await write_blob(content, self.directory, suffix=suffix) message_id = self.record.identity.native_id - self.attachments.append( - GmailAttachmentEntity( - breadcrumbs=[ - Breadcrumb( - entity_id=f"msg_{message_id}", - name="Email message", - entity_type="GmailMessageEntity", - ) - ], - attachment_key=f"attach_{message_id}_{hashlib.sha256(path.encode()).hexdigest()}", - filename=filename, - message_id=message_id, - attachment_id=part.get("body", {}).get("attachmentId") - or part.get("partId") - or path, - thread_id=self.record.payload["threadId"], - url=f"https://mail.google.com/mail/u/0/#inbox/{message_id}", - size=len(content), - file_type=suffix.lstrip("."), - mime_type=mime, - local_path=str(local), - ) + entity = GmailAttachmentEntity( + breadcrumbs=[ + Breadcrumb( + entity_id=f"msg_{message_id}", + name="Email message", + entity_type="GmailMessageEntity", + ) + ], + attachment_key=f"attach_{message_id}_{hashlib.sha256(path.encode()).hexdigest()}", + filename=filename, + message_id=message_id, + attachment_id=part.get("body", {}).get("attachmentId") or part.get("partId") or path, + thread_id=self.record.payload["threadId"], + url=f"https://mail.google.com/mail/u/0/#inbox/{message_id}", + size=len(content), + file_type=suffix.lstrip("."), + mime_type=mime, + local_path=str(local), ) + self.attachments.append(ProjectionInput(part=descriptor, entity=entity)) async def render(self, part: dict, path: str) -> str: """Select MIME alternatives, preserve mixed bodies, and materialize attachments.""" @@ -134,7 +141,7 @@ async def render(self, part: dict, path: str) -> str: text = _decode_text(content, part) return text if mime == "text/html" else "
" + html.escape(text) + "
" - async def map(self) -> tuple[BaseEntity, ...]: + 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 { @@ -157,13 +164,27 @@ async def map(self) -> tuple[BaseEntity, ...]: 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 (entity, *self.attachments) + return ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, + key="/payload/body", + kind="body", + media_type="text/html", + extension=".html", + ), + entity=entity, + ), + *self.attachments, + ) + ) async def map_gmail( record: SourceRecord, storage: StorageBackend, directory: Path, -) -> tuple[BaseEntity, ...]: +) -> 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/projection_inputs.py b/backend/airweave/domains/entities/canonical/projection_inputs.py new file mode 100644 index 0000000000..0d335ecb77 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_inputs.py @@ -0,0 +1,26 @@ +"""Typed mapper outputs; source descriptors survive filtering and entity-ID stamping.""" + +from pydantic import BaseModel, ConfigDict + +from airweave.domains.entities.canonical.extraction_models import ExtractionPart +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 + + +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 index 31fc9ec669..f8e2df5804 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -13,9 +13,12 @@ from pydantic import JsonValue from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event +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.storage.protocols import StorageBackend -from airweave.platform.entities._base import BaseEntity +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, @@ -238,43 +241,88 @@ async def map_record( record: SourceRecord, source_name: str, storage: StorageBackend, -) -> AsyncIterator[tuple[BaseEntity, ...]]: +) -> AsyncIterator[ProjectionInputs]: """Keep verified local blob files alive only while strict projection consumes them.""" if record.deleted_at is not None or record.content_access != "available": raise ProjectionMappingError("Unavailable records cannot be projected") if excluded_from_search(record, source_name): - yield () + yield ProjectionInputs(parts=()) return with TemporaryDirectory(prefix="airweave-projection-") as temporary: directory = Path(temporary) if source_name == "gmail": from airweave.domains.entities.canonical.gmail_projection import map_gmail - entities = await map_gmail(record, storage, directory) - elif source_name == "github": - from airweave.domains.entities.canonical.github_projection import map_github - - entities = await map_github(record, storage, directory) - elif source_name == "linear": - from airweave.domains.entities.canonical.linear_projection import map_linear - - entities = await map_linear(record, storage, directory) - elif source_name == "attio": - from airweave.domains.entities.canonical.attio_projection import map_attio - - entities = map_attio(record) - elif source_name == "notion": - from airweave.domains.entities.canonical.notion_projection import map_notion - - entities = map_notion(record) - elif source_name == "google_drive": - entities = await _drive(record, storage, directory) - elif source_name == "google_calendar": - entities = _calendar(record) - elif source_name == "slack": - entities = _slack(record) - elif source_name == "wispr": - entities = _wispr(record) + yield await map_gmail(record, storage, directory) + return else: - raise ProjectionMappingError("Source has no audited search projection mapper") - yield entities + entities = await _map_entities(record, source_name, storage, directory) + parts = tuple(_projection_input(index, entity) for index, entity in enumerate(entities)) + if ( + source_name == "github" + and record.identity.record_type == "file" + and record.completeness != "complete" + ): + parts += ( + ProjectionInput( + part=ExtractionPart( + part_index=len(parts), + key="/file/content", + kind="file", + extension=Path(record.identity.native_id).suffix.lower() or None, + ), + entity=None, + ), + ) + yield ProjectionInputs(parts=parts) + + +async def _map_entities( + record: SourceRecord, source_name: str, storage: StorageBackend, directory: Path +) -> tuple[BaseEntity, ...]: + """Retain existing audited provider mappings with one part per explicit entity.""" + if source_name == "github": + from airweave.domains.entities.canonical.github_projection import map_github + + entities = await map_github(record, storage, directory) + elif source_name == "linear": + from airweave.domains.entities.canonical.linear_projection import map_linear + + entities = await map_linear(record, storage, directory) + elif source_name == "attio": + from airweave.domains.entities.canonical.attio_projection import map_attio + + entities = map_attio(record) + elif source_name == "notion": + from airweave.domains.entities.canonical.notion_projection import map_notion + + entities = map_notion(record) + elif source_name == "google_drive": + entities = await _drive(record, storage, directory) + elif source_name == "google_calendar": + entities = _calendar(record) + elif source_name == "slack": + entities = _slack(record) + elif source_name == "wispr": + entities = _wispr(record) + else: + raise ProjectionMappingError("Source has no audited search projection mapper") + return entities + + +def _projection_input(index: int, entity: BaseEntity) -> ProjectionInput: + """Retain native mapper identity before generation-specific stamping.""" + populate_base_fields(entity) + file = entity if isinstance(entity, FileEntity) else None + return ProjectionInput( + part=ExtractionPart( + part_index=index, + key=entity.entity_id, + kind="file" if file is not None else "record", + media_type=file.mime_type if file is not None else None, + extension=Path(file.local_path).suffix.lower() + if file is not None and file.local_path + else None, + ), + entity=entity, + ) diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index f7dd2bff4d..800cc794ce 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -3,11 +3,12 @@ from datetime import datetime, timedelta, timezone from uuid import UUID -from sqlalchemy import and_, or_, select +from sqlalchemy import and_, exists, or_, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event +from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.projection_models import ( ProjectionDocument, ProjectionLocator, @@ -31,6 +32,18 @@ def publication_matches(locator: ProjectionLocator): Entity.indexed_generation == locator.generation, Entity.deleted_at.is_(None), content_is_available(), + exists( + select(ProjectionGeneration.id).where( + ProjectionGeneration.id == locator.generation, + ProjectionGeneration.retired_at.is_(None), + or_( + ProjectionGeneration.extraction_coverage.is_(None), + ProjectionGeneration.extraction_coverage.contains( + {"parts": [{"part_index": locator.part_index, "outcome": "indexed"}]} + ), + ), + ) + ), ) @@ -121,9 +134,12 @@ async def prepare( generation: UUID, collection_id: UUID, documents: tuple[ProjectionDocument, ...], + *, + coverage: ExtractionCoverage | None = None, ) -> bool: """Commit exact immutable deletion identities before the first remote feed.""" identities = set() + indexed_parts = set() for document in documents: try: locator = projection_document_locator( @@ -140,6 +156,12 @@ async def prepare( if not valid or (document.schema_name, document.document_id) in identities: raise ValueError("Projection manifest must contain unique IDs of this generation") identities.add((document.schema_name, document.document_id)) + indexed_parts.add(locator.part_index) + if coverage is not None and indexed_parts != { + p.part_index for p in coverage.parts if p.outcome == "indexed" + }: + raise ValueError("Projection documents must cover exactly the indexed parts") + extraction = coverage.persisted() if coverage is not None else None manifest = [document.model_dump() for document in documents] async with UnitOfWork(db): if await self._current(db, work) is None: @@ -148,6 +170,7 @@ async def prepare( if prior is not None: if ( prior.documents != manifest + or prior.extraction_coverage != extraction or prior.record_id != work.record.id or prior.revision != work.record.revision or prior.pipeline_version != work.pipeline_version @@ -166,6 +189,7 @@ async def prepare( revision=work.record.revision, pipeline_version=work.pipeline_version, documents=manifest, + extraction_coverage=extraction, next_gc_at=datetime.now(timezone.utc) + timedelta(hours=1), ) ) @@ -183,9 +207,12 @@ async def publish( exclusion = work.record.identity.record_type == "event" and is_cancelled_recurring_event( work.record.payload ) - no_content = (work.record.deleted_at is not None or exclusion - or work.record.identity.record_type == "event_occurrence") - if chunk_count < 0 or (not no_content and chunk_count == 0): + no_content = ( + work.record.deleted_at is not None + or exclusion + or work.record.identity.record_type == "event_occurrence" + ) + if chunk_count < 0: raise ValueError("Active records require a nonempty complete projection") if no_content and chunk_count != 0: raise ValueError("Deleted records and recurrence exclusions must publish no content") @@ -211,6 +238,21 @@ async def publish( or len(attempt.documents) != chunk_count ): return False + coverage = ( + ExtractionCoverage.model_validate(attempt.extraction_coverage) + if attempt.extraction_coverage is not None + else None + ) + if ( + not no_content + and chunk_count == 0 + and (coverage is None or not coverage.parts or coverage.status != "unavailable") + ): + raise ValueError( + "Active records require indexed content or explicit unavailable extraction" + ) + if no_content and coverage is not None and coverage.parts: + raise ValueError("Excluded records must have empty extraction coverage") if work.previous_generation is not None: previous = await db.get(ProjectionGeneration, work.previous_generation) if previous is not None: @@ -231,3 +273,31 @@ async def fail(self, db: AsyncSession, work: ProjectionWork, message: str) -> No if entity is not None: entity.projection_error = message[:1000] await db.flush() + + +async def current_extraction( + db: AsyncSession, organization_id: UUID, sync_id: UUID, record_id: UUID, revision: int +) -> ExtractionCoverage | None: + """Only the visible current revision/pipeline/publication exposes extraction evidence.""" + raw = await db.scalar( + select(ProjectionGeneration.extraction_coverage) + .join(Entity, Entity.indexed_generation == ProjectionGeneration.id) + .join(Sync, Sync.id == Entity.sync_id) + .where( + Entity.id == record_id, + Entity.record_revision == revision, + Entity.organization_id == organization_id, + Entity.sync_id == sync_id, + Sync.organization_id == organization_id, + ProjectionGeneration.organization_id == organization_id, + ProjectionGeneration.record_id == Entity.id, + ProjectionGeneration.revision == Entity.record_revision, + ProjectionGeneration.pipeline_version == Sync.index_pipeline_version, + ProjectionGeneration.retired_at.is_(None), + Entity.indexed_revision == Entity.record_revision, + Entity.indexed_pipeline_version == Sync.index_pipeline_version, + Entity.deleted_at.is_(None), + content_is_available(), + ) + ) + return ExtractionCoverage.model_validate(raw) if raw is not None else None diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index ddbb09a20c..19d52e55f8 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -8,7 +8,12 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.core.logging import ContextualLogger +from airweave.domains.entities.canonical.extraction_models import ( + ExtractionCoverage, + ExtractionOutcome, +) from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.projection_inputs import ProjectionInputs from airweave.domains.entities.canonical.projection_models import ( ProjectionBatchResult, ProjectionDocument, @@ -19,6 +24,7 @@ from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata from airweave.domains.storage.protocols import StorageBackend +from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields from airweave.platform.destinations.vespa.destination import VespaDestination @@ -88,34 +94,17 @@ async def project_one( generation = uuid4() chunks = [] + coverage = ExtractionCoverage(parts=()) no_documents = work.record.deleted_at is not None or excluded_from_search( work.record, source_name ) if not no_documents: async with map_record(work.record, source_name, self._storage) as mapped: - if not mapped: + if not mapped.parts: raise ValueError("Projection mapper returned no required content") - for part_index, entity in enumerate(mapped): - populate_base_fields(entity) - if entity.airweave_system_metadata is None: - entity.airweave_system_metadata = AirweaveSystemMetadata() - locator = ProjectionLocator( - record_id=work.record.id, - revision=work.record.revision, - pipeline_version=work.pipeline_version, - generation=generation, - part_index=part_index, - ) - entity.entity_id = locator.encode() - meta = entity.airweave_system_metadata - meta.source_name = source_name - meta.entity_type = type(entity).__name__ - meta.sync_id = work.record.sync_id - meta.sync_job_id = None - meta.db_entity_id = work.record.id - stamp_search_metadata(meta, work.record) + selected, coverage = _select_inputs(mapped, work, source_name, generation) chunks = await self._processor.process( - list(mapped), + selected, ProjectionContext(logger, source_name), ProjectionRuntime(StrictProjectionTracker()), strict=True, @@ -134,14 +123,14 @@ async def project_one( ) async with self._sessions() as db: if not await self._store.prepare( - db, work, generation, destination.collection_id, manifest + db, work, generation, destination.collection_id, manifest, coverage=coverage ): return False await destination.feed_prepared(prepared) if no_documents: async with self._sessions() as db: if not await self._store.prepare( - db, work, generation, destination.collection_id, () + db, work, generation, destination.collection_id, (), coverage=coverage ): return False async with self._sessions() as db: @@ -190,3 +179,47 @@ async def batch( superseded=superseded, failed=failed, ) + + +def _select_inputs( + mapped: ProjectionInputs, work: ProjectionWork, source_name: str, generation: UUID +) -> tuple[list[BaseEntity], ExtractionCoverage]: + """Classify only deterministic omissions, preserving stable pre-filter part ordinals.""" + selected = [] + outcomes = [] + for item in mapped.parts: + part, entity = item.part, item.entity + if entity is None: + outcome, reason = "unavailable_original", "original_not_captured" + elif ( + part.kind == "file" + and part.extension is not None + and part.extension not in SUPPORTED_FILE_EXTENSIONS + ): + outcome, reason = "unsupported", "unsupported_format" + else: + outcome, reason = "indexed", None + selected.append(entity) + outcomes.append(ExtractionOutcome(**part.model_dump(), outcome=outcome, reason=reason)) + if outcome != "indexed": + continue + part_index = part.part_index + populate_base_fields(entity) + if entity.airweave_system_metadata is None: + entity.airweave_system_metadata = AirweaveSystemMetadata() + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=part_index, + ) + entity.entity_id = locator.encode() + meta = entity.airweave_system_metadata + meta.source_name = source_name + meta.entity_type = type(entity).__name__ + meta.sync_id = work.record.sync_id + meta.sync_job_id = None + meta.db_entity_id = work.record.id + stamp_search_metadata(meta, work.record) + return selected, ExtractionCoverage(parts=tuple(outcomes)) diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index 204e30b753..fa21dc5498 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -7,11 +7,18 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue from airweave.domains.entities.canonical.coverage import CaptureCoverage +from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.platform.sources.records.workspace_manifest import WorkspaceManifestV1 +class IndexedRecordRead(SourceRecord): + """Current original plus independently attested current extraction coverage.""" + + extraction: ExtractionCoverage | None = None + + class DocumentRead(BaseModel): """One authorized stored revision's native document and representation coverage.""" diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index 302d46cde4..15dee1cd85 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -49,6 +49,7 @@ async def test_cancelled_exception_read_project_reinstate_and_revoke(database, s if w.record.id == row.id ) async with map_record(current, "google_calendar", AsyncMock()) as mapped: + mapped = mapped.entities assert mapped == () processor = MagicMock() processor.process = AsyncMock() @@ -89,6 +90,7 @@ async def test_cancelled_exception_read_project_reinstate_and_revoke(database, s ) current = await queries.read(db, fence.organization_id, fence.sync_id, row.id) async with map_record(current, "google_calendar", AsyncMock()) as mapped: + mapped = mapped.entities assert len(mapped) == 1 and mapped[0].status == "confirmed" async with database() as db: await capture.capture(db, CaptureBatch(fence=fence, records=(excluded,))) @@ -143,8 +145,10 @@ async def test_live_event_cannot_publish_empty_or_exclusion_nonempty(database, s ) } ) - with pytest.raises(ValueError, match="nonempty"): - await store.publish(db, active, uuid4(), 0) + generation = uuid4() + assert await store.prepare(db, active, generation, uuid4(), ()) + with pytest.raises(ValueError, match="indexed content"): + await store.publish(db, active, generation, 0) async def test_operational_occurrences_publish_zero_documents_without_retry(database, source): diff --git a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py new file mode 100644 index 0000000000..a4ab72a964 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py @@ -0,0 +1,345 @@ +"""Real SQL and conversion; synthetic originals, embeddings and remote feed.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import fitz +import pytest +from sqlalchemy import select, update + +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.extraction_models import ( + ExtractionCoverage, + ExtractionOutcome, +) +from airweave.domains.entities.canonical.projection_models import ProjectionLocator +from airweave.domains.entities.canonical.projection_store import ( + CanonicalProjectionStore, + current_extraction, + publication_matches, +) +from airweave.domains.entities.canonical.projector import CanonicalProjector +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_gmail_projection import part +from airweave.domains.search.owned import OwnedSearchService +from airweave.domains.search.owned_models import OwnedSearchRequest +from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor +from airweave.models import Entity, Sync +from airweave.models.projection_generation import ProjectionGeneration +from airweave.platform.destinations.vespa.transformer import EntityTransformer + + +def original(*attachments, completeness="complete"): + return observation( + identity=RecordIdentity(record_type="message", native_id="m1"), + completeness=completeness, + payload={ + "id": "m1", + "threadId": "t1", + "internalDate": "1000", + "payload": { + "mimeType": "multipart/mixed", + "parts": [ + part(b"Intact fundraising email body with useful context."), + *attachments, + ], + }, + }, + ) + + +def destination(): + result = MagicMock(collection_id=uuid4(), feed_prepared=AsyncMock()) + result.prepare_documents = lambda chunks: { + "base_entity": [ + EntityTransformer(collection_id=result.collection_id).transform(c) for c in chunks + ] + } + return result + + +def projector(database, storage): + return CanonicalProjector( + CanonicalProjectionStore(), + database, + ChunkEmbedProcessor(ConverterRegistry(), FakeDenseEmbedder(), FakeSparseEmbedder()), + storage, + ) + + +async def test_message_pdf_video_and_missing_part_publish_truthful_coverage( + database, source, tmp_path +): + service, fence = source + pdf = fitz.open() + page = pdf.new_page() + page.insert_text( + (40, 40), "Retained PDF investment discussion and detailed fundraising plans. " * 4 + ) + data = pdf.tobytes() + pdf.close() + item = original( + part(data, mime="application/pdf", filename="deck.pdf"), + part(b"synthetic video bytes", mime="video/mp4", filename="meeting.mp4"), + { + "mimeType": "application/pdf", + "filename": "large.pdf", + "body": {"attachmentId": "missing", "size": 500000000}, + }, + completeness="partial", + ) + await capture(database, service, fence, item) + store = CanonicalProjectionStore() + target = destination() + result = await projector(database, FilesystemBackend(tmp_path)).batch( + fence.organization_id, fence.sync_id, "gmail", target, 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 [p.outcome for p in coverage.parts] == [ + "indexed", + "indexed", + "unsupported", + "unavailable_original", + ] + assert coverage.parts[1].key == "/payload/parts/1" + assert ( + coverage.parts[1].kind == "file" and coverage.parts[1].media_type == "application/pdf" + ) + assert row.completeness == "partial" and row.source_payload == item.payload + assert not await store.pending(db, fence.organization_id, fence.sync_id) + locator = ProjectionLocator( + record_id=row.id, + revision=row.record_revision, + pipeline_version=row.indexed_pipeline_version, + generation=row.indexed_generation, + part_index=1, + ) + assert ( + await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator))) + == row.id + ) + assert ( + await db.scalar( + select(Entity.id) + .join(Sync) + .where(publication_matches(locator.model_copy(update={"part_index": 2}))) + ) + is None + ) + ctx = MagicMock() + ctx.organization.id = fence.organization_id + counts = await OwnedSearchService._coverage( + db, ctx, OwnedSearchRequest(query="fundraising", sync_ids=(fence.sync_id,)) + ) + assert counts[0].partially_indexed_records == 1 and counts[0].pending_records == 0 + # New native revision withdraws every old part and its coverage before new extraction. + await capture( + database, service, fence, original(part(b"new video", mime="video/mp4", filename="new.mp4")) + ) + async with database() as db: + assert ( + await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator))) + is None + ) + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + is None + ) + assert ( + await projector(database, FilesystemBackend(tmp_path)).batch( + fence.organization_id, fence.sync_id, "gmail", target, logger + ) + ).published == 1 + async with database() as db: + assert ( + await db.scalar(select(Entity.id).join(Sync).where(publication_matches(locator))) + is None + ) + old = await db.get(ProjectionGeneration, locator.generation) + assert old.retired_at is not None + + +async def test_supported_attachment_conversion_failure_and_feed_failure_stay_pending( + database, source, tmp_path +): + service, fence = source + await capture( + database, + service, + fence, + original(part(b"broken binary\x00", mime="application/pdf", filename="broken.pdf")), + ) + target = destination() + project = projector(database, FilesystemBackend(tmp_path)) + assert ( + await project.batch(fence.organization_id, fence.sync_id, "gmail", target, logger) + ).failed == 1 + target.feed_prepared.assert_not_awaited() + await capture( + database, service, fence, original(part(b"video", mime="video/mp4", filename="clip.mp4")) + ) + target.feed_prepared.side_effect = ConnectionError("synthetic interruption") + assert ( + await project.batch(fence.organization_id, fence.sync_id, "gmail", target, logger) + ).failed == 1 + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id)) + assert row.indexed_generation is None + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + is None + ) + assert ( + len(await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id)) + == 1 + ) + target.feed_prepared.side_effect = None + assert ( + await project.batch(fence.organization_id, fence.sync_id, "gmail", target, logger) + ).published == 1 + + +async def test_unsupported_only_publication_is_explicit_and_retries_on_pipeline_upgrade( + database, source, tmp_path +): + service, fence = source + from airweave.domains.entities.canonical.tests.test_gmail_projection import blob, record + + captured = record({}, sync_id=fence.sync_id) + ref = blob(captured, b"video bytes") + storage = FilesystemBackend(tmp_path) + await storage.write_file(ref.key, b"video bytes") + await capture( + database, + service, + fence, + observation( + identity=RecordIdentity(record_type="file", native_id="video"), + payload={"id": "video", "name": "movie.mp4", "mimeType": "video/mp4"}, + blobs=(ref,), + ), + ) + target = destination() + assert ( + await projector(database, storage).batch( + fence.organization_id, fence.sync_id, "google_drive", target, logger + ) + ).published == 1 + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id)) + assert row.indexed_chunk_count == 0 and row.completeness == "complete" + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + ).status == "unavailable" + await db.execute( + update(Sync).where(Sync.id == fence.sync_id).values(index_pipeline_version=2) + ) + await db.commit() + assert ( + len(await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id)) + == 1 + ) + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + is None + ) + + +def test_coverage_rejects_missing_duplicate_and_contradictory_parts(): + one = ExtractionOutcome(part_index=0, key="/body", kind="body", outcome="indexed") + for parts in [(one, one), (one.model_copy(update={"part_index": 1}),)]: + with pytest.raises(ValueError): + ExtractionCoverage(parts=parts) + with pytest.raises(ValueError): + ExtractionOutcome(part_index=0, key="a", kind="file", outcome="unsupported", reason=None) + + +async def test_coverage_manifest_is_exact_immutable_and_capture_fences_publication( + database, source +): + from airweave.domains.entities.canonical.projection_models import ( + ProjectionDocument, + scope_projection_document_id, + ) + + service, fence = source + await capture(database, service, fence, original()) + store = CanonicalProjectionStore() + async with database() as db: + work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] + generation, collection = uuid4(), uuid4() + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=0, + ) + documents = ( + ProjectionDocument( + schema_name="base_entity", + document_id=scope_projection_document_id( + fence.sync_id, collection, f"Entity_{locator.encode()}__chunk_0" + ), + ), + ) + indexed = ExtractionOutcome(part_index=0, key="/body", kind="body", outcome="indexed") + async with database() as db: + with pytest.raises(ValueError, match="exactly"): + await store.prepare( + db, work, generation, collection, documents, coverage=ExtractionCoverage(parts=()) + ) + assert await store.prepare( + db, + work, + generation, + collection, + documents, + coverage=ExtractionCoverage(parts=(indexed,)), + ) + with pytest.raises(ValueError, match="immutable"): + await store.prepare( + db, + work, + generation, + collection, + documents, + coverage=ExtractionCoverage( + parts=(indexed.model_copy(update={"key": "different"}),) + ), + ) + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, work.record.id, work.record.revision + ) + is None + ) + await capture( + database, service, fence, original(part(b"video", mime="video/mp4", filename="clip.mp4")) + ) + async with database() as db: + assert not await store.publish(db, work, generation, 1) + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, work.record.id, work.record.revision + ) + is None + ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py index 46bdec7632..62792da390 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -76,6 +76,7 @@ async def test_plain_charset_body_and_quoted_address_are_preserved(tmp_path): ) before = source.model_dump() entities = await map_gmail(source, AsyncMock(), tmp_path) + entities = entities.entities assert len(entities) == 1 assert entities[0].to == ['"Doe, Jane" ', "bob@example.com"] assert "café <hello>" in next(tmp_path.glob("*.html")).read_text() @@ -94,6 +95,7 @@ async def test_alternative_prefers_full_html_not_plain_duplicate_or_snippet(tmp_ } ) 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 @@ -118,6 +120,7 @@ async def test_external_blob_and_duplicate_filename_attachments_get_distinct_ide 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 @@ -174,6 +177,7 @@ async def test_missing_attachment_keeps_available_body_searchable_and_original_p 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 @@ -209,6 +213,7 @@ async def test_related_html_branch_does_not_get_replaced_by_plain_alternative(tm } ) entities = await map_gmail(source, AsyncMock(), tmp_path) + entities = entities.entities assert len(entities) == 2 from pathlib import Path diff --git a/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py b/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py index 8ddc0417c8..de75b08699 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_linear_projection.py @@ -72,6 +72,7 @@ async def test_native_issue_and_comment_render_without_provider_or_blob_access() 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() @@ -84,6 +85,7 @@ async def test_native_issue_and_comment_render_without_provider_or_blob_access() 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 @@ -111,6 +113,7 @@ async def test_retained_file_verified_and_disposable(): 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" diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py index c871845a10..5ea7f278ff 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py @@ -76,6 +76,7 @@ def block(kind, content, **fields): async def mapped(record): storage = AsyncMock() async with map_record(record, "notion", storage) as entities: + entities = entities.entities assert len(entities) == 1 entity = entities[0] assert not storage.mock_calls diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index e434dd109c..32fa48c63f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -22,6 +22,7 @@ from airweave.domains.search.owned import OwnedSearchService from airweave.domains.search.types import SearchResults 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.vector_db_deployment_metadata import VectorDbDeploymentMetadata @@ -109,6 +110,13 @@ async def session(): async with database() as db: yield db + from airweave.domains.entities.canonical.query import CanonicalQueryService + from airweave.domains.entities.canonical.query_store import CanonicalQueryStore + from airweave.domains.entities.canonical.store import CanonicalRecordStore + + app.dependency_overrides[deps.get_canonical_query_service] = lambda: CanonicalQueryService( + CanonicalRecordStore(), CanonicalQueryStore(), "test-extraction-key" + ) app.dependency_overrides[get_db] = session app.dependency_overrides[deps.get_context] = lambda: ctx app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace(owned_search=service) @@ -485,4 +493,64 @@ async def test_capture_discovery_coverage_is_separate_from_retrieval( "policies": {"event": "discovery_only"}, "discovery": "incomplete", } - assert not body["retrieval_incomplete"] + assert body["sources"][0]["extraction_unknown_records"] == 1 + assert body["retrieval_incomplete"] + + +async def test_partial_extraction_http_hit_original_and_nonindexed_part_gate( + database, indexed, http_search +): + from airweave.domains.entities.canonical.extraction_models import ( + ExtractionCoverage, + ExtractionOutcome, + ) + from airweave.models.projection_generation import ProjectionGeneration + + fence, locator, _ = indexed + coverage = ExtractionCoverage( + parts=( + ExtractionOutcome(part_index=0, key="/body", kind="body", outcome="indexed"), + ExtractionOutcome( + part_index=1, + key="/payload/parts/1", + kind="file", + media_type="video/mp4", + extension=".mp4", + outcome="unsupported", + reason="unsupported_format", + ), + ) + ) + async with database() as db: + await db.execute( + update(ProjectionGeneration) + .where(ProjectionGeneration.id == locator.generation) + .values(extraction_coverage=coverage.persisted()) + ) + await db.commit() + client, vector, _, _, _ = http_search + valid = hit(fence, locator.encode()) + invalid = hit(fence, locator.model_copy(update={"part_index": 1}).encode()) + invalid.textual_representation = "unsupported attachment must never escape" + vector.seed_results(SearchResults(results=[valid, invalid])) + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + assert response.status_code == 200, response.text + payload = response.json() + assert len(payload["items"]) == 1 and payload["retrieval_incomplete"] + assert payload["items"][0]["extraction"] == coverage.model_dump(mode="json") + assert "unsupported attachment must never escape" not in str(payload["items"]) + assert payload["sources"][0]["partially_indexed_records"] == 1 + response = await client.get(f"/sync/{fence.sync_id}/records/{locator.record_id}") + assert response.status_code == 200, response.text + assert response.json()["completeness"] == "complete" + assert response.json()["extraction"] == coverage.model_dump(mode="json") + async with database() as db: + await db.execute( + update(Entity).where(Entity.id == locator.record_id).values(record_revision=2) + ) + await db.commit() + response = await client.get(f"/sync/{fence.sync_id}/records/{locator.record_id}") + assert response.json()["extraction"] is None diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index e2c4403a8a..502e6a02a2 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -10,6 +10,7 @@ from airweave.api.context import ApiContext 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 publication_matches from airweave.domains.entities.canonical.requests import RecordIdentity @@ -30,6 +31,7 @@ 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 @@ -110,7 +112,13 @@ async def search( or exclusions > 0 or postfiltered > 0 or len(ranked) > request.limit - or any(row.pending_records for row in sources), + 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 @@ -241,6 +249,15 @@ async def _enrich( ) ).all() 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 @@ -254,6 +271,13 @@ async def _enrich( ): 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 @@ -275,6 +299,7 @@ async def _enrich( 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" @@ -383,9 +408,25 @@ async def _coverage( 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)) + 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, @@ -396,7 +437,7 @@ async def _coverage( ) .group_by(Entity.sync_id) ) - counts = {sync: (active, pending) for sync, active, pending in rows} + 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( @@ -404,6 +445,9 @@ async def _coverage( 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 index 48facdb70f..62ce497f61 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -6,6 +6,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator from airweave.domains.entities.canonical.coverage import CaptureCoverage +from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.search.types import RetrievalStrategy @@ -56,6 +57,7 @@ class OwnedSearchHit(BaseModel): source_created_at: AwareDatetime | None source_updated_at: AwareDatetime | None completeness: Literal["complete", "partial", "metadata_only"] + extraction: ExtractionCoverage | None = None class OwnedSearchCoverage(BaseModel): @@ -65,6 +67,9 @@ class OwnedSearchCoverage(BaseModel): 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 OwnedSearchResponse(BaseModel): diff --git a/backend/airweave/models/projection_generation.py b/backend/airweave/models/projection_generation.py index c87764ef1d..36dfd7d591 100644 --- a/backend/airweave/models/projection_generation.py +++ b/backend/airweave/models/projection_generation.py @@ -21,6 +21,9 @@ class ProjectionGeneration(Base): 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 + ) 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) 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/tests/unit/domains/entities/test_attio_projection.py b/backend/tests/unit/domains/entities/test_attio_projection.py index 0a1b2b0ae0..e99ed8df2d 100644 --- a/backend/tests/unit/domains/entities/test_attio_projection.py +++ b/backend/tests/unit/domains/entities/test_attio_projection.py @@ -62,6 +62,7 @@ async def test_membership_search_keeps_entry_identity_and_never_fetches_record() before = captured.model_dump() storage = AsyncMock() async with map_record(captured, "attio", storage) as entities: + entities = entities.entities assert len(entities) == 1 result = entities[0] assert result.native_id == ENTRY @@ -89,6 +90,7 @@ async def test_note_uses_current_plaintext_field_and_its_record_parent(): RecordIdentity(record_type="record", native_id=RECORD, container_id=OBJECT), ) async with map_record(captured, "attio", AsyncMock()) as entities: + entities = entities.entities assert entities[0].text == "Discussed fundraising" assert entities[0].title == "Call" @@ -131,6 +133,7 @@ async def test_record_attributes_and_native_url(url): pass else: async with map_record(captured, "attio", AsyncMock()) as entities: + entities = entities.entities assert entities[0].title == "Sam Example" assert "custom_score" in entities[0].text assert entities[0].web_url == url @@ -149,7 +152,9 @@ async def test_two_memberships_for_one_crm_record_remain_distinct(): RecordIdentity(record_type="list", native_id=second_list), ) async with map_record(first, "attio", AsyncMock()) as one: + one = one.entities async with map_record(second, "attio", AsyncMock()) as two: + two = two.entities assert one[0].native_id != two[0].native_id assert RECORD in one[0].text and RECORD in two[0].text assert one[0].original_kind == two[0].original_kind == "entry" @@ -166,6 +171,7 @@ async def test_container_projection_does_not_invent_children(kind): payload.update(name="Fundraising") captured = original(kind, native_id, payload) async with map_record(captured, "attio", AsyncMock()) as entities: + entities = entities.entities assert len(entities) == 1 assert "metadata only" in entities[0].content_coverage assert entities[0].web_url == "" diff --git a/backend/tests/unit/domains/entities/test_github_projection.py b/backend/tests/unit/domains/entities/test_github_projection.py index aaf630ef76..cc3b446559 100644 --- a/backend/tests/unit/domains/entities/test_github_projection.py +++ b/backend/tests/unit/domains/entities/test_github_projection.py @@ -51,6 +51,7 @@ async def test_native_conversation_text_and_exact_kind(kind): ) storage = AsyncMock() async with map_record(item, "github", storage) as views: + views = views.entities assert views[0].text == "Native **Markdown**" assert views[0].resource_kind == kind assert views[0].state == "APPROVED" @@ -89,6 +90,7 @@ async def test_code_uses_verified_owned_blob_and_temporary_file(): storage = AsyncMock() storage.read_file.return_value = content async with map_record(item, "github", storage) as views: + views = views.entities local = Path(views[0].local_path) assert local.read_bytes() == content assert views[0].commit_id == "b" * 40 @@ -121,4 +123,7 @@ async def test_missing_complete_bytes_keeps_projection_pending(): pass partial = item.model_copy(update={"completeness": "partial"}) async with map_record(partial, "github", AsyncMock()) as views: + assert views.parts[-1].entity is None and views.parts[-1].part.kind == "file" + assert views.parts[-1].part.key == "/file/content" + views = views.entities assert views[0].content_coverage == "partial" and views[0].text == "" diff --git a/backend/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index e8bae1ccf9..78b757da04 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -51,6 +51,7 @@ async def test_calendar_keeps_master_recurrence_and_native_container(): } item = record("event", payload, native_id="event", container_id="calendar") async with map_record(item, "google_calendar", AsyncMock()) as entities: + entities = entities.entities assert entities[0].calendar_key == "calendar" assert entities[0].recurrence == ["RRULE:FREQ=WEEKLY"] assert entities[0].start_date == "2026-10-01" @@ -67,6 +68,7 @@ async def test_slack_history_shape_uses_container_without_provider_lookup(): ) storage = AsyncMock() async with map_record(item, "slack", storage) as entities: + entities = entities.entities assert entities[0].channel_id == "C1" assert entities[0].text == "hello" assert entities[0].updated_at.timestamp() == 124 @@ -107,6 +109,7 @@ async def test_wispr_ranges_do_not_duplicate_default_content(): ] } async with map_record(record("meeting", payload), "wispr", AsyncMock()) as entities: + entities = entities.entities assert entities[0].notes == "notes" assert entities[0].transcript == "abc\ndef" assert "guidance" not in entities[0].transcript @@ -156,6 +159,7 @@ async def test_drive_uses_verified_owned_bytes_and_cleans_materialization(): storage = AsyncMock() storage.read_file.return_value = content async with map_record(item, "google_drive", storage) as entities: + entities = entities.entities path = Path(entities[0].local_path) assert path.read_bytes() == content assert path.name == digest + ".txt" @@ -209,6 +213,7 @@ async def test_owned_drive_spreadsheet_uses_real_converter(): storage = AsyncMock() storage.read_file.return_value = content async with map_record(item, "google_drive", storage) as entities: + entities = entities.entities path = entities[0].local_path assert Path(path).suffix == ".xlsx" tracker = AsyncMock() @@ -255,6 +260,7 @@ async def test_wispr_scratchpad_retains_all_text_ranges_without_meeting_fields() before = original.model_dump() storage = AsyncMock() async with map_record(original, "wispr", storage) as entities: + entities = entities.entities assert entities[0].note_id == "note" assert entities[0].content == "abc\ndef" assert entities[0].web_url == "" diff --git a/backend/tests/unit/domains/entities/test_workspace_docs.py b/backend/tests/unit/domains/entities/test_workspace_docs.py index 642effcb1e..bc323a2a93 100644 --- a/backend/tests/unit/domains/entities/test_workspace_docs.py +++ b/backend/tests/unit/domains/entities/test_workspace_docs.py @@ -146,6 +146,7 @@ async def test_recursive_tabs_preserve_original_structure_without_edit_revision( async def test_projection_reads_native_once_and_never_indexes_export(): record, storage, _ = capture() async with map_record(record, "google_drive", storage) as entities: + entities = entities.entities assert len(entities) == 1 entity = entities[0] assert entity.file_id == "doc" From 4b1f67ca4ab67ab5ef210f900f2844dc06b1367b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:11:15 -0700 Subject: [PATCH 143/266] Fence source construction with authoritative job generation admission --- .../airweave/domains/sync_pipeline/factory.py | 8 ++- .../sync_pipeline/tests/test_factory.py | 62 ++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index 60b62391b7..2ce2a3307a 100644 --- a/backend/airweave/domains/sync_pipeline/factory.py +++ b/backend/airweave/domains/sync_pipeline/factory.py @@ -35,7 +35,7 @@ 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 +from airweave.domains.entities.canonical.store import CanonicalRecordStore, StaleWriter from airweave.domains.entities.protocols import ( EntityCountRepositoryProtocol, EntityRepositoryProtocol, @@ -169,7 +169,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 @@ -192,6 +196,8 @@ 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) canonical_source = ( source_result.source diff --git a/backend/airweave/domains/sync_pipeline/tests/test_factory.py b/backend/airweave/domains/sync_pipeline/tests/test_factory.py index 4e396d275a..2daf27e40a 100644 --- a/backend/airweave/domains/sync_pipeline/tests/test_factory.py +++ b/backend/airweave/domains/sync_pipeline/tests/test_factory.py @@ -117,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, @@ -161,6 +168,11 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): 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" @@ -217,6 +229,8 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): ) assert orchestrator is not None + 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() @@ -613,3 +627,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() From 63c99f836ddeb255fce9c53cc6b28481b3794458 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:14:56 -0700 Subject: [PATCH 144/266] Fence capture by verified connection generation and integrate extraction fixtures --- .../domains/entities/canonical/store.py | 54 ++++++++++++- .../entities/canonical/tests/conftest.py | 2 + .../tests/test_provisioning_fence.py | 75 +++++++++++++++++++ .../canonical/tests/test_slack_recovery.py | 73 +++++++++--------- backend/airweave/platform/configs/config.py | 12 ++- backend/tests/live/canonical_capture.py | 2 + backend/tests/live/provider_lifecycle.py | 2 + backend/tests/live/provider_sample.py | 15 +++- 8 files changed, 193 insertions(+), 42 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_provisioning_fence.py diff --git a/backend/airweave/domains/entities/canonical/store.py b/backend/airweave/domains/entities/canonical/store.py index 7b53a34a10..a5abb34057 100644 --- a/backend/airweave/domains/entities/canonical/store.py +++ b/backend/airweave/domains/entities/canonical/store.py @@ -59,6 +59,16 @@ class WriterBusy(CanonicalStoreError): code = "writer_busy" +def _require_current_generation(sync: Sync, job: SyncJob) -> None: + """A prior credential or unverified connection cannot acquire or use a writer.""" + generation = sync.provisioning_generation + if job.provisioning_generation != generation or ( + generation > 0 + and (sync.provisioning_ready_generation != generation or sync.status != "active") + ): + raise StaleWriter("Source connection generation is not admitted for capture") + + def capture_fingerprint(record: CaptureRecord) -> str: """Observation times never create changes; sparse tombstone payloads do.""" material = record.model_dump( @@ -220,6 +230,35 @@ async def _sync( raise SourceNotFound("Source does not exist in this organization") return sync + async def admit_job( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID, job_id: UUID + ) -> int: + """Read one current job/source generation before external source construction. + + This is an admission snapshot, not a replacement for the locked write fence. + Callers compare it again after construction; every write rechecks independently. + """ + row = ( + await db.execute( + select(Sync, SyncJob) + .join(SyncJob, SyncJob.sync_id == Sync.id) + .where( + Sync.id == sync_id, + Sync.organization_id == organization_id, + SyncJob.id == job_id, + SyncJob.organization_id == organization_id, + ) + .execution_options(populate_existing=True) + ) + ).one_or_none() + if row is None: + raise StaleWriter("Source job is unavailable for capture") + sync, job = row + if job.status not in ("pending", "running"): + raise StaleWriter("Source run is no longer active") + _require_current_generation(sync, job) + return sync.provisioning_generation + async def activate_writer( self, db: AsyncSession, @@ -248,6 +287,7 @@ async def activate_writer( ).scalar_one_or_none() if job is None or job.status not in ("pending", "running"): raise StaleWriter("Writer must be an active job of this source") + _require_current_generation(sync, job) identical = ( sync.writer_job_id == job_id and sync.writer_attempt_id == attempt_id @@ -287,11 +327,19 @@ async def _fenced_sync(self, db: AsyncSession, fence: WriterFence) -> Sync: or sync.writer_attempt_number != fence.attempt_number ): raise StaleWriter("Source run has been superseded") - status = await db.scalar( - select(SyncJob.status).where(SyncJob.id == fence.job_id).with_for_update() + job = await db.scalar( + select(SyncJob) + .where( + SyncJob.id == fence.job_id, + SyncJob.sync_id == fence.sync_id, + SyncJob.organization_id == fence.organization_id, + ) + .with_for_update() + .execution_options(populate_existing=True) ) - if status not in ("pending", "running"): + if job is None or job.status not in ("pending", "running"): raise StaleWriter("Source run is no longer active") + _require_current_generation(sync, job) return sync async def capture(self, db: AsyncSession, batch: CaptureBatch) -> CaptureResult: diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 757c9ce474..e4b5abe289 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -278,6 +278,8 @@ async def database(request): await connection.run_sync(seed_flat_scan, request.param) await connection.run_sync(migrate, "0007_scan_scope_owner.py") await connection.run_sync(migrate, "0008_scope_execution.py") + await connection.run_sync(migrate, "0009_extraction_coverage.py") + await connection.run_sync(migrate, "0010_owned_provisioning.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() 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_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index 2ee1adcf76..b5a122f503 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -15,6 +15,7 @@ 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"}]} @@ -23,13 +24,16 @@ REPLIES = {"messages": [MESSAGE, {"ts": "1.1", "thread_ts": "1", "text": "Synthetic reply"}]} -def runner(database, source, responses, attempt=1, service=None): +async def runner(database, source, responses, attempt=1, service=None): 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") + 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, @@ -67,7 +71,7 @@ async def saved(database): async def test_cancel_after_history_resumes_pending_reply_without_repeating_history( database, source ): - first, _, _ = runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()]) + first, _, _ = await runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()]) with pytest.raises(asyncio.CancelledError): await run(first) rows, cursor, scans = await saved(database) @@ -75,7 +79,7 @@ async def test_cancel_after_history_resumes_pending_reply_without_repeating_hist 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, _ = runner(database, source, [ROOT, REPLIES], attempt=2) + second, connector, _ = await runner(database, source, [ROOT, REPLIES], attempt=2) await run(second) assert [call.args[0].split("/")[-1] for call in connector._get.call_args_list] == [ "conversations.list", @@ -95,13 +99,13 @@ async def commit_scan_page(self, db, request): raise ConnectionError("synthetic lost commit acknowledgement") return result - first, connector, _ = runner( + 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, _ = runner(database, source, [ROOT, REPLIES], attempt=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") @@ -113,13 +117,13 @@ async def test_expired_reply_cursor_restarts_whole_scope_and_never_partial_recon "messages": [{"ts": "1.old", "thread_ts": "1"}], "response_metadata": {"next_cursor": "expired"}, } - first, _, _ = runner(database, source, [ROOT, HISTORY, first_reply, asyncio.CancelledError()]) + 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, _ = runner( + second, connector, _ = await runner( database, source, [ROOT, SlackApiError("invalid_cursor"), HISTORY, REPLIES], attempt=2 ) await run(second) @@ -135,7 +139,7 @@ async def test_expired_reply_cursor_restarts_whole_scope_and_never_partial_recon async def test_repeated_invalid_cursor_fails_without_checkpoint_or_absence( database, source, error_code ): - first, _, _ = runner( + first, _, _ = await runner( database, source, [ROOT, HISTORY, SlackApiError(error_code), HISTORY, SlackApiError(error_code)], @@ -153,14 +157,14 @@ async def test_repeated_invalid_cursor_fails_without_checkpoint_or_absence( async def test_membership_loss_on_retry_hides_pending_children_only_after_confirmation( database, source, accessible ): - first, _, _ = runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()]) + 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 = runner(database, source, responses, attempt=2) + second, connector, pipeline = await runner(database, source, responses, attempt=2) if accessible: with pytest.raises(ValueError, match="omitted"): await run(second) @@ -180,14 +184,14 @@ async def test_membership_loss_on_retry_hides_pending_children_only_after_confir async def test_crash_after_scans_then_lost_final_ack_same_job_does_not_start_new_cycle( database, source ): - first, _, _ = runner(database, source, [ROOT, HISTORY, REPLIES]) + 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, _ = runner(database, source, [ROOT], attempt=2) + second, connector, _ = await runner(database, source, [ROOT], attempt=2) await run(second) assert connector._get.await_count == 1 - third, connector, _ = runner(database, source, [], attempt=3) + third, connector, _ = await runner(database, source, [], attempt=3) await run(third) assert connector._get.await_count == 0 _, cursor, _ = await saved(database) @@ -198,7 +202,7 @@ async def test_new_job_starts_fresh_cycle_without_overriding_successful_old_job( from airweave.domains.entities.canonical.store import StaleWriter from airweave.models.sync_job import SyncJob - first, _, _ = runner(database, source, [ROOT, HISTORY, REPLIES]) + first, _, _ = await runner(database, source, [ROOT, HISTORY, REPLIES]) await run(first) _, cursor, _ = await saved(database) prior_cycle = cursor["canonical_cycle"]["version"]["cycle_id"] @@ -216,11 +220,11 @@ async def test_new_job_starts_fresh_cycle_without_overriding_successful_old_job( ) ) await db.commit() - stale, _, _ = runner(database, source, [], attempt=2) + 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, _ = runner(database, new_source, [ROOT, HISTORY, REPLIES]) + new, connector, _ = await runner(database, new_source, [ROOT, HISTORY, REPLIES]) await run(new) assert connector._get.await_count == 3 _, cursor, _ = await saved(database) @@ -231,7 +235,7 @@ async def test_new_job_starts_fresh_cycle_without_overriding_successful_old_job( async def test_access_lost_during_reply_withdraws_root_and_children(database, source): - first, _, _ = runner(database, source, [ROOT, HISTORY, SlackApiError("not_in_channel")]) + 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" @@ -258,7 +262,7 @@ async def test_oversized_page_or_pending_thread_queue_fails_before_commit( for i in range(count) ] } - first, _, _ = runner(database, source, [ROOT, oversized]) + 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) @@ -282,7 +286,7 @@ async def test_interrupted_omission_confirmations_remove_nothing_then_retry(data for value in ("C1", "C2") ), ) - first, _, _ = runner( + first, _, _ = await runner( database, source, [ @@ -296,7 +300,7 @@ async def test_interrupted_omission_confirmations_remove_nothing_then_retry(data rows, cursor, _ = await saved(database) assert all(row.deleted_at is None for row in rows) assert "canonical_checkpoint" not in cursor - second, connector, _ = runner( + second, connector, _ = await runner( database, source, [{"channels": []}, SlackApiError("channel_not_found"), SlackApiError("channel_not_found")], @@ -327,7 +331,7 @@ async def test_cleanup_budget_stop_preserves_reconciling_state_and_resumes(datab for i in range(251) ), ) - first, _, _ = runner(database, source, [ROOT, {"messages": []}]) + first, _, _ = await runner(database, source, [ROOT, {"messages": []}]) checks = 0 async def limit(): @@ -344,7 +348,7 @@ async def limit(): 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, _ = runner(database, source, [ROOT], attempt=2) + second, connector, _ = await runner(database, source, [ROOT], attempt=2) await run(second) assert connector._get.await_count == 1 rows, cursor, _ = await saved(database) @@ -363,11 +367,11 @@ async def test_restored_root_restarts_withdrawn_child_scope_in_same_cycle( HISTORY, REPLIES if completed_child else SlackApiError("not_in_channel"), ] - first, _, _ = runner(database, source, first_responses) + first, _, _ = await runner(database, source, first_responses) await run(first, checkpoint=False) restore_attempt = 2 if completed_child: - removed, _, _ = runner( + removed, _, _ = await runner( database, source, [{"channels": []}, SlackApiError("channel_not_found")], attempt=2 ) await run(removed, checkpoint=False) @@ -375,7 +379,7 @@ async def test_restored_root_restarts_withdrawn_child_scope_in_same_cycle( 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, _ = runner( + restored, connector, _ = await runner( database, source, [ROOT, HISTORY, REPLIES], attempt=restore_attempt ) await run(restored) @@ -390,9 +394,9 @@ async def test_restored_root_restarts_withdrawn_child_scope_in_same_cycle( async def test_root_revival_and_child_invalidation_roll_back_together(database, source): from airweave.domains.entities.canonical.scan_store import CanonicalScanStore - first, _, _ = runner(database, source, [ROOT, HISTORY, REPLIES]) + first, _, _ = await runner(database, source, [ROOT, HISTORY, REPLIES]) await run(first, checkpoint=False) - removed, _, _ = runner( + removed, _, _ = await runner( database, source, [{"channels": []}, SlackApiError("channel_not_found")], attempt=2 ) await run(removed, checkpoint=False) @@ -407,7 +411,7 @@ async def page(self, db, request): service = CanonicalCaptureService(source[0].store) service.scans = FailAfterPage(service.store) - restoring, _, _ = runner(database, source, [ROOT], attempt=3, service=service) + 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) @@ -422,13 +426,13 @@ async def page(self, db, request): async def test_unavailable_parent_without_tombstone_also_invalidates_child_progress( database, source ): - first, _, _ = runner(database, source, [ROOT, HISTORY, SlackApiError("not_in_channel")]) + 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, _ = runner(database, source, [ROOT, HISTORY, REPLIES], attempt=2) + 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) @@ -440,7 +444,7 @@ async def test_scope_withdrawal_keeps_uncertain_reason_without_claiming_access_r ): from airweave.domains.entities.canonical.page_source import ScopeAccessLost - instance, _, pipeline = runner(database, source, [ROOT, HISTORY]) + instance, _, pipeline = await runner(database, source, [ROOT, HISTORY]) original_page = pipeline.page_source.capture_page async def unavailable(scope, continuation, *, files, parent=None): @@ -473,7 +477,7 @@ async def test_old_provider_route_failure_cannot_withdraw_changed_same_fence_par from airweave.domains.entities.canonical.scan_store import ScanConflict service, fence = source - instance, connector, pipeline = runner(database, source, [ROOT]) + instance, connector, pipeline = await runner(database, source, [ROOT]) original_page = connector.capture_page async def change_during_network(scope, continuation, *, files, parent=None): @@ -508,7 +512,7 @@ async def change_during_network(scope, continuation, *, files, parent=None): async def test_missing_queued_thread_restarts_inventory_before_absence(database, source): - instance, connector, _ = runner( + instance, connector, _ = await runner( database, source, [ @@ -546,7 +550,7 @@ async def test_native_slack_dates_reach_persisted_search_fields_and_date_filters from airweave.platform.destinations.vespa.transformer import EntityTransformer from airweave.platform.entities._base import AirweaveSystemMetadata - first, _, _ = runner( + first, _, _ = await runner( database, source, [ @@ -573,6 +577,7 @@ async def test_native_slack_dates_reach_persisted_search_fields_and_date_filters # 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)) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index d5aa8088fa..56b997ebc8 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -796,9 +796,17 @@ def validate_host_ssrf(cls, v: str) -> str: class SlackConfig(SourceConfig): - """Slack configuration schema.""" + """Trusted native workspace/user pair for owned Slack capture.""" - pass + 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): diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index f6dae31521..9f1b89a949 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -318,6 +318,8 @@ async def main(): "0006_record_visibility.py", "0007_scan_scope_owner.py", "0008_scope_execution.py", + "0009_extraction_coverage.py", + "0010_owned_provisioning.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index e2f26d1c8d..b444da0aee 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -895,6 +895,8 @@ async def main(): "0006_record_visibility.py", "0007_scan_scope_owner.py", "0008_scope_execution.py", + "0009_extraction_coverage.py", + "0010_owned_provisioning.py", ): await connection.run_sync(harness.migrate, migration) organization_id, sync_id = uuid4(), uuid4() diff --git a/backend/tests/live/provider_sample.py b/backend/tests/live/provider_sample.py index e89afd24cb..baf0417ab1 100644 --- a/backend/tests/live/provider_sample.py +++ b/backend/tests/live/provider_sample.py @@ -116,15 +116,24 @@ async def rest_source( GoogleCalendarSource, calendar_config or GoogleCalendarConfig(), ), - "slack": (SlackSource, SlackConfig()), + "slack": ( + SlackSource, + SlackConfig( + expected_team_id=os.environ["LIVE_SLACK_TEAM_ID"], + expected_user_id=os.environ["LIVE_SLACK_USER_ID"], + ) + if name == "slack" + else SlackConfig(), + ), }[name] source = await source_type.create( auth=auth, logger=logging.getLogger("probe"), http_client=wrapped, config=config ) - # Google sources attest their explicitly pinned native identity during create. - if name not in {"gmail", "google_calendar", "google_drive"}: + # Sources attest their explicitly pinned native identity during create. + if name not in {"gmail", "google_calendar", "google_drive", "slack"}: await verify_rest_identity(name, source, expected_email) identity_kind = { + "slack": "provider_team_user_id", "google_drive": "provider_permission_id", "google_calendar": "provider_primary_calendar_id", }.get(name, "provider_email") From 39fd26ecf1c88e271ad32db6ba742c7797b545c4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:31:17 -0700 Subject: [PATCH 145/266] Provision owned sources with durable generation and execution recovery --- backend/airweave/api/v1/api.py | 5 + backend/airweave/api/v1/endpoints/admin.py | 4 + .../api/v1/endpoints/owned_provisioning.py | 45 +++ backend/airweave/core/container/container.py | 2 + backend/airweave/core/container/factory.py | 10 + backend/airweave/crud/crud_sync_job.py | 37 ++ .../domains/owned_provisioning/__init__.py | 1 + .../domains/owned_provisioning/guard.py | 23 ++ .../domains/owned_provisioning/models.py | 79 ++++ .../domains/owned_provisioning/service.py | 208 +++++++++++ .../domains/owned_provisioning/store.py | 209 +++++++++++ .../owned_provisioning/tests/conftest.py | 3 + .../tests/test_provisioning.py | 352 ++++++++++++++++++ .../domains/source_connections/create.py | 36 +- .../domains/source_connections/delete.py | 3 + .../domains/source_connections/protocols.py | 11 + .../source_connections/tests/test_delete.py | 31 +- .../source_connections/tests/test_update.py | 74 ++-- .../domains/source_connections/update.py | 3 + backend/airweave/domains/sources/lifecycle.py | 64 ++-- backend/airweave/domains/syncs/protocols.py | 1 + backend/airweave/domains/syncs/service.py | 3 +- .../airweave/domains/temporal/protocols.py | 3 + .../domains/temporal/schedule_service.py | 18 +- backend/airweave/domains/temporal/service.py | 2 + .../tests/test_temporal_schedule_service.py | 99 ++++- backend/airweave/models/owned_provisioning.py | 1 + .../versions/0010_owned_provisioning.py | 1 + docs/owned-sources/provisioning.md | 25 ++ 29 files changed, 1271 insertions(+), 82 deletions(-) create mode 100644 backend/airweave/api/v1/endpoints/owned_provisioning.py create mode 100644 backend/airweave/domains/owned_provisioning/__init__.py create mode 100644 backend/airweave/domains/owned_provisioning/guard.py create mode 100644 backend/airweave/domains/owned_provisioning/models.py create mode 100644 backend/airweave/domains/owned_provisioning/service.py create mode 100644 backend/airweave/domains/owned_provisioning/store.py create mode 100644 backend/airweave/domains/owned_provisioning/tests/conftest.py create mode 100644 backend/airweave/domains/owned_provisioning/tests/test_provisioning.py create mode 100644 docs/owned-sources/provisioning.md diff --git a/backend/airweave/api/v1/api.py b/backend/airweave/api/v1/api.py index 14461ec3d3..0d6f7557b1 100644 --- a/backend/airweave/api/v1/api.py +++ b/backend/airweave/api/v1/api.py @@ -15,6 +15,7 @@ file_retrieval, health, organizations, + owned_provisioning, records, search, search_legacy, @@ -65,3 +66,7 @@ api_router.include_router(calendar_records.router, prefix="/sync", tags=["records"]) + +api_router.include_router( + owned_provisioning.router, prefix="/owned-sources", tags=["owned-sources"] +) diff --git a/backend/airweave/api/v1/endpoints/admin.py b/backend/airweave/api/v1/endpoints/admin.py index 6718fa4f1e..1ac61933c7 100644 --- a/backend/airweave/api/v1/endpoints/admin.py +++ b/backend/airweave/api/v1/endpoints/admin.py @@ -1212,7 +1212,11 @@ async def resync_with_execution_config( ctx.logger.info(f"Admin resync job will be tagged with: {tags}") async with UnitOfWork(db) as uow: + from airweave.crud.crud_sync_job import lock_sync_for_job + + locked_sync = await lock_sync_for_job(uow.session, sync_obj.organization_id, sync_id) sync_job_obj = SyncJob( + provisioning_generation=locked_sync.provisioning_generation, sync_id=sync_id, organization_id=sync_obj.organization_id, status=SyncJobStatus.PENDING, diff --git a/backend/airweave/api/v1/endpoints/owned_provisioning.py b/backend/airweave/api/v1/endpoints/owned_provisioning.py new file mode 100644 index 0000000000..6b52682371 --- /dev/null +++ b/backend/airweave/api/v1/endpoints/owned_provisioning.py @@ -0,0 +1,45 @@ +"""Backend-only account provisioning; org keys retain their existing broad authority.""" + +from uuid import UUID + +from fastapi import Depends, HTTPException +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api import deps +from airweave.api.context import ApiContext +from airweave.api.router import TrailingSlashRouter +from airweave.core.container import Container +from airweave.db.session import get_db +from airweave.domains.owned_provisioning.models import EnsureSource, ProvisionedSource + +router = TrailingSlashRouter() + + +def backend_actor(ctx: ApiContext = Depends(deps.get_context)) -> ApiContext: + """A user session cannot act as Almanac's credential/identity authority.""" + if not ctx.is_api_key_auth: + raise HTTPException(status_code=403, detail="Owned provisioning requires a backend API key") + return ctx + + +@router.put("/{account_id}", response_model=ProvisionedSource) +async def ensure_owned_source( + account_id: UUID, + request: EnsureSource, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> ProvisionedSource: + """Retry the identical generation after uncertain outcomes; IDs remain stable.""" + return await container.owned_provisioning.ensure(db, ctx, account_id, request) + + +@router.get("/{account_id}", response_model=ProvisionedSource) +async def read_owned_source( + account_id: UUID, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> ProvisionedSource: + """Read provisioning state without contacting the provider or starting capture.""" + return await container.owned_provisioning.get(db, ctx, account_id) diff --git a/backend/airweave/core/container/container.py b/backend/airweave/core/container/container.py index a86adde5f4..00ab5fda18 100644 --- a/backend/airweave/core/container/container.py +++ b/backend/airweave/core/container/container.py @@ -74,6 +74,7 @@ OrganizationServiceProtocol, UserOrganizationRepositoryProtocol, ) +from airweave.domains.owned_provisioning.service import OwnedProvisioningService from airweave.domains.search.owned import OwnedSearchService from airweave.domains.search.protocols import ( AgenticSearchServiceProtocol, @@ -253,6 +254,7 @@ async def my_endpoint(event_bus: EventBus = Inject(EventBus)): # Search domain (v2 tiers) owned_search: OwnedSearchService + owned_provisioning: OwnedProvisioningService instant_search: InstantSearchServiceProtocol classic_search: ClassicSearchServiceProtocol agentic_search: AgenticSearchServiceProtocol diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 7485ccfce4..21cf55c676 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -103,6 +103,8 @@ from airweave.domains.organizations.protocols import UserOrganizationRepositoryProtocol from airweave.domains.organizations.repository import OrganizationRepository as OrgRepo from airweave.domains.organizations.repository import UserOrganizationRepository +from airweave.domains.owned_provisioning.service import OwnedProvisioningService +from airweave.domains.owned_provisioning.store import ProvisioningStore 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.agentic.service import AgenticSearchService @@ -615,6 +617,14 @@ def create_container(settings: Settings) -> Container: user_service=user_service, email_service=email_service, owned_search=search_deps["owned_search"], + owned_provisioning=OwnedProvisioningService( + store=ProvisioningStore(create_service, source_validation), + lifecycle=source_deps["source_lifecycle_service"], + jobs=source_deps["sync_job_repo"], + syncs=source_deps["sync_repo"], + schedules=sync_deps["temporal_schedule_service"], + workflows=sync_deps["temporal_workflow_service"], + ), instant_search=search_deps["instant_search"], classic_search=search_deps["classic_search"], agentic_search=search_deps["agentic_search"], diff --git a/backend/airweave/crud/crud_sync_job.py b/backend/airweave/crud/crud_sync_job.py index fa906b6820..89590122bf 100644 --- a/backend/airweave/crud/crud_sync_job.py +++ b/backend/airweave/crud/crud_sync_job.py @@ -4,19 +4,56 @@ from typing import Optional from uuid import UUID +from fastapi import HTTPException from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from airweave.core.context import BaseContext from airweave.crud._base_organization import CRUDBaseOrganization +from airweave.db.unit_of_work import UnitOfWork from airweave.models.sync import Sync from airweave.models.sync_job import SyncJob from airweave.schemas.sync_job import SyncJobCreate, SyncJobUpdate +async def lock_sync_for_job(db: AsyncSession, organization: UUID, sync_id: UUID) -> Sync: + """Serialize admission with provisioning changes, including privileged admin runs.""" + sync = await db.scalar( + select(Sync) + .where(Sync.id == sync_id, Sync.organization_id == organization) + .with_for_update() + .execution_options(populate_existing=True) + ) + if sync is None: + raise HTTPException(status_code=404, detail="Sync not found") + if sync.provisioning_generation and ( + sync.provisioning_generation != sync.provisioning_ready_generation + or sync.status != "active" + ): + raise HTTPException(status_code=409, detail="Owned source is not ready for capture") + return sync + + class CRUDSyncJob(CRUDBaseOrganization[SyncJob, SyncJobCreate, SyncJobUpdate]): """CRUD operations for sync jobs.""" + async def create( + self, + db: AsyncSession, + *, + obj_in: SyncJobCreate, + ctx: BaseContext, + uow: Optional[UnitOfWork] = None, + skip_validation: bool = False, + ) -> SyncJob: + """Stamp job admission under the same lock used by account generation changes.""" + sync = await lock_sync_for_job(db, ctx.organization.id, obj_in.sync_id) + values = obj_in.model_dump(exclude_unset=True) + values["provisioning_generation"] = sync.provisioning_generation + return await super().create( + db, obj_in=values, ctx=ctx, uow=uow, skip_validation=skip_validation + ) + async def get(self, db: AsyncSession, id: UUID, ctx: BaseContext) -> SyncJob | None: """Get a sync job by ID.""" stmt = ( 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..57a63f6f00 --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/guard.py @@ -0,0 +1,23 @@ +"""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 + + +async def require_unmanaged_source(db: AsyncSession, source_id: UUID, organization: UUID) -> None: + """Owned source identity/configuration is writable only through ensure.""" + 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" + ) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py new file mode 100644 index 0000000000..28f3993b48 --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -0,0 +1,79 @@ +"""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, 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"] + expected_identity: str = Field(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) + + @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.""" + 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", "disconnected"] + source: ManagedSource | None = None + + @model_validator(mode="after") + def active_source(self) -> "EnsureSource": + """Disconnect never accepts a replacement identity or credential.""" + if self.state == "active" and self.source is None: + raise ValueError("An active account requires its verified source specification") + if self.state == "disconnected" and self.source is not None: + raise ValueError("Disconnect 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", "disconnected"] + source_connection_id: UUID | None + sync_id: UUID | None + expected_identity: str | None diff --git a/backend/airweave/domains/owned_provisioning/service.py b/backend/airweave/domains/owned_provisioning/service.py new file mode 100644 index 0000000000..89008d717a --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/service.py @@ -0,0 +1,208 @@ +"""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 +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") + spec = EnsureSource.model_validate(row.request_payload) + 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 "disconnected" + ), + source_connection_id=row.source_connection_id, + sync_id=row.sync_id, + expected_identity=spec.source.expected_identity if spec.source else None, + ) + + 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() + 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 == "disconnected": + 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..2235ced77b --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -0,0 +1,209 @@ +"""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 +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.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 == "disconnected" + ): + return + previous = EnsureSource.model_validate(row.request_payload) + if request.state == "active" and row.source_connection_id is not None: + if previous.state == "disconnected": + raise HTTPException( + status_code=409, detail="Connect a new account after disconnect" + ) + if ( + previous.source.provider, + previous.source.expected_identity, + previous.source.collection, + ) != ( + request.source.provider, + request.source.expected_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() + ) + 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..530354a78b --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/tests/conftest.py @@ -0,0 +1,3 @@ +"""Reuse the isolated canonical PostgreSQL migration fixture.""" + +from airweave.domains.entities.canonical.tests.conftest import database # noqa: F401 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..8a2d26b94d --- /dev/null +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -0,0 +1,352 @@ +"""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 +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" diff --git a/backend/airweave/domains/source_connections/create.py b/backend/airweave/domains/source_connections/create.py index 63bc07acae..76a7217ba5 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 @@ -163,6 +164,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, @@ -388,7 +408,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 +456,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 @@ -464,6 +490,7 @@ async def _create_with_auth_provider( run_immediately=bool(obj_in.sync_immediately), ctx=ctx, uow=uow, + **({"defer_execution": True} if transaction else {}), ) await uow.session.flush() @@ -485,11 +512,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( diff --git a/backend/airweave/domains/source_connections/delete.py b/backend/airweave/domains/source_connections/delete.py index 8a25a14169..aed95f15a5 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,8 @@ async def delete( if not source_conn: raise NotFoundException("Source connection not found") + await require_unmanaged_source(db, id, ctx.organization.id) + 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/tests/test_delete.py b/backend/airweave/domains/source_connections/tests/test_delete.py index a70bf44288..17d69140d0 100644 --- a/backend/airweave/domains/source_connections/tests/test_delete.py +++ b/backend/airweave/domains/source_connections/tests/test_delete.py @@ -32,6 +32,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 +47,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 +130,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 +156,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 +187,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 +211,4 @@ 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()) diff --git a/backend/airweave/domains/source_connections/tests/test_update.py b/backend/airweave/domains/source_connections/tests/test_update.py index 581cec880c..0449431731 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,7 +196,7 @@ 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) @@ -212,10 +218,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 +320,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 +343,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 +380,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 +447,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 +473,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 +501,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 +530,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 +611,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 +654,5 @@ 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 diff --git a/backend/airweave/domains/source_connections/update.py b/backend/airweave/domains/source_connections/update.py index 670220f2ab..facced9c4c 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, @@ -93,6 +94,8 @@ async def update( if not source_conn: raise NotFoundException("Source connection not found") + await require_unmanaged_source(db, id, ctx.organization.id) + # Update fields update_data = obj_in.model_dump(exclude_unset=True) diff --git a/backend/airweave/domains/sources/lifecycle.py b/backend/airweave/domains/sources/lifecycle.py index b7b45d37d6..f548d902d0 100644 --- a/backend/airweave/domains/sources/lifecycle.py +++ b/backend/airweave/domains/sources/lifecycle.py @@ -166,38 +166,42 @@ async def create( 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) + + # 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, + ) - return source + # 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, diff --git a/backend/airweave/domains/syncs/protocols.py b/backend/airweave/domains/syncs/protocols.py index 456f3109f4..416c2e8e5e 100644 --- a/backend/airweave/domains/syncs/protocols.py +++ b/backend/airweave/domains/syncs/protocols.py @@ -123,6 +123,7 @@ async def create( run_immediately: bool, ctx: ApiContext, uow: UnitOfWork, + defer_execution: bool = False, ) -> SyncProvisionResult: """Create sync + optional job + Temporal schedule atomically.""" ... diff --git a/backend/airweave/domains/syncs/service.py b/backend/airweave/domains/syncs/service.py index f27f0aee97..009aac357e 100644 --- a/backend/airweave/domains/syncs/service.py +++ b/backend/airweave/domains/syncs/service.py @@ -109,6 +109,7 @@ async def create( run_immediately: bool, ctx: ApiContext, uow: UnitOfWork, + defer_execution: bool = False, ) -> SyncProvisionResult: """Create sync + optional job + Temporal schedule atomically. @@ -143,7 +144,7 @@ async def create( 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, 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/models/owned_provisioning.py b/backend/airweave/models/owned_provisioning.py index 80229008be..3857c1f5ea 100644 --- a/backend/airweave/models/owned_provisioning.py +++ b/backend/airweave/models/owned_provisioning.py @@ -30,6 +30,7 @@ class OwnedProvisioning(OrganizationBase): 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__ = ( diff --git a/backend/alembic/versions/0010_owned_provisioning.py b/backend/alembic/versions/0010_owned_provisioning.py index e296ab62bb..4bd5dba62c 100644 --- a/backend/alembic/versions/0010_owned_provisioning.py +++ b/backend/alembic/versions/0010_owned_provisioning.py @@ -44,6 +44,7 @@ def upgrade(): ), 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", diff --git a/docs/owned-sources/provisioning.md b/docs/owned-sources/provisioning.md new file mode 100644 index 0000000000..f91a137408 --- /dev/null +++ b/docs/owned-sources/provisioning.md @@ -0,0 +1,25 @@ +# Owned account provisioning + +This backend-only API is implemented; it has not been deployed or qualified against live Temporal/provider services. Almanac still needs durable delivery from its account authority before Connect automatically provisions capture. + +`PUT /api/v1/owned-sources/{account_uuid}` commits a desired account generation and reconciles it. `GET` reads the committed delivery state. Both require an organization API key, whose existing broad source-management authority remains unchanged. Never expose this key or endpoint directly to browsers. The server fixes the namespace to `almanac`. + +An active request contains `generation`, `state: "active"`, and `source` with `provider`, `expected_identity`, `collection`, `auth_provider`, `connected_account_id`, `auth_config_id`, `user_id`, `config`, and five-field `cron`. Composio selectors contain no provider tokens. The configured auth-provider connection must belong to the request organization. Native validation checks the expected identity before admitting capture. Supported providers are Gmail (mailbox), Google Calendar (primary calendar ID), and Google Drive (permission ID). Slack and Wispr identity contracts remain future work. + +A disconnect contains only `generation` and `state: "disconnected"`. Active credential rotation preserves source/sync IDs and retained originals. Disconnect is a tombstone for this account UUID, matching Almanac's account lifecycle; a subsequent new Connect uses a new UUID. Deletion of retained records is not part of disconnect. + +Same generation and identical payload reuse the same source, sync and initial job. Changed payload at the same generation or an older generation returns 409. A higher generation cannot change provider, native identity or collection. Ordinary source update/delete routes reject provisioned sources to prevent bypassing this lifecycle. + +The response includes account/organization IDs, current/observed generations, source/sync IDs, expected identity and `pending`, `ready`, or `disconnected`. `ready` means native identity was verified and scheduling/start was acknowledged; it does **not** mean records were captured or indexed. Read source coverage for that evidence. Do not publish an account binding merely because IDs exist in a pending response. + +## Recovery ownership + +Database source creation and desired-generation changes commit atomically before external execution. Provider verification occurs outside the relationship lock and must pass a generation comparison afterward. Jobs copy the current generation under the sync lock; unverified, paused and stale managed generations cannot enter capture. Canonical write fences prevent already-running older jobs from committing after rotation/disconnect. + +The bounded 20-second execution phase serializes cancellation and scheduler changes under the relationship lock. Timeouts release the lock without acknowledging cleanup; cancellation job IDs and the desired generation remain committed for retry. Deterministic schedules are removed even when their database link is missing, then recreated with current generation arguments. The initial workflow uses its durable job UUID and rejects duplicate workflow reuse. + +After an uncertain response, Almanac must retry the identical PUT until acknowledged, or deliver a newer generation. GET does not retry side effects. This API introduces no second queue or autonomous reconciliation worker: the planned Almanac durable account intent owns delivery/retry. Native verification and infrastructure errors may leave a pending intent, and live operational retry/alert behavior remains to be integrated there. + +## Local verification + +The provisioning tests use migrated PostgreSQL and actual source/sync creation, with synthetic identities and mocked provider/Temporal I/O. They cover lost replies, transaction rollback, rotation cleanup timeout, disconnect racing native verification, organization/API-key boundaries, and legacy mutation rejection. A separate actual TemporalScheduleService test uses fake remote handles to prove orphan schedule replacement and relinking. These tests are not live deployment qualification. From 21fb93c18a7c2bd4117d75fadb80b0573cfe3744 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:38:44 -0700 Subject: [PATCH 146/266] Support reversible owned capture pause with preserved native identity --- .../domains/owned_provisioning/models.py | 32 ++++++++-- .../domains/owned_provisioning/service.py | 20 ++++-- .../domains/owned_provisioning/store.py | 26 +++++--- .../tests/test_provisioning.py | 62 +++++++++++++++++++ .../versions/0010_owned_provisioning.py | 3 +- docs/owned-sources/provisioning.md | 21 ++++++- 6 files changed, 143 insertions(+), 21 deletions(-) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 28f3993b48..73f6de1928 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -52,16 +52,16 @@ class EnsureSource(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) generation: int = Field(ge=1, strict=True) - state: Literal["active", "disconnected"] + state: Literal["active", "paused", "disconnected"] source: ManagedSource | None = None @model_validator(mode="after") def active_source(self) -> "EnsureSource": - """Disconnect never accepts a replacement identity or credential.""" + """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 == "disconnected" and self.source is not None: - raise ValueError("Disconnect does not accept new credentials or identity") + if self.state != "active" and self.source is not None: + raise ValueError("Inactive state does not accept new credentials or identity") return self @@ -73,7 +73,29 @@ class ProvisionedSource(BaseModel): organization_id: UUID generation: int observed_generation: int - state: Literal["pending", "ready", "disconnected"] + state: Literal["pending", "ready", "paused", "disconnected"] source_connection_id: UUID | None sync_id: UUID | None expected_identity: str | None + + +def native_identity(provider: str, config: dict) -> str: + """Read the original attested principal from the protected source config.""" + from airweave.platform.configs.config import ( + GmailConfig, + GoogleCalendarConfig, + GoogleDriveConfig, + ) + + 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 _: + raise ValueError("Unsupported owned source provider") + if identity is None: + raise ValueError("Owned source has no trusted native identity") + return identity diff --git a/backend/airweave/domains/owned_provisioning/service.py b/backend/airweave/domains/owned_provisioning/service.py index 89008d717a..edb034f472 100644 --- a/backend/airweave/domains/owned_provisioning/service.py +++ b/backend/airweave/domains/owned_provisioning/service.py @@ -12,7 +12,11 @@ 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 +from airweave.domains.owned_provisioning.models import ( + EnsureSource, + ProvisionedSource, + native_identity, +) 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 @@ -69,7 +73,11 @@ async def get(self, db: AsyncSession, ctx: ApiContext, account: UUID) -> Provisi ) if row is None: raise HTTPException(status_code=404, detail="Owned account source not found") - spec = EnsureSource.model_validate(row.request_payload) + source = ( + await db.get(SourceConnection, row.source_connection_id) + if row.source_connection_id + else None + ) return ProvisionedSource( account_id=row.account_id, organization_id=row.organization_id, @@ -80,11 +88,13 @@ async def get(self, db: AsyncSession, ctx: ApiContext, account: UUID) -> Provisi if row.observed_generation != row.generation else "ready" if row.desired_state == "active" - else "disconnected" + else row.desired_state ), source_connection_id=row.source_connection_id, sync_id=row.sync_id, - expected_identity=spec.source.expected_identity if spec.source else None, + expected_identity=native_identity(source.short_name, source.config_fields) + if source + else None, ) async def reconcile( @@ -155,7 +165,7 @@ async def _execution( status_code=503, detail="Owned capture cancellation is pending" ) row.cancellation_job_ids = [] - if row.desired_state == "disconnected": + 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 diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py index 2235ced77b..b39d584d18 100644 --- a/backend/airweave/domains/owned_provisioning/store.py +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -12,7 +12,7 @@ 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 +from airweave.domains.owned_provisioning.models import EnsureSource, native_identity from airweave.domains.source_connections.protocols import SourceConnectionCreateServiceProtocol from airweave.domains.sources.protocols import SourceValidationServiceProtocol from airweave.models.connection import Connection @@ -92,19 +92,27 @@ async def ensure( 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 == "disconnected" + row.sync_id is not None or row.desired_state != "active" ): return - previous = EnsureSource.model_validate(row.request_payload) + 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: - if previous.state == "disconnected": - raise HTTPException( - status_code=409, detail="Connect a new account after disconnect" + 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 ( - previous.source.provider, - previous.source.expected_identity, - previous.source.collection, + source.short_name, + native_identity(source.short_name, source.config_fields), + source.readable_collection_id, ) != ( request.source.provider, request.source.expected_identity, diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index 8a2d26b94d..ce9f6978d6 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -350,3 +350,65 @@ async def test_creation_failure_rolls_back_relationship_source_and_sync(database 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})) diff --git a/backend/alembic/versions/0010_owned_provisioning.py b/backend/alembic/versions/0010_owned_provisioning.py index 4bd5dba62c..086b489fae 100644 --- a/backend/alembic/versions/0010_owned_provisioning.py +++ b/backend/alembic/versions/0010_owned_provisioning.py @@ -59,7 +59,8 @@ def upgrade(): ), sa.CheckConstraint("client_namespace = 'almanac'", name="ck_owned_provisioning_namespace"), sa.CheckConstraint( - "desired_state IN ('active', 'disconnected')", name="ck_owned_provisioning_state" + "desired_state IN ('active', 'paused', 'disconnected')", + name="ck_owned_provisioning_state", ), ) diff --git a/docs/owned-sources/provisioning.md b/docs/owned-sources/provisioning.md index f91a137408..a575200db1 100644 --- a/docs/owned-sources/provisioning.md +++ b/docs/owned-sources/provisioning.md @@ -2,7 +2,7 @@ This backend-only API is implemented; it has not been deployed or qualified against live Temporal/provider services. Almanac still needs durable delivery from its account authority before Connect automatically provisions capture. -`PUT /api/v1/owned-sources/{account_uuid}` commits a desired account generation and reconciles it. `GET` reads the committed delivery state. Both require an organization API key, whose existing broad source-management authority remains unchanged. Never expose this key or endpoint directly to browsers. The server fixes the namespace to `almanac`. +`PUT /owned-sources/{account_uuid}` commits a desired account generation and reconciles it. `GET` reads the committed delivery state. Both require an organization API key, whose existing broad source-management authority remains unchanged. Never expose this key or endpoint directly to browsers. The server fixes the namespace to `almanac`. An active request contains `generation`, `state: "active"`, and `source` with `provider`, `expected_identity`, `collection`, `auth_provider`, `connected_account_id`, `auth_config_id`, `user_id`, `config`, and five-field `cron`. Composio selectors contain no provider tokens. The configured auth-provider connection must belong to the request organization. Native validation checks the expected identity before admitting capture. Supported providers are Gmail (mailbox), Google Calendar (primary calendar ID), and Google Drive (permission ID). Slack and Wispr identity contracts remain future work. @@ -23,3 +23,22 @@ After an uncertain response, Almanac must retry the identical PUT until acknowle ## Local verification The provisioning tests use migrated PostgreSQL and actual source/sync creation, with synthetic identities and mocked provider/Temporal I/O. They cover lost replies, transaction rollback, rotation cleanup timeout, disconnect racing native verification, organization/API-key boundaries, and legacy mutation rejection. A separate actual TemporalScheduleService test uses fake remote handles to prove orphan schedule replacement and relinking. These tests are not live deployment qualification. + +## Reversible pause + +`state: paused` accepts no source specification, fences existing jobs and removes +schedules using the same durable cleanup as disconnect. It retains the source/sync +IDs and original protected SourceConnection config. A higher-generation active +request revalidates that exact provider/native identity/collection before rotating +credentials and resuming. The stop request does not replace that authority. +Inactive responses return the preserved expected identity when a source exists. +An initially paused account may create its first source later; disconnected is +terminal even if no source was ever created. Pause is for configuration disable; +disconnect is for account removal. + +Pause qualification: seven real PostgreSQL provisioning tests passed (including +active → paused → same-ID active, mismatched native identity rejected, terminal +existing/initial disconnect, and initially paused → first active). Provider native +validation and Temporal calls are simulated; no live account operations occurred. +Log: `/tmp/provisioning-pause-tests.log`. Migration 0010's state constraint was +extended before deployment; deployed schemas would require a new migration. From d5090002f10ee03ce847100147701444ae9e7fe7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:44:04 -0700 Subject: [PATCH 147/266] Require owned provisioning and scheduling recovery checks in CI --- .github/workflows/owned-source-store.yml | 2 ++ Makefile | 8 ++++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 6be1cac93b..16b23be45d 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -74,6 +74,8 @@ jobs: run: make test-store - name: Verify source capture and transport run: make test-capture + - name: Verify connection provisioning and recovery + run: make test-provisioning - name: Verify indexed retrieval and visibility run: make test-search - name: Verify service authentication diff --git a/Makefile b/Makefile index 474600ccd9..0613b8a598 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest IMAGE ?= almanac-source-store:local -.PHONY: help setup check test-store test-capture test-search test-auth test-health test-index validate-deploy build +.PHONY: help setup check test-store test-provisioning test-capture test-search test-auth test-health test-index validate-deploy build help: @echo 'setup Install the locked backend and development dependencies' @@ -37,6 +37,10 @@ test-store: test-capture: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py +test-provisioning: + @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/owned_provisioning/tests airweave/domains/source_connections/tests airweave/domains/temporal/tests airweave/domains/sources/tests/test_lifecycle.py airweave/domains/sync_pipeline/tests/test_factory.py + test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search @@ -50,4 +54,4 @@ test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py tests/integration/test_canonical_search_prefilters.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py -check: test-store test-capture test-search test-auth test-health +check: test-store test-provisioning test-capture test-search test-auth test-health From 25a0561972597e0b7ce773a690e3b9f3460db130 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:55:20 -0700 Subject: [PATCH 148/266] Provision Slack capture with native workspace and user binding --- .../domains/owned_provisioning/models.py | 26 ++++++++++-- .../domains/owned_provisioning/service.py | 10 +++-- .../domains/owned_provisioning/store.py | 6 +-- .../tests/test_provisioning.py | 41 +++++++++++++++++++ docs/owned-sources/provisioning.md | 17 ++++++++ 5 files changed, 90 insertions(+), 10 deletions(-) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 73f6de1928..b8509ae064 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -11,8 +11,9 @@ 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"] + provider: Literal["gmail", "google_calendar", "google_drive", "slack"] 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) @@ -21,6 +22,13 @@ class ManagedSource(BaseModel): 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): + """Slack workspace membership requires its native user, not broker user_id.""" + if (self.provider == "slack") != (self.expected_user_identity is not None): + raise ValueError("Only Slack requires an expected native user identity") + return self + @field_validator("cron") @classmethod def valid_schedule(cls, value: str) -> str: @@ -31,6 +39,12 @@ def valid_schedule(cls, value: str) -> str: def source_config(self) -> dict[str, JsonValue]: """Expected identity cannot be overridden inside unstructured provider config.""" + 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", @@ -77,16 +91,19 @@ class ProvisionedSource(BaseModel): source_connection_id: UUID | None sync_id: UUID | None expected_identity: str | None + expected_user_identity: str | None = None -def native_identity(provider: str, config: dict) -> str: +def native_principal(provider: str, config: dict) -> tuple[str, str | None]: """Read the original attested principal from the protected source config.""" from airweave.platform.configs.config import ( GmailConfig, GoogleCalendarConfig, GoogleDriveConfig, + SlackConfig, ) + user = None match provider: case "gmail": identity = GmailConfig.model_validate(config).expected_mailbox @@ -94,8 +111,11 @@ def native_identity(provider: str, config: dict) -> str: identity = GoogleCalendarConfig.model_validate(config).expected_primary_calendar_id case "google_drive": identity = GoogleDriveConfig.model_validate(config).expected_permission_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 + return identity, user diff --git a/backend/airweave/domains/owned_provisioning/service.py b/backend/airweave/domains/owned_provisioning/service.py index edb034f472..5bf7da6950 100644 --- a/backend/airweave/domains/owned_provisioning/service.py +++ b/backend/airweave/domains/owned_provisioning/service.py @@ -15,7 +15,7 @@ from airweave.domains.owned_provisioning.models import ( EnsureSource, ProvisionedSource, - native_identity, + native_principal, ) from airweave.domains.owned_provisioning.store import ProvisioningStore, locked_intent from airweave.domains.sources.protocols import SourceLifecycleServiceProtocol @@ -78,6 +78,9 @@ async def get(self, db: AsyncSession, ctx: ApiContext, account: UUID) -> Provisi 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, @@ -92,9 +95,8 @@ async def get(self, db: AsyncSession, ctx: ApiContext, account: UUID) -> Provisi ), source_connection_id=row.source_connection_id, sync_id=row.sync_id, - expected_identity=native_identity(source.short_name, source.config_fields) - if source - else None, + expected_identity=identity, + expected_user_identity=user, ) async def reconcile( diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py index b39d584d18..dbbacb0fcb 100644 --- a/backend/airweave/domains/owned_provisioning/store.py +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -12,7 +12,7 @@ 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_identity +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 @@ -111,11 +111,11 @@ async def ensure( ) if ( source.short_name, - native_identity(source.short_name, source.config_fields), + native_principal(source.short_name, source.config_fields), source.readable_collection_id, ) != ( request.source.provider, - request.source.expected_identity, + (request.source.expected_identity, request.source.expected_user_identity), request.source.collection, ): raise HTTPException( diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index ce9f6978d6..2d3b233aac 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -412,3 +412,44 @@ async def test_initial_pause_can_activate_but_initial_disconnect_is_terminal(dat 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") diff --git a/docs/owned-sources/provisioning.md b/docs/owned-sources/provisioning.md index a575200db1..59bab6b403 100644 --- a/docs/owned-sources/provisioning.md +++ b/docs/owned-sources/provisioning.md @@ -42,3 +42,20 @@ existing/initial disconnect, and initially paused → first active). Provider na validation and Temporal calls are simulated; no live account operations occurred. Log: `/tmp/provisioning-pause-tests.log`. Migration 0010's state constraint was extended before deployment; deployed schemas would require a new migration. + +## Slack principal + +Slack requests use `provider: slack`, `expected_identity` for the native workspace +ID and `expected_user_identity` for the native Slack user. The latter is required +only for Slack and is distinct from Composio's `user_id`. Both overwrite conflicting +free-form config fields, feed the existing native `auth.test` capture fence, survive +pause acknowledgements, and must match on resume. Response acknowledgements include +both fields; Google responses have null expected_user_identity. The Almanac adapter +maps `slack_composio` to this provider and verifies both acknowledgement identities. +This admits capture; it does not claim complete Slack incremental edit/delete coverage. + +Verification: eight real-SQL provisioning tests passed (native/Temporal calls mocked), +including wrong-user/same-workspace rejection and pause/resume preservation. Almanac +adapter's 18 tests passed, including rejection of a wrong native-user acknowledgement. +Wispr remains outside automatic provisioning pending native-principal qualification: +the saved catalogue exposes GET_ACCOUNT_INFO but its output schema is unspecified. From 498c4cd8c3476dbc09f4978a01e7b58ff0739c68 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 14:56:56 -0700 Subject: [PATCH 149/266] Record bounded Wispr identity probe and native attestation limitation --- docs/owned-sources/provisioning.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/owned-sources/provisioning.md b/docs/owned-sources/provisioning.md index 59bab6b403..34a13f4e1a 100644 --- a/docs/owned-sources/provisioning.md +++ b/docs/owned-sources/provisioning.md @@ -59,3 +59,22 @@ including wrong-user/same-workspace rejection and pause/resume preservation. Alm adapter's 18 tests passed, including rejection of a wrong native-user acknowledgement. Wispr remains outside automatic provisioning pending native-principal qualification: the saved catalogue exposes GET_ACCOUNT_INFO but its output schema is unspecified. + +### Wispr identity probe (2026-09-30) + +A bounded live read used the existing saved pilot, verified its exact Composio +connected-account ID, ACTIVE state, toolkit and auth-config, then executed only +`WISPR_FLOW_MCP_GET_ACCOUNT_INFO` in an account-pinned session with connection +management disabled. The successful native response contained exactly `name: str` +and `aliases: list[str]`. No values, tokens, names or aliases were logged or stored. +The first session request returned 201 and was incorrectly treated as incomplete +by the disposable diagnostic; a second accepted session completed the single +account-info execution. No inventory read, account grant or provider mutation ran. + +This response supplies display identity, not a stable native principal. It cannot +prove a reconnect preserves the same Wispr account. The automatic native-attested +path remains unsupported for Wispr; the existing explicitly broker-bound pilot +is not upgraded by this probe. A future alternative requires either documented +stable native identity or an explicit product policy for broker-only authority +without silently permitting cross-credential continuity. Discovery/deletion and +raw editor/dictation limits remain unchanged. From 4adfc99465c8800aed6086226dbd5e1422f68fd5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:03:28 -0700 Subject: [PATCH 150/266] Add strict idempotent exact-object storage deletion for projection cleanup --- backend/airweave/adapters/storage/aws_s3.py | 14 ++++ .../airweave/adapters/storage/azure_blob.py | 12 +++ .../airweave/adapters/storage/filesystem.py | 9 +++ backend/airweave/adapters/storage/gcp_gcs.py | 16 ++++ backend/airweave/domains/storage/protocols.py | 8 ++ .../storage/tests/test_exact_delete.py | 79 +++++++++++++++++++ 6 files changed, 138 insertions(+) create mode 100644 backend/airweave/domains/storage/tests/test_exact_delete.py diff --git a/backend/airweave/adapters/storage/aws_s3.py b/backend/airweave/adapters/storage/aws_s3.py index bb2f8371e7..8b4a45cb79 100644 --- a/backend/airweave/adapters/storage/aws_s3.py +++ b/backend/airweave/adapters/storage/aws_s3.py @@ -166,6 +166,20 @@ async def exists(self, path: str) -> bool: except Exception: return False + async def delete_file(self, path: str) -> None: + """S3 DeleteObject is idempotent and never interprets the key as a prefix.""" + from botocore.exceptions import ClientError + + try: + client = await self._get_client() + await client.delete_object(Bucket=self.bucket, Key=self._resolve(path)) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "NoSuchKey": + return + raise StorageException("Exact S3 object deletion failed") from None + except Exception: + raise StorageException("Exact S3 object deletion failed") from None + async def delete(self, path: str) -> bool: """Delete object or all objects under prefix.""" key = self._resolve(path) diff --git a/backend/airweave/adapters/storage/azure_blob.py b/backend/airweave/adapters/storage/azure_blob.py index 5060f810de..fc18a425a3 100644 --- a/backend/airweave/adapters/storage/azure_blob.py +++ b/backend/airweave/adapters/storage/azure_blob.py @@ -170,6 +170,18 @@ async def exists(self, path: str) -> bool: except Exception: return False + async def delete_file(self, path: str) -> None: + """Delete one exact blob; snapshots/leases errors remain explicit failures.""" + from azure.core.exceptions import ResourceNotFoundError + + try: + container = await self._get_container_client() + await container.get_blob_client(self._resolve(path)).delete_blob() + except ResourceNotFoundError: + return + except Exception: + raise StorageException("Exact Azure blob deletion failed") from None + async def delete(self, path: str) -> bool: """Delete blob or all blobs under prefix. diff --git a/backend/airweave/adapters/storage/filesystem.py b/backend/airweave/adapters/storage/filesystem.py index 0f8e0d8490..ce383ccc21 100644 --- a/backend/airweave/adapters/storage/filesystem.py +++ b/backend/airweave/adapters/storage/filesystem.py @@ -121,6 +121,15 @@ async def exists(self, path: str) -> bool: full_path = self._resolve(path) return await aiofiles.os.path.exists(full_path) + async def delete_file(self, path: str) -> None: + """Unlink one exact file; never recursively remove a directory.""" + try: + await aiofiles.os.remove(self._resolve(path)) + except FileNotFoundError: + return + except OSError: + raise StorageException("Exact file deletion failed") from None + async def delete(self, path: str) -> bool: """Delete file or directory from filesystem.""" full_path = self._resolve(path) diff --git a/backend/airweave/adapters/storage/gcp_gcs.py b/backend/airweave/adapters/storage/gcp_gcs.py index a666007e2a..4bdbaf48dd 100644 --- a/backend/airweave/adapters/storage/gcp_gcs.py +++ b/backend/airweave/adapters/storage/gcp_gcs.py @@ -198,6 +198,22 @@ def _exists(): except Exception: return False + async def delete_file(self, path: str) -> None: + """Delete one exact object off-loop; never enumerate a prefix.""" + import asyncio + + from google.cloud.exceptions import NotFound + + def remove(): + self._get_bucket().blob(self._resolve(path)).delete() + + try: + await asyncio.to_thread(remove) + except NotFound: + return + except Exception: + raise StorageException("Exact GCS object deletion failed") from None + async def delete(self, path: str) -> bool: """Delete blob or all blobs under prefix.""" import asyncio diff --git a/backend/airweave/domains/storage/protocols.py b/backend/airweave/domains/storage/protocols.py index 202db64196..c2d1973707 100644 --- a/backend/airweave/domains/storage/protocols.py +++ b/backend/airweave/domains/storage/protocols.py @@ -92,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() From fcc1ec0a35403e36cde1709157a9d73a292b10e1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:08:19 -0700 Subject: [PATCH 151/266] Retain complete derived text with index publication and fenced reads --- backend/airweave/api/v1/endpoints/records.py | 67 ++++ .../entities/canonical/projection_gc.py | 15 +- .../entities/canonical/projection_models.py | 1 + .../entities/canonical/projection_store.py | 8 + .../domains/entities/canonical/projector.py | 34 ++- .../entities/canonical/tests/conftest.py | 1 + .../tests/test_calendar_exclusions.py | 8 +- .../canonical/tests/test_notion_projection.py | 2 +- .../canonical/tests/test_projection.py | 17 +- .../canonical/tests/test_retained_text.py | 289 ++++++++++++++++++ .../entities/canonical/text_artifacts.py | 52 ++++ .../domains/entities/canonical/text_models.py | 75 +++++ .../domains/entities/canonical/text_query.py | 225 ++++++++++++++ .../sync_pipeline/pipeline/text_builder.py | 39 ++- .../sync_pipeline/pipeline/text_models.py | 22 ++ .../sync_pipeline/processors/chunk_embed.py | 30 +- .../sync_pipeline/tests/test_text_builder.py | 13 +- .../cleanup_projection_generations.py | 9 +- .../worker/tests/test_worker_ocr_guard.py | 14 +- .../domains/temporal/worker/wiring.py | 2 +- .../airweave/models/projection_generation.py | 3 + .../versions/0011_retained_projection_text.py | 28 ++ docs/owned-sources/representations.md | 75 +++++ 23 files changed, 996 insertions(+), 33 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_retained_text.py create mode 100644 backend/airweave/domains/entities/canonical/text_artifacts.py create mode 100644 backend/airweave/domains/entities/canonical/text_models.py create mode 100644 backend/airweave/domains/entities/canonical/text_query.py create mode 100644 backend/airweave/domains/sync_pipeline/pipeline/text_models.py create mode 100644 backend/alembic/versions/0011_retained_projection_text.py create mode 100644 docs/owned-sources/representations.md diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index 163b240c77..4587df6b1e 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -25,6 +25,7 @@ RecordPage, ) from airweave.domains.entities.canonical.store import CanonicalStoreError +from airweave.domains.entities.canonical.text_models import TextRead, TextRepresentationList from airweave.domains.search.owned_models import OwnedSearchRequest, OwnedSearchResponse router = TrailingSlashRouter() @@ -50,6 +51,7 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> "stale_record_revision": 409, "blob_unavailable": 503, "document_unavailable": 409, + "text_unavailable": 409, "document_incomplete": 409, "calendar_changed_restart": 409, "calendar_read_incomplete": 409, @@ -172,3 +174,68 @@ async def record_blob( "ETag": '"' + sha256 + '"', }, ) + + +@router.get( + "/{sync_id}/records/{record_id}/text-representations", response_model=TextRepresentationList +) +async def list_text_representations( + sync_id: UUID, + record_id: UUID, + response: Response, + revision: int = Query(ge=1), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), + container: Container = Depends(deps.get_container), +) -> TextRepresentationList: + """Describe derived text only from the current authorized index publication.""" + from airweave.domains.entities.canonical.text_query import CanonicalTextReader + + response.headers["Cache-Control"] = "private, no-store" + response.headers["X-Content-Type-Options"] = "nosniff" + return await CanonicalTextReader(service, container.storage_backend).list( + db, + ctx.organization.id, + sync_id, + record_id, + revision, + ) + + +@router.get( + "/{sync_id}/records/{record_id}/text-representations/{representation_id}", + response_model=TextRead, +) +async def read_text_representation( + sync_id: UUID, + record_id: UUID, + response: Response, + representation_id: UUID, + generation: UUID, + revision: int = Query(ge=1), + offset: int = Query(0, ge=0), + limit: int = Query(12000, ge=1, le=100000), + view: Literal["content", "index"] = "content", + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), + container: Container = Depends(deps.get_container), +) -> TextRead: + """Read complete retained converter text in bounded character ranges.""" + from airweave.domains.entities.canonical.text_query import CanonicalTextReader + + response.headers["Cache-Control"] = "private, no-store" + response.headers["X-Content-Type-Options"] = "nosniff" + return await CanonicalTextReader(service, container.storage_backend).read( + db, + ctx.organization.id, + sync_id, + record_id, + revision, + generation, + representation_id, + offset=offset, + limit=limit, + view=view, + ) diff --git a/backend/airweave/domains/entities/canonical/projection_gc.py b/backend/airweave/domains/entities/canonical/projection_gc.py index c7614c1b98..11edcb2c14 100644 --- a/backend/airweave/domains/entities/canonical/projection_gc.py +++ b/backend/airweave/domains/entities/canonical/projection_gc.py @@ -12,6 +12,7 @@ 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 @@ -85,10 +86,20 @@ async def claim( 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] @@ -126,8 +137,8 @@ async def acknowledge( row.next_gc_at = now + timedelta(minutes=5) return row.gc_error = None - row.delete_cursor += len(page.documents) - if row.delete_cursor >= len(row.documents): + 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 diff --git a/backend/airweave/domains/entities/canonical/projection_models.py b/backend/airweave/domains/entities/canonical/projection_models.py index f2fecac0dc..c23c74a77e 100644 --- a/backend/airweave/domains/entities/canonical/projection_models.py +++ b/backend/airweave/domains/entities/canonical/projection_models.py @@ -78,6 +78,7 @@ class ProjectionCleanupPage(BaseModel): attempt: UUID cursor: int documents: tuple[ProjectionDocument, ...] + artifact_keys: tuple[str, ...] = () def scope_projection_document_id(sync_id: UUID, collection_id: UUID, document_id: str) -> str: diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index 800cc794ce..92c62632d9 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -16,6 +16,8 @@ projection_document_locator, ) from airweave.domains.entities.canonical.store import content_is_available, source_record +from airweave.domains.entities.canonical.text_artifacts import text_manifest +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 @@ -136,6 +138,7 @@ async def prepare( documents: tuple[ProjectionDocument, ...], *, coverage: ExtractionCoverage | None = None, + text_representations: tuple[TextArtifact, ...] | None = None, ) -> bool: """Commit exact immutable deletion identities before the first remote feed.""" identities = set() @@ -161,6 +164,9 @@ async def prepare( p.part_index for p in coverage.parts if p.outcome == "indexed" }: raise ValueError("Projection documents must cover exactly the indexed parts") + text_descriptors = text_manifest( + text_representations, indexed_parts, work.record.sync_id, generation + ) extraction = coverage.persisted() if coverage is not None else None manifest = [document.model_dump() for document in documents] async with UnitOfWork(db): @@ -171,6 +177,7 @@ async def prepare( if ( prior.documents != manifest or prior.extraction_coverage != extraction + or prior.text_representations != text_descriptors or prior.record_id != work.record.id or prior.revision != work.record.revision or prior.pipeline_version != work.pipeline_version @@ -190,6 +197,7 @@ async def prepare( pipeline_version=work.pipeline_version, documents=manifest, extraction_coverage=extraction, + text_representations=text_descriptors, next_gc_at=datetime.now(timezone.utc) + timedelta(hours=1), ) ) diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 19d52e55f8..1845279cee 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -23,6 +23,7 @@ ) from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata +from airweave.domains.entities.canonical.text_artifacts import prepare_text from airweave.domains.storage.protocols import StorageBackend from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor @@ -103,11 +104,16 @@ async def project_one( if not mapped.parts: raise ValueError("Projection mapper returned no required content") selected, coverage = _select_inputs(mapped, work, source_name, generation) - chunks = await self._processor.process( - selected, - ProjectionContext(logger, source_name), - ProjectionRuntime(StrictProjectionTracker()), + context = ProjectionContext(logger, source_name) + runtime = ProjectionRuntime(StrictProjectionTracker()) + built = await self._processor.build_text(selected, context, runtime) + artifacts = prepare_text(built.representations, generation) + chunks = await self._processor.process_built_text( + built.entities, + context, + runtime, strict=True, + expected_ids={entity.entity_id for entity in selected}, ) _stamp_chunks(chunks, work.record) prepared = destination.prepare_documents(chunks) @@ -123,14 +129,30 @@ async def project_one( ) async with self._sessions() as db: if not await self._store.prepare( - db, work, generation, destination.collection_id, manifest, coverage=coverage + db, + work, + generation, + destination.collection_id, + manifest, + coverage=coverage, + text_representations=tuple(item for item, _ in artifacts), ): return False + for artifact, content in artifacts: + await self._storage.write_file( + artifact.storage_key(work.record.sync_id, generation), content + ) await destination.feed_prepared(prepared) if no_documents: async with self._sessions() as db: if not await self._store.prepare( - db, work, generation, destination.collection_id, (), coverage=coverage + db, + work, + generation, + destination.collection_id, + (), + coverage=coverage, + text_representations=(), ): return False async with self._sessions() as db: diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index e4b5abe289..37328d3238 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -280,6 +280,7 @@ async def database(request): await connection.run_sync(migrate, "0008_scope_execution.py") await connection.run_sync(migrate, "0009_extraction_coverage.py") await connection.run_sync(migrate, "0010_owned_provisioning.py") + await connection.run_sync(migrate, "0011_retained_projection_text.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index 15dee1cd85..7ae56e5ce9 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -52,13 +52,13 @@ async def test_cancelled_exception_read_project_reinstate_and_revoke(database, s mapped = mapped.entities assert mapped == () processor = MagicMock() - processor.process = AsyncMock() + processor.build_text = AsyncMock() destination = MagicMock() destination.collection_id = uuid4() destination.feed_prepared = AsyncMock() projector = CanonicalProjector(projections, database, processor, AsyncMock()) assert await projector.project_one(work, "google_calendar", destination, MagicMock()) - processor.process.assert_not_awaited() + processor.build_text.assert_not_awaited() destination.feed_prepared.assert_not_awaited() async with database() as db: assert all( @@ -175,13 +175,13 @@ async def test_operational_occurrences_publish_zero_documents_without_retry(data if w.record.id == row.id ) processor = MagicMock() - processor.process = AsyncMock() + processor.build_text = AsyncMock() destination = MagicMock() destination.collection_id = uuid4() destination.feed_prepared = AsyncMock() projector = CanonicalProjector(projections, database, processor, AsyncMock()) assert await projector.project_one(work, "google_calendar", destination, MagicMock()) - processor.process.assert_not_awaited() + processor.build_text.assert_not_awaited() destination.feed_prepared.assert_not_awaited() async with database() as db: assert all( diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py index 5ea7f278ff..37dea6ddf7 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py @@ -306,7 +306,7 @@ async def test_unknown_notion_block_stays_pending_in_actual_projector(database, fence.organization_id, fence.sync_id, "notion", destination, MagicMock() ) assert result.failed == 1 and result.published == 0 - processor.process.assert_not_called() + processor.build_text.assert_not_called() destination.feed_prepared.assert_not_called() async with database() as db: pending = await store.pending(db, fence.organization_id, fence.sync_id) diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py index 6a94751a20..692d97f715 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py @@ -128,7 +128,7 @@ async def test_owned_payload_mapper_to_publication_pipeline(database, source): work = (await pending(database, fence))[0] processor = MagicMock() - async def process(entities, context, runtime, *, strict): + 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 @@ -139,7 +139,18 @@ async def process(entities, context, runtime, *, strict): chunk.entity_id += "__chunk_0" return [chunk] - processor.process = process + from airweave.domains.sync_pipeline.pipeline.text_models import BuiltText, BuiltTextBatch + + async def build_text(entities, context, runtime): + 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 = uuid4() destination.prepare_documents = lambda chunks: { @@ -166,7 +177,7 @@ async def feed_prepared(documents): assert document.fields["airweave_system_metadata_sync_id"] == str(fence.sync_id) destination.feed_prepared = AsyncMock(side_effect=feed_prepared) - storage = MagicMock() + 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() 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..4442e432bc --- /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 +from uuid import uuid4 + +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 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 + 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=uuid4(), 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/text_artifacts.py b/backend/airweave/domains/entities/canonical/text_artifacts.py new file mode 100644 index 0000000000..3d4622beae --- /dev/null +++ b/backend/airweave/domains/entities/canonical/text_artifacts.py @@ -0,0 +1,52 @@ +"""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, + ) + 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..429afd1bae --- /dev/null +++ b/backend/airweave/domains/entities/canonical/text_models.py @@ -0,0 +1,75 @@ +"""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) + + @model_validator(mode="after") + def boundary(self): + """The converter boundary must lie inside the retained indexing text.""" + 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["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..44f2b156af --- /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="extracted_text" if artifact.content_start is not None else "generated_text", + 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/sync_pipeline/pipeline/text_builder.py b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py index 5e01b89213..9c3440105a 100644 --- a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py +++ b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py @@ -9,6 +9,7 @@ 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 from airweave.platform.entities._base import BaseEntity, CodeFileEntity, FileEntity, WebEntity if TYPE_CHECKING: @@ -62,6 +63,16 @@ 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", + ) -> BuiltTextBatch: + """Retain exact converted text and content boundaries before chunking clears them.""" + content_starts: dict[str, int] = {} source_name = sync_context.source_short_name # Step 1: Build metadata section for all entities @@ -71,13 +82,25 @@ async def build_for_batch( 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), + ) + for entity in entities + ), + ) # ------------------------------------------------------------------------------------ # Metadata Building @@ -316,12 +339,14 @@ async def _convert_partitions( self, converter_groups: Dict[Any, List[Tuple[BaseEntity, str]]], 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 @@ -333,7 +358,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 @@ -343,6 +370,7 @@ async def _convert_sub_batch( converter: Any, sub_batch: List[Tuple[BaseEntity, str]], sync_context: "ProcessingContext", + content_starts: dict[str, int], ) -> List[BaseEntity]: """Convert a sub-batch of entities using the given converter. @@ -350,6 +378,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 @@ -374,7 +403,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 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..45e5cc3b09 --- /dev/null +++ b/backend/airweave/domains/sync_pipeline/pipeline/text_models.py @@ -0,0 +1,22 @@ +"""One conversion result, shared by indexing and optional retained representation reads.""" + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.platform.entities._base import BaseEntity + + +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) + + +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 1d52e46bd6..b2eb7d5e57 100644 --- a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py +++ b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py @@ -19,6 +19,7 @@ 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 from airweave.domains.sync_pipeline.processors.utils import filter_empty_representations from airweave.platform.entities._base import BaseEntity, CodeFileEntity @@ -62,6 +63,33 @@ async def process( # 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", + ) -> BuiltTextBatch: + """Convert once, exposing the complete pre-chunk representation for retention.""" + return await self._text_builder.build_with_text(entities, sync_context, runtime) + + 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" @@ -93,7 +121,7 @@ async def process( 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 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/temporal/activities/cleanup_projection_generations.py b/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py index c6a46941b9..e810c8b789 100644 --- a/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py +++ b/backend/airweave/domains/temporal/activities/cleanup_projection_generations.py @@ -1,17 +1,22 @@ """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.""" @@ -36,10 +41,12 @@ async def run(self) -> None: 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) + remaining -= len(page.documents) + len(page.artifact_keys) finally: await client.close() 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 ad47e2a915..2ecc85cbaf 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 @@ -14,7 +14,10 @@ @pytest.mark.asyncio -async def test_worker_composes_without_ocr_and_starts_supported_activities(monkeypatch, tmp_path): +@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={ @@ -37,8 +40,8 @@ async def test_worker_composes_without_ocr_and_starts_supported_activities(monke def health_client(*args, **kwargs): return original_client( transport=httpx.MockTransport( - lambda request: httpx.Response(200) - if request.url.path == "/health" + lambda request: httpx.Response(health_status) + if request.url.path == "/.well-known/ready" else httpx.Response(503) ), **kwargs, @@ -59,6 +62,11 @@ def health_client(*args, **kwargs): lifecycle = MagicMock(start=AsyncMock(), stop=AsyncMock()) monkeypatch.setattr(worker, "TemporalWorker", lambda _: lifecycle) monkeypatch.setattr(worker.signal, "signal", lambda *_: None) + 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 diff --git a/backend/airweave/domains/temporal/worker/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py index 6fff40c9f9..a51c34d850 100644 --- a/backend/airweave/domains/temporal/worker/wiring.py +++ b/backend/airweave/domains/temporal/worker/wiring.py @@ -61,7 +61,7 @@ def create_activities() -> list: logger.debug("Wiring activities with container dependencies") return [ - CleanupProjectionGenerationsActivity().run, + CleanupProjectionGenerationsActivity(container.storage_backend).run, ProjectCanonicalRecordsActivity( projector=CanonicalProjector( CanonicalProjectionStore(), diff --git a/backend/airweave/models/projection_generation.py b/backend/airweave/models/projection_generation.py index 36dfd7d591..7bb510f27e 100644 --- a/backend/airweave/models/projection_generation.py +++ b/backend/airweave/models/projection_generation.py @@ -24,6 +24,9 @@ class ProjectionGeneration(Base): 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) 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/docs/owned-sources/representations.md b/docs/owned-sources/representations.md new file mode 100644 index 0000000000..99d7b86d94 --- /dev/null +++ b/docs/owned-sources/representations.md @@ -0,0 +1,75 @@ +# Retained text implementation plan + +Goal: read complete derived text from retained originals, with the exact revision, +pipeline and publication generation used by search. Original provider records and +bytes remain authoritative. No additional transcript store, source fetch or CLI +vocabulary is introduced. + +The text builder will record complete indexing text and its converter-content +boundary at construction. Indexing and artifact persistence consume that same +result, without rerunning conversion or reconstructing text from snippets. Generated +metadata remains distinguishable; content reads default to the converter output. +Existing PDF converters flatten page boundaries, so source page anchors remain +explicitly unavailable. Character ranges identify the declared derived text only. + +Generation manifests will retain immutable representation descriptors, not bodies. +The existing object store holds UTF-8 bytes. A manifest commits before external +writes; publication uses the existing compare-and-swap after blob and index writes. +Reads use the same publication/ancestor-visibility predicate before and after +bounded, checksum-verified storage access. Retired-generation garbage collection +owns derived blobs, including abandoned attempts and repeated late writes. + +Scope: full PDF/readable text and declared generated record representations. +No transcription, video understanding, visual fidelity or source-page coordinates +are implied. Existing generations without representation metadata stay explicitly +unavailable until reprojected. + +Proof required: converter runs once; text beyond a search excerpt remains readable; +generated/content boundary survives marker-like user content; prepared but +unpublished bytes cannot be read; revision/pipeline/revocation changes deny reads; +corrupt blobs fail closed; retirement reclaims artifacts and retries failures. +Synthetic provider/PDF and migrated SQL tests are separate from live qualification. + +## Implemented internal read contract + +`GET /sync/{sync_id}/records/{record_id}/text-representations?revision=N` +returns `record_id`, `revision`, `status` (`available` or `unavailable`) and +`representations`. Each descriptor contains an opaque `id`, parent `record_id`, +`revision`, `generation`, `pipeline_version`, source-local `part_key`, +`kind` (`extracted_text` or `generated_text`), `media_type: text/markdown`, +`content_characters` (null for generated-only text), `index_characters`, and +`source_anchors: unavailable`. Storage keys and internal part ordinals are omitted. + +`GET` of that path plus `/{representation_id}` requires `revision` and `generation`. +It accepts `offset` (default 0), `limit` (default 12000, maximum 100000), and +`view=content|index` (default content). Offsets count Unicode code points in that +view. The response contains `representation`, `view`, `offset`, `total_characters`, +`text`, and nullable `next_offset`. Follow next_offset with the same immutable +reference/view. The default excludes generated metadata; generated-only records +require explicit index view. An unavailable/changed generation returns 409, and +withdrawn content returns 404. Neither falls back to a provider or search excerpt. + +These are internal service APIs for Almanac's established read surface, not new +public CLI commands. Existing organization-scoped backend authorization applies. +Almanac must enforce its user/account binding before calling them. + +UTF-8 artifacts use the existing per-file storage size bound; publication fails +instead of silently truncating an oversized conversion. Reads verify byte length, +SHA256 and character count. Migration 0011 adds nullable body-free descriptors; +old generations require explicit reprojection before this representation exists. + +## Verification evidence + +Local synthetic tests exercise an actual three-page PDF converter, migrated +PostgreSQL publication, filesystem blobs and authenticated ASGI read routes. +Chunking/embedding are fixed substitutes: this is a content/authorization proof, +not a search-quality or live-deployment claim. Tests cover unpublished artifacts, +full content versus metadata, corruption, mid-read revision/pipeline changes, +parent withdrawal, failed cleanup and late blob recreation. Strict exact-object +storage deletion propagates failure; cloud version retention remains backend policy. + +The canonical/pipeline/wiring regression run passed 376, skipped 2, with one +30-second child-process timeout; that unchanged lifecycle test passed in 9.18s +on isolated rerun. Final focused retained-text/exclusion/worker checks passed 34. +The worker readiness mock now uses the real `/.well-known/ready` path and tests +both healthy 200 and unavailable 503. No private provider calls or deployment ran. From f36abeb4c34313b68aa4aa93d3bbb222d6d13da0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:16:26 -0700 Subject: [PATCH 152/266] Include retained text schema in live capture proof --- backend/tests/live/canonical_capture.py | 1 + 1 file changed, 1 insertion(+) diff --git a/backend/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index 9f1b89a949..191ab4485c 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -320,6 +320,7 @@ async def main(): "0008_scope_execution.py", "0009_extraction_coverage.py", "0010_owned_provisioning.py", + "0011_retained_projection_text.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: From bb49abe238017354c8e5643041e6644baa25dfe0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:35:51 -0700 Subject: [PATCH 153/266] Add retained local evaluation through owned capture and search APIs --- backend/tests/live/README.md | 66 ++++++ backend/tests/live/evaluation.py | 349 +++++++++++++++++++++++++++++++ 2 files changed, 415 insertions(+) create mode 100644 backend/tests/live/evaluation.py diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index f281c39616..30e4285097 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -401,3 +401,69 @@ Parent independently verified zero live schemas and private-live directories and removed the immutable code archive. The earlier20-request failed batch remains historical evidence; this successful trial had its own20-request cap. Safe results: [evidence](evidence/wispr-recovery-20260930.json), committed `78e7f37`. + +## Retained browser evaluation (explicit opt-in) + +`evaluation.py capture --directory /absolute/new/private/directory` retains a +bounded Gmail + Drive sample for the local UI. It reuses `provider_lifecycle.child` and its production durable-page driver, +the production projector/converters, real local MiniLM/BM25 embeddings, and Vespa. +The default `canonical_capture.py` still always removes its schema/files. This +separate evaluation command requires the same private `sync_tests` Unix socket, +provider environment, and additionally `LIVE_DRIVE_PERMISSION_ID`. Each provider +trial is bounded to 180 requests, 250 records, 180 seconds and 50 MiB of blob writes; +one file is bounded to 10 MiB. Gmail uses `newer_than:1d smaller:100K`. Committed +pages remain searchable if a later page hits a bound; interrupted capture is explicit +in the printed lifecycle evidence and must not be represented as a complete sync. + +Run from `backend` with `PYTHONPATH=.`. The named directory must not already exist; +it is created mode0700 with a mode0600 `manifest.json`. This contains service keys +and source identities and must never be committed or copied into public logs. +Failed evaluations also retain their explicitly requested directory/schema for +inspection and cleanup. API keys expire after one day. Native identity verification +happens before retained binding metadata is created. + +Local dependencies are isolated Vespa on `127.0.0.1:8081/19071`, the repository's +MiniLM inference image on `127.0.0.1:9878`, and Redis on `127.0.0.1:16379`. Deploy +`vespa/app` with dimension384 before capture. Model inference is real; no model API +key or paid model call is necessary. PDF text extraction uses the existing local +extractor; scanned PDFs needing OCR are not claimed covered without an OCR service. + +`evaluation.py serve --directory SAME_DIRECTORY` serves the actual Airweave API +router on `127.0.0.1:18081`, with its production container and persisted encrypted +organization API keys. It does not override authentication. Only DB session ownership +is redirected to the explicit isolated schema. No Temporal schedules or worker are +started. This is a bounded capture/search/read evaluation, **not full synchronization +or automatic Connect qualification**. Almanac must use its own disposable database, +explicit development actor, real source bindings and existing hosted HTTP proxy. +Development-token authentication is not proof of WorkOS login. + +Cleanup after stopping the API: delete each evaluation sync's Vespa documents +through `VespaDestination.delete_by_sync_id`, drop only the manifest's validated +`canonical_eval_<32 hex>` schema from the explicitly private `sync_tests` database, +and remove the exact private evaluation directory and provider-environment file. +The disposable evaluation containers can then be removed by their exact names; +never run global Docker pruning or alter the product database. PostgreSQL WAL and +Docker filesystem blocks may retain deleted bytes until those disposable resources +are destroyed. The private corpus remains on this machine until that cleanup. + + +Verified local evaluation, 2026-09-30 (before subsequent extraction repairs): + +- Real MiniLM 384-dimensional and BM25 inference passed; Vespa schema was deployed. +- Actual API-key middleware denied absent keys (401), invalid keys (403), and + admitted the persisted organization keys (200). No context override was used. +- Gmail: 50 originals from the explicit one-day/100 KiB filter, 47 published text + representations, 6 original blobs verified by HTTP SHA256. Three projections + remain pending: two native MIME size discrepancies and one image-conversion failure. +- Drive: 4 committed originals before a file exceeded the 10 MiB trial limit; + 3 published representations, 5 original blobs verified by HTTP SHA256. Content + includes native Google Doc data and DOCX/PPTX/XLSX; no PDF was present in this sample. +- Exact record reads and keyword search returned actual retained records for both + providers. This does not establish full-account coverage, scheduling recovery, + WorkOS login, cloud deployment, or overall search quality. + +Use `evaluation.py project --directory SAME_DIRECTORY` to process pending stored +records without provider access. It traverses all pending pages once; failures stay +pending and are reported, rather than being silently marked indexed. The service +reads current SQL publication state, so reprojection becomes visible without an API +restart. diff --git a/backend/tests/live/evaluation.py b/backend/tests/live/evaluation.py new file mode 100644 index 0000000000..5d6b883a2c --- /dev/null +++ b/backend/tests/live/evaluation.py @@ -0,0 +1,349 @@ +# ruff: noqa: E402 +"""Opt-in retained local evaluation, using production capture, index and HTTP services.""" + +import argparse +import asyncio +import json +import os +import secrets +from datetime import timedelta +from pathlib import Path +from uuid import UUID, uuid4 + +os.environ.update( + DENSE_EMBEDDER="local_minilm", + EMBEDDING_DIMENSIONS="384", + SPARSE_EMBEDDER="fastembed_bm25", + TEXT2VEC_INFERENCE_URL="http://127.0.0.1:9878", + REDIS_PORT="16379", + DISABLE_RATE_LIMIT="true", +) + +import canonical_capture as capture # noqa: E402 +from pydantic import BaseModel, ConfigDict, Field, SecretStr +from sqlalchemy import func, select, text +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from airweave.core import credentials +from airweave.core.config import AuthMode, settings +from airweave.core.datetime_utils import utc_now_naive +from airweave.core.logging import logger +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.models import Entity, Organization, Sync +from airweave.models.api_key import APIKey +from airweave.models.collection import Collection +from airweave.models.source_connection import SourceConnection +from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata +from airweave.platform.configs.config import GmailConfig, GoogleDriveConfig +from airweave.platform.destinations.vespa.destination import VespaDestination + + +class Source(BaseModel): + model_config = ConfigDict(extra="forbid") + provider: str + account_id: UUID = Field(default_factory=uuid4) + organization_id: UUID + sync_id: UUID + source_connection_id: UUID + external_account_id: str + label: str + api_key: SecretStr + + +class Manifest(BaseModel): + model_config = ConfigDict(extra="forbid") + schema_name: str = Field(pattern=r"^canonical_eval_[a-f0-9]{32}$") + sources: list[Source] = Field(default_factory=list) + encryption_key: SecretStr + state_secret: SecretStr + origin: str = "http://127.0.0.1:18081" + + +def private_root(value: str, *, create: bool = False) -> Path: + root = Path(value).absolute() + if root.is_symlink(): + raise ValueError("Evaluation directory must not be a symlink") + if create: + root.mkdir(mode=0o700, parents=False, exist_ok=False) + if not root.is_dir() or root.stat().st_uid != os.getuid() or root.stat().st_mode & 0o077: + raise ValueError("Evaluation directory must be private and current-user owned") + return root + + +def save(root: Path, manifest: Manifest) -> None: + """Secrets enter only the explicit mode0600 handoff, never stdout.""" + value = manifest.model_dump(mode="json") + value["encryption_key"] = manifest.encryption_key.get_secret_value() + value["state_secret"] = manifest.state_secret.get_secret_value() + for row, source in zip(value["sources"], manifest.sources, strict=True): + row["api_key"] = source.api_key.get_secret_value() + target = root / "manifest.json" + target.write_text(json.dumps(value)) + target.chmod(0o600) + + +def configure(root: Path, manifest: Manifest): + """Production container; no authentication dependency overrides.""" + from airweave.core import container + + settings.AUTH_MODE = AuthMode.API_KEY + settings.ENCRYPTION_KEY = manifest.encryption_key.get_secret_value() + settings.STATE_SECRET = manifest.state_secret.get_secret_value() + settings.DENSE_EMBEDDER = "local_minilm" + settings.EMBEDDING_DIMENSIONS = 384 + settings.SPARSE_EMBEDDER = "fastembed_bm25" + settings.TEXT2VEC_INFERENCE_URL = "http://127.0.0.1:9878" + settings.VESPA_URL = "http://127.0.0.1" + settings.VESPA_PORT = 8081 + settings.STORAGE_BACKEND = "filesystem" + settings.STORAGE_PATH = str(root / "blobs") + settings.REDIS_PORT = 16379 + settings.DISABLE_RATE_LIMIT = True + container.initialize_container(settings) + return container.container + + +async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): + runtime = configure(root, manifest) + projector = CanonicalProjector( + CanonicalProjectionStore(), + sessions, + ChunkEmbedProcessor( + runtime.converter_registry, runtime.dense_embedder, runtime.sparse_embedder + ), + runtime.storage_backend, + ) + for provider, native in ( + ("gmail", os.environ["LIVE_EXPECTED_EMAIL"]), + ("google_drive", os.environ["LIVE_DRIVE_PERMISSION_ID"]), + ): + from provider_lifecycle import child + + async with sessions() as db: + organization = Organization(name="Private bounded UI evaluation") + db.add(organization) + await db.flush() + sync = Sync(organization_id=organization.id, name="Bounded live " + provider) + db.add(sync) + await db.commit() + # Production durable-page lifecycle; bounds interrupt without inventing completion. + outcome = await child( + { + "schema": manifest.schema_name, + "root": str(root), + "provider": provider, + "organization_id": str(organization.id), + "sync_id": str(sync.id), + "request_limit": 180, + "record_limit": 250, + "timeout": 180, + "blob_byte_limit": 50 * 1024 * 1024, + "file_byte_limit": 10 * 1024 * 1024, + "query": "newer_than:1d smaller:100K", + } + ) + async with sessions() as db: + count = await db.scalar( + select(func.count()).select_from(Entity).where(Entity.sync_id == sync.id) + ) + if not count: + raise ValueError("Bounded lifecycle retained no original records") + evidence = {"records": count, "lifecycle_exit": outcome, "bounded_evaluation": True} + deployment = await db.scalar(select(VectorDbDeploymentMetadata)) + if deployment is None: + deployment = VectorDbDeploymentMetadata( + dense_embedder="local_minilm", + embedding_dimensions=384, + sparse_embedder="fastembed_bm25", + ) + db.add(deployment) + await db.flush() + if ( + deployment.dense_embedder, + deployment.embedding_dimensions, + deployment.sparse_embedder, + ) != ("local_minilm", 384, "fastembed_bm25"): + raise ValueError("Evaluation embedding deployment mismatch") + collection = Collection( + name="Bounded private evaluation", + readable_id="evaluation-" + uuid4().hex, + organization_id=sync.organization_id, + vector_db_deployment_metadata_id=deployment.id, + ) + db.add(collection) + await db.flush() + connection = SourceConnection( + name="Bounded " + provider, + short_name=provider, + config_fields=( + GmailConfig(expected_mailbox=native, gmail_query="newer_than:1d smaller:100K") + if provider == "gmail" + else GoogleDriveConfig(expected_permission_id=native) + ).model_dump(mode="json"), + organization_id=sync.organization_id, + readable_collection_id=collection.readable_id, + sync_id=sync.id, + is_authenticated=True, + ) + db.add(connection) + key = secrets.token_urlsafe(32) + db.add( + APIKey( + organization_id=sync.organization_id, + encrypted_key=credentials.encrypt({"key": key}), + expiration_date=utc_now_naive() + timedelta(days=1), + ) + ) + await db.commit() + source = Source( + provider=provider, + organization_id=sync.organization_id, + sync_id=sync.id, + source_connection_id=connection.id, + external_account_id=native, + label=os.environ["LIVE_EXPECTED_EMAIL"], + api_key=SecretStr(key), + ) + manifest.sources.append(source) + save(root, manifest) + destination = await VespaDestination.create( + collection_id=collection.id, organization_id=sync.organization_id, logger=logger + ) + try: + after = None + while True: + batch = await projector.batch( + sync.organization_id, sync.id, provider, destination, logger, after_id=after + ) + print( + json.dumps( + { + "provider": provider, + "capture": evidence, + "projection": batch.model_dump(mode="json"), + } + ), + flush=True, + ) + if not batch.has_more: + break + after = batch.after_id + + finally: + await destination.close_connection() + + +async def project(root: Path, manifest: Manifest, sessions): + runtime = configure(root, manifest) + projector = CanonicalProjector( + CanonicalProjectionStore(), + sessions, + ChunkEmbedProcessor( + runtime.converter_registry, runtime.dense_embedder, runtime.sparse_embedder + ), + runtime.storage_backend, + ) + for source in manifest.sources: + async with sessions() as db: + collection_id = await db.scalar( + select(Collection.id) + .join( + SourceConnection, + SourceConnection.readable_collection_id == Collection.readable_id, + ) + .where(SourceConnection.id == source.source_connection_id) + ) + destination = await VespaDestination.create( + collection_id=collection_id, organization_id=source.organization_id, logger=logger + ) + try: + after = None + while True: + batch = await projector.batch( + source.organization_id, + source.sync_id, + source.provider, + destination, + logger, + after_id=after, + ) + print( + json.dumps( + {"provider": source.provider, "projection": batch.model_dump(mode="json")} + ), + flush=True, + ) + if not batch.has_more: + break + after = batch.after_id + finally: + await destination.close_connection() + + +async def serve(root: Path, manifest: Manifest, sessions): + import uvicorn + from fastapi import FastAPI + + from airweave.api.v1.api import api_router + from airweave.api.v1.endpoints.records import record_error_response + from airweave.db import session + from airweave.domains.entities.canonical.store import CanonicalStoreError + + configure(root, manifest) + session.AsyncSessionLocal = sessions + app = FastAPI() + app.include_router(api_router) + app.add_exception_handler(CanonicalStoreError, record_error_response) + await uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=18081, access_log=False, log_config=None) + ).serve() + + +async def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=("capture", "project", "serve")) + parser.add_argument("--directory", required=True) + args = parser.parse_args() + os.umask(0o077) + root = private_root(args.directory, create=args.command == "capture") + if args.command == "capture": + from cryptography.fernet import Fernet + + manifest = Manifest( + schema_name="canonical_eval_" + uuid4().hex, + encryption_key=SecretStr(Fernet.generate_key().decode()), + state_secret=SecretStr(secrets.token_urlsafe(32)), + ) + save(root, manifest) + else: + manifest = Manifest.model_validate_json((root / "manifest.json").read_text()) + url = capture.test_database_url() + engine = create_async_engine( + url, connect_args={"server_settings": {"search_path": manifest.schema_name}} + ) + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + if args.command == "capture": + async with engine.begin() as db: + await db.execute(text(f'CREATE SCHEMA "{manifest.schema_name}"')) + for path in sorted((Path(__file__).parents[2] / "alembic/versions").glob("*.py")): + if path.name[:4].isdigit(): + await db.run_sync(capture.migrate, path.name) + capture.StoragePaths.TEMP_PROCESSING = str(root / "downloads") + await capture_corpus(root, manifest, sessions, engine) + elif args.command == "project": + await project(root, manifest, sessions) + else: + await serve(root, manifest, sessions) + finally: + await engine.dispose() + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except Exception as error: + print(json.dumps({"failed": True, "error_type": type(error).__name__})) + raise SystemExit(1) from None From f63f36e20430c6369ce4593268f994ebef792381 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:36:31 -0700 Subject: [PATCH 154/266] Project retained Gmail bodies despite native size drift and unavailable OCR --- .../entities/canonical/gmail_projection.py | 5 +- .../domains/entities/canonical/projector.py | 13 +++-- .../tests/test_extraction_coverage.py | 50 +++++++++++++++++++ .../canonical/tests/test_gmail_projection.py | 21 ++++++-- .../sync_pipeline/processors/chunk_embed.py | 5 ++ docs/owned-sources/representations.md | 19 +++++++ 6 files changed, 104 insertions(+), 9 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py index 1c9002fb6b..e437e90f7d 100644 --- a/backend/airweave/domains/entities/canonical/gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -57,8 +57,9 @@ async def _body( content = base64.b64decode( encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True ) - if len(content) != int(body.get("size", 0)): - raise ValueError("Gmail MIME size differs from committed provider metadata") + # Native size is retained metadata, not our integrity authority: live Gmail + # responses can disagree with their decoded data. External bytes have already + # passed canonical size/SHA verification; inline bytes came from retained JSON. return content diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 1845279cee..765e09dfe5 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -25,7 +25,6 @@ from airweave.domains.entities.canonical.search_metadata import stamp_search_metadata from airweave.domains.entities.canonical.text_artifacts import prepare_text from airweave.domains.storage.protocols import StorageBackend -from airweave.domains.sync_pipeline.file_types import SUPPORTED_FILE_EXTENSIONS from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields from airweave.platform.destinations.vespa.destination import VespaDestination @@ -103,7 +102,9 @@ async def project_one( async with map_record(work.record, source_name, self._storage) as mapped: if not mapped.parts: raise ValueError("Projection mapper returned no required content") - selected, coverage = _select_inputs(mapped, work, source_name, generation) + selected, coverage = _select_inputs( + mapped, work, source_name, generation, self._processor.supports_file_extension + ) context = ProjectionContext(logger, source_name) runtime = ProjectionRuntime(StrictProjectionTracker()) built = await self._processor.build_text(selected, context, runtime) @@ -204,7 +205,11 @@ async def batch( def _select_inputs( - mapped: ProjectionInputs, work: ProjectionWork, source_name: str, generation: UUID + mapped: ProjectionInputs, + work: ProjectionWork, + source_name: str, + generation: UUID, + supports_file_extension: Callable[[str], bool], ) -> tuple[list[BaseEntity], ExtractionCoverage]: """Classify only deterministic omissions, preserving stable pre-filter part ordinals.""" selected = [] @@ -216,7 +221,7 @@ def _select_inputs( elif ( part.kind == "file" and part.extension is not None - and part.extension not in SUPPORTED_FILE_EXTENSIONS + and not supports_file_extension(part.extension) ): outcome, reason = "unsupported", "unsupported_format" else: diff --git a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py index a4ab72a964..10c90e3cd6 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py +++ b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py @@ -343,3 +343,53 @@ async def test_coverage_manifest_is_exact_immutable_and_capture_fences_publicati ) is None ) + + +async def test_inline_image_without_ocr_keeps_body_partial_but_converter_failure_is_fatal( + database, source, tmp_path +): + service, fence = source + item = original(part(b"retained image bytes", mime="image/png", filename="inline.png")) + item.payload["payload"]["mimeType"] = "multipart/related" + await capture(database, service, fence, item) + storage = FilesystemBackend(tmp_path) + target = destination() + result = await projector(database, storage).batch( + fence.organization_id, fence.sync_id, "gmail", target, 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 [p.outcome for p in coverage.parts] == ["indexed", "unsupported"] + assert coverage.parts[1].media_type == "image/png" + assert row.indexed_chunk_count > 0 + assert row.source_payload == item.payload + # Enabling OCR changes the pipeline; an actual converter failure must not + # masquerade as an unsupported-format success or preserve stale coverage. + await db.execute( + update(Sync).where(Sync.id == fence.sync_id).values(index_pipeline_version=2) + ) + await db.commit() + ocr = MagicMock(convert_batch=AsyncMock(side_effect=lambda paths: dict.fromkeys(paths))) + configured = CanonicalProjector( + CanonicalProjectionStore(), + database, + ChunkEmbedProcessor(ConverterRegistry(ocr), FakeDenseEmbedder(), FakeSparseEmbedder()), + storage, + ) + target.feed_prepared.reset_mock() + failed = await configured.batch(fence.organization_id, fence.sync_id, "gmail", target, logger) + assert failed.failed == 1 and failed.published == 0 + ocr.convert_batch.assert_awaited_once() + target.feed_prepared.assert_not_awaited() + async with database() as db: + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + is None + ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py index 62792da390..c5ede3d4aa 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -107,7 +107,7 @@ async def test_external_blob_and_duplicate_filename_attachments_get_distinct_ide { "mimeType": "multipart/mixed", "parts": [ - {"mimeType": "text/html", "body": {"attachmentId": "body", "size": 11}}, + {"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"), ], @@ -115,8 +115,8 @@ async def test_external_blob_and_duplicate_filename_attachments_get_distinct_ide ) data = b"body" source = source.model_copy(update={"blobs": (blob(source, data, "/payload/parts/0/body"),)}) - # Provider's declared size must match actual content. - source.payload["payload"]["parts"][0]["body"]["size"] = len(data) + # 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) @@ -219,3 +219,18 @@ async def test_related_html_branch_does_not_get_replaced_by_plain_alternative(tm 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/sync_pipeline/processors/chunk_embed.py b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py index b2eb7d5e57..024f88f5aa 100644 --- a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py +++ b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py @@ -40,10 +40,15 @@ 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], diff --git a/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index 99d7b86d94..73dc3d1a0a 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -73,3 +73,22 @@ The canonical/pipeline/wiring regression run passed 376, skipped 2, with one on isolated rerun. Final focused retained-text/exclusion/worker checks passed 34. The worker readiness mock now uses the real `/.well-known/ready` path and tests both healthy 200 and unavailable 503. No private provider calls or deployment ran. + +## Gmail size metadata and configured extraction coverage + +| Native capability/resource | Product name | Access operation | Stored representation | Sync/change guarantee | Proof | Gap | +| --- | --- | --- | --- | --- | --- | --- | +| Inline MIME body / external body bytes | Email text | Retained record projection | Native JSON unchanged; external blob verified against canonical actual byte size and SHA256 | Projection does not refetch or modify capture | Synthetic declared-size discrepancy succeeds; invalid base64 and corrupt canonical bytes still fail | Live retained inline bodies showed native size one byte smaller than decoded data; no upstream cause established | +| Inline images and attachments | Email extraction parts | Configured converter registry | Available body text plus explicit unsupported part coverage when converter absent | Reprojection follows existing pipeline-version lifecycle | Real SQL partial publication without OCR; configured converter failure still blocks publication | Adding OCR requires pipeline-version change to retry prior unsupported parts; no image understanding claim | + +Google documents MIME `size` as byte count, but it is provider metadata rather +than our storage-integrity checksum. Projection preserves the reported value +and consumes the retained base64 data or verified canonical blob without an +additional provider-size equality requirement. External attachment acquisition +keeps its existing checks; this change does not claim that path was requalified. +[Native MessagePartBody contract](https://developers.google.com/workspace/gmail/api/reference/rest/v1/users.messages.attachments#MessagePartBody). + +Unsupported means no configured converter for that part. It does not turn an +installed converter's failure into success: malformed PDFs, OCR errors, and +other attempted-conversion failures still leave publication pending. Scanned +PDFs with a configured PDF converter but no OCR fallback remain in that category. From 0f9427eac78aa076cc515fb7f15212f974f220b4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:41:29 -0700 Subject: [PATCH 155/266] Report uncaptured Drive originals as unavailable extraction --- .../entities/canonical/projection_mappers.py | 23 +++++++ .../tests/test_extraction_coverage.py | 64 +++++++++++++++++++ docs/owned-sources/representations.md | 13 ++++ 3 files changed, 100 insertions(+) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index f8e2df5804..6034103105 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -248,6 +248,29 @@ async def map_record( if excluded_from_search(record, source_name): yield ProjectionInputs(parts=()) return + if ( + source_name == "google_drive" + and record.identity.record_type == "file" + and record.completeness == "metadata_only" + and not record.blobs + and record.payload.get("mimeType") != "application/vnd.google-apps.folder" + ): + # Capture explicitly omitted bytes. Publish the omission, not invented + # title-only document content; complete records still require their blob. + yield ProjectionInputs( + parts=( + ProjectionInput( + part=ExtractionPart( + part_index=0, + key=record.identity.native_id, + kind="file", + media_type=_string(record.payload.get("mimeType")) or None, + ), + entity=None, + ), + ) + ) + return with TemporaryDirectory(prefix="airweave-projection-") as temporary: directory = Path(temporary) if source_name == "gmail": diff --git a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py index 10c90e3cd6..88804cec63 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py +++ b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py @@ -393,3 +393,67 @@ async def test_inline_image_without_ocr_keeps_body_partial_but_converter_failure ) is None ) + + +async def test_drive_metadata_only_reports_missing_original_then_new_bytes_are_indexed( + database, source, tmp_path +): + from airweave.domains.entities.canonical.tests.test_gmail_projection import blob, record + + service, fence = source + identity = RecordIdentity(record_type="file", native_id="retained-file") + payload = {"id": identity.native_id, "name": "Notes.txt", "mimeType": "text/plain"} + metadata = observation(identity=identity, payload=payload, completeness="metadata_only") + await capture(database, service, fence, metadata) + storage = FilesystemBackend(tmp_path) + target = destination() + project = projector(database, storage) + assert ( + await project.batch(fence.organization_id, fence.sync_id, "google_drive", target, logger) + ).published == 1 + 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 == "unavailable" and row.indexed_chunk_count == 0 + assert coverage.parts[0].outcome == "unavailable_original" + assert coverage.parts[0].reason == "original_not_captured" + content = b"Recovered original text contains the complete available document." + ref = blob(record({}, sync_id=fence.sync_id), content).model_copy( + update={"media_type": "text/plain"} + ) + await storage.write_file(ref.key, content) + await capture( + database, service, fence, observation(identity=identity, payload=payload, blobs=(ref,)) + ) + assert ( + await project.batch(fence.organization_id, fence.sync_id, "google_drive", target, logger) + ).published == 1 + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id)) + assert row.indexed_chunk_count > 0 + previous = ProjectionLocator( + record_id=row.id, + revision=row.record_revision, + pipeline_version=row.indexed_pipeline_version, + generation=row.indexed_generation, + part_index=0, + ) + # Losing retained body availability withdraws the prior text generation. + await capture(database, service, fence, metadata) + assert ( + await project.batch(fence.organization_id, fence.sync_id, "google_drive", target, logger) + ).published == 1 + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.sync_id == fence.sync_id)) + assert row.indexed_chunk_count == 0 + assert ( + await db.scalar(select(Entity.id).join(Sync).where(publication_matches(previous))) + is None + ) + assert ( + await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + ).status == "unavailable" diff --git a/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index 73dc3d1a0a..c98a4b425b 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -92,3 +92,16 @@ Unsupported means no configured converter for that part. It does not turn an installed converter's failure into success: malformed PDFs, OCR errors, and other attempted-conversion failures still leave publication pending. Scanned PDFs with a configured PDF converter but no OCR fallback remain in that category. + +### Drive originals omitted during capture + +A non-folder Drive file explicitly captured as `metadata_only` with no blobs +publishes `unavailable_original` / `original_not_captured` extraction evidence, +zero searchable chunks, and unavailable text. It does not index a title as if it +were the document. A complete record missing its required blob remains an error. +The retained evaluation presentation failure was this omission, not a converter +failure: native size exceeded the initial capture limit and no bytes were stored. + +| Native capability/resource | Product name | Access operation | Stored representation | Sync/change guarantee | Proof | Gap | +| --- | --- | --- | --- | --- | --- | --- | +| File metadata without captured body | Original / extraction availability | Existing record and extraction reads | Native metadata plus explicit unavailable-original coverage | Later captured bytes form a new revision; loss of body retires previous text eligibility | Real SQL metadata-only → retained text → metadata-only lifecycle; complete/missing-blob remains error | Actual recapture depends on provider access and configured limits; no converter success claimed | From 48af3ec8bee594838e9011a869478492d5481a8c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:50:28 -0700 Subject: [PATCH 156/266] Extend retained trials with native Calendar and Slack identities --- .../tests/test_drive_lifecycle_probe.py | 38 ++++++++ backend/tests/live/README.md | 28 +++++- backend/tests/live/evaluation.py | 92 +++++++++++++++---- backend/tests/live/provider_lifecycle.py | 10 +- 4 files changed, 146 insertions(+), 22 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py index 75d14429e6..5af3ba79d5 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_lifecycle_probe.py @@ -88,3 +88,41 @@ async def test_drive_lifecycle_initial_then_fresh_process_changes(database, sour assert first["resumed_changes_requests"] == 0 and second["resumed_changes_requests"] == 1 assert first["observed_change_sequence"] == second["observed_change_sequence"] assert first["payload_revision_digest"] == second["payload_revision_digest"] + + +async def test_checkpoint_validation_distinguishes_full_resume_from_changes(): + code, output, error = await subprocess_script( + r""" +import sys +from types import SimpleNamespace +sys.path.insert(0, 'tests/live') +import provider_lifecycle as lifecycle +attempt = SimpleNamespace(id="attempt") +for prior_phase, prior_mode, final_mode, changes_requests in ( + ("active", "full", "full", 0), + ("complete", "full", "changes", 1), + ("active", "changes", "changes", 1), +): + previous = {"canonical_cycle": {"phase": prior_phase, "mode": prior_mode}} + saved = { + "canonical_cycle": {"phase": "complete", "mode": final_mode, + "completed_job_id": "job", "last_full_capture": {"cycle_id": "full"}, + "promoted_checkpoint": {"checkpoint": {"value": {"page_token": "next"}}}}, + "canonical_checkpoint": {"writer_attempt_id": "attempt", "observed_change_sequence": 4}, + } + counters = {"started": 0, "completed": 0, "resumed_changes_requests": changes_requests} + lifecycle.validate_checkpoint("google_drive", {}, counters, saved, previous, True, attempt, 4) +# Loading a prior cursor must not excuse a missing promotion after completion. +saved["canonical_cycle"]["promoted_checkpoint"] = None +try: + lifecycle.validate_checkpoint("google_drive", {}, counters, saved, previous, True, attempt, 4) +except AssertionError: + pass +else: + raise AssertionError("Missing checkpoint promotion was accepted") +print("checkpoint modes verified") +""", + {}, + ) + assert code == 0, output + error + assert "checkpoint modes verified" in output diff --git a/backend/tests/live/README.md b/backend/tests/live/README.md index 30e4285097..fa881b8374 100644 --- a/backend/tests/live/README.md +++ b/backend/tests/live/README.md @@ -410,7 +410,7 @@ the production projector/converters, real local MiniLM/BM25 embeddings, and Vesp The default `canonical_capture.py` still always removes its schema/files. This separate evaluation command requires the same private `sync_tests` Unix socket, provider environment, and additionally `LIVE_DRIVE_PERMISSION_ID`. Each provider -trial is bounded to 180 requests, 250 records, 180 seconds and 50 MiB of blob writes; +trial is bounded to 180 requests, 500 records, 180 seconds and 50 MiB of blob writes; one file is bounded to 10 MiB. Gmail uses `newer_than:1d smaller:100K`. Committed pages remain searchable if a later page hits a bound; interrupted capture is explicit in the printed lifecycle evidence and must not be represented as a complete sync. @@ -467,3 +467,29 @@ records without provider access. It traverses all pending pages once; failures s pending and are reported, rather than being silently marked indexed. The service reads current SQL publication state, so reprojection becomes visible without an API restart. + + +`evaluation.py extend --directory SAME_DIRECTORY --provider google_calendar|slack` +adds one native-attested source while preserving existing keys and source identities. +Extensions run sequentially because the private manifest has one owner. Calendar +uses its verified primary ID and a seven-day occurrence materialization window; +original masters are also retained by the production connector. Slack uses the +verified workspace/user pair, with `external_user_id` retained for the Almanac binding. +Source construction starts at the current canonical search metadata pipeline version; +older retained trials upgrade using `plan_reprojection` followed by `project`. + +The Calendar trial retained 183 visible originals (190 observations including removals) +and projected every record. The Slack trial stopped at its record budget after 56 +requests: 460 committed originals were all projected, with capture still explicitly +active/incomplete. Wispr's separate private broker-attested trial stopped on an upstream +tool error with a rate signal after 18 requests/284 observations; no native account +identity or Almanac binding is claimed for it. + +Drive's second bounded trial resumed the same unfinished full-scan cursor with a +200 MiB file/512 MiB total blob limit after checking at least 10 GiB free disk. It +retained 14 PDFs among the captured originals. The source job completed and SQL showed +a complete full cycle with a promoted checkpoint. The former live-helper assertion +incorrectly assumed any loaded cursor implies a changes pass; it now distinguishes +unfinished full-scan continuation from subsequent/resumed changes while still requiring +checkpoint promotion. This was a harness assertion failure, separately recorded from +the actual completed source job. PDF text/read verification remains a separate check. diff --git a/backend/tests/live/evaluation.py b/backend/tests/live/evaluation.py index 5d6b883a2c..9e62ffa863 100644 --- a/backend/tests/live/evaluation.py +++ b/backend/tests/live/evaluation.py @@ -6,8 +6,9 @@ import json import os import secrets -from datetime import timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path +from tempfile import NamedTemporaryFile from uuid import UUID, uuid4 os.environ.update( @@ -30,13 +31,19 @@ from airweave.core.logging import logger from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore from airweave.domains.entities.canonical.projector import CanonicalProjector +from airweave.domains.entities.canonical.search_metadata import SEARCH_METADATA_PIPELINE_VERSION from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor from airweave.models import Entity, Organization, Sync from airweave.models.api_key import APIKey from airweave.models.collection import Collection from airweave.models.source_connection import SourceConnection from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata -from airweave.platform.configs.config import GmailConfig, GoogleDriveConfig +from airweave.platform.configs.config import ( + GmailConfig, + GoogleCalendarConfig, + GoogleDriveConfig, + SlackConfig, +) from airweave.platform.destinations.vespa.destination import VespaDestination @@ -48,6 +55,7 @@ class Source(BaseModel): sync_id: UUID source_connection_id: UUID external_account_id: str + external_user_id: str | None = None label: str api_key: SecretStr @@ -79,9 +87,15 @@ def save(root: Path, manifest: Manifest) -> None: value["state_secret"] = manifest.state_secret.get_secret_value() for row, source in zip(value["sources"], manifest.sources, strict=True): row["api_key"] = source.api_key.get_secret_value() - target = root / "manifest.json" - target.write_text(json.dumps(value)) - target.chmod(0o600) + with NamedTemporaryFile(mode="w", dir=root, prefix=".manifest-", delete=False) as staged: + temporary = Path(staged.name) + try: + json.dump(value, staged) + staged.flush() + os.fsync(staged.fileno()) + temporary.replace(root / "manifest.json") + finally: + temporary.unlink(missing_ok=True) def configure(root: Path, manifest: Manifest): @@ -105,7 +119,33 @@ def configure(root: Path, manifest: Manifest): return container.container -async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): +def capture_configuration(provider): + if provider == "gmail": + native = os.environ["LIVE_EXPECTED_EMAIL"] + return ( + GmailConfig(expected_mailbox=native, gmail_query="newer_than:1d smaller:100K"), + native, + None, + ) + if provider == "google_drive": + native = os.environ["LIVE_DRIVE_PERMISSION_ID"] + return GoogleDriveConfig(expected_permission_id=native), native, None + if provider == "slack": + native, user = os.environ["LIVE_SLACK_TEAM_ID"], os.environ["LIVE_SLACK_USER_ID"] + return SlackConfig(expected_team_id=native, expected_user_id=user), native, user + if provider == "google_calendar": + native = os.environ["LIVE_CALENDAR_PRIMARY_ID"] + now = datetime.now(timezone.utc).replace(microsecond=0) + config = GoogleCalendarConfig( + expected_primary_calendar_id=native, + calendar_ids=[native], + occurrence_window={"start": now, "end": now + timedelta(days=7)}, + ) + return config, native, None + raise ValueError("Provider requires a native identity before evaluation binding") + + +async def capture_corpus(root: Path, manifest: Manifest, sessions, providers): runtime = configure(root, manifest) projector = CanonicalProjector( CanonicalProjectionStore(), @@ -115,17 +155,21 @@ async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): ), runtime.storage_backend, ) - for provider, native in ( - ("gmail", os.environ["LIVE_EXPECTED_EMAIL"]), - ("google_drive", os.environ["LIVE_DRIVE_PERMISSION_ID"]), - ): + for provider in providers: + if any(source.provider == provider for source in manifest.sources): + raise ValueError("Source already captured; resume its existing cursor instead") + config, native, native_user = capture_configuration(provider) from provider_lifecycle import child async with sessions() as db: organization = Organization(name="Private bounded UI evaluation") db.add(organization) await db.flush() - sync = Sync(organization_id=organization.id, name="Bounded live " + provider) + sync = Sync( + organization_id=organization.id, + name="Bounded live " + provider, + index_pipeline_version=SEARCH_METADATA_PIPELINE_VERSION, + ) db.add(sync) await db.commit() # Production durable-page lifecycle; bounds interrupt without inventing completion. @@ -137,11 +181,16 @@ async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): "organization_id": str(organization.id), "sync_id": str(sync.id), "request_limit": 180, - "record_limit": 250, + "record_limit": 500, "timeout": 180, "blob_byte_limit": 50 * 1024 * 1024, "file_byte_limit": 10 * 1024 * 1024, "query": "newer_than:1d smaller:100K", + **( + {"calendar_config": config.model_dump(mode="json")} + if provider == "google_calendar" + else {} + ), } ) async with sessions() as db: @@ -177,11 +226,7 @@ async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): connection = SourceConnection( name="Bounded " + provider, short_name=provider, - config_fields=( - GmailConfig(expected_mailbox=native, gmail_query="newer_than:1d smaller:100K") - if provider == "gmail" - else GoogleDriveConfig(expected_permission_id=native) - ).model_dump(mode="json"), + config_fields=config.model_dump(mode="json"), organization_id=sync.organization_id, readable_collection_id=collection.readable_id, sync_id=sync.id, @@ -203,6 +248,7 @@ async def capture_corpus(root: Path, manifest: Manifest, sessions, engine): sync_id=sync.id, source_connection_id=connection.id, external_account_id=native, + external_user_id=native_user, label=os.environ["LIVE_EXPECTED_EMAIL"], api_key=SecretStr(key), ) @@ -303,9 +349,12 @@ async def serve(root: Path, manifest: Manifest, sessions): async def main(): parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("command", choices=("capture", "project", "serve")) + parser.add_argument("command", choices=("capture", "extend", "project", "serve")) parser.add_argument("--directory", required=True) + parser.add_argument("--provider", choices=("google_calendar", "slack")) args = parser.parse_args() + if (args.command == "extend") != (args.provider is not None): + parser.error("--provider is required only for extend") os.umask(0o077) root = private_root(args.directory, create=args.command == "capture") if args.command == "capture": @@ -332,7 +381,12 @@ async def main(): if path.name[:4].isdigit(): await db.run_sync(capture.migrate, path.name) capture.StoragePaths.TEMP_PROCESSING = str(root / "downloads") - await capture_corpus(root, manifest, sessions, engine) + await capture_corpus(root, manifest, sessions, ("gmail", "google_drive")) + elif args.command == "extend": + if args.provider is None: + raise ValueError("Extend requires an explicit provider") + capture.StoragePaths.TEMP_PROCESSING = str(root / "downloads") + await capture_corpus(root, manifest, sessions, (args.provider,)) elif args.command == "project": await project(root, manifest, sessions) else: diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index b444da0aee..b588f1e9ec 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -247,9 +247,15 @@ def validate_checkpoint(name, manifest, counters, saved, previous, loaded, attem assert cycle["phase"] == "complete" and cycle["completed_job_id"] assert cycle["last_full_capture"] and cycle["promoted_checkpoint"] assert cycle["promoted_checkpoint"]["checkpoint"]["value"]["page_token"] - assert cycle["mode"] == ("changes" if loaded else "full") + previous_cycle = previous.get("canonical_cycle", {}) + changes_expected = loaded and ( + previous_cycle.get("phase") == "complete" or previous_cycle.get("mode") == "changes" + ) + # A durable cursor can also be an unfinished full enumeration. Loading it + # resumes full capture; it is not evidence that a changes pass began. + assert cycle["mode"] == ("changes" if changes_expected else "full") assert counters["started"] == counters["completed"] == 0 - assert bool(counters["resumed_changes_requests"]) == loaded + assert bool(counters["resumed_changes_requests"]) == changes_expected elif name in {"slack", "wispr"}: assert saved["canonical_cycle"]["phase"] == "complete" assert saved["canonical_cycle"]["completed_job_id"] From bf4aa3bd10b7383c42ce293df2b7b7a5957ae9dc Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 15:58:50 -0700 Subject: [PATCH 157/266] Batch exact publication checks across search candidates --- .../entities/canonical/projection_store.py | 50 +++++++++++---- .../canonical/tests/test_owned_search.py | 50 +++++++++++++++ .../domains/search/canonical_visibility.py | 61 ++++++++++++------- backend/airweave/domains/search/owned.py | 19 +----- 4 files changed, 131 insertions(+), 49 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index 92c62632d9..85d20fa9b7 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -1,9 +1,11 @@ """Canonical pending work and atomic publication. No network I/O under row locks.""" +from collections import defaultdict +from collections.abc import Iterable from datetime import datetime, timedelta, timezone from uuid import UUID -from sqlalchemy import and_, exists, or_, select +from sqlalchemy import and_, exists, or_, select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork @@ -25,23 +27,49 @@ def publication_matches(locator: ProjectionLocator): """Add request organization/source scope separately; gate before text leaves retrieval.""" + return publications_match((locator,)) + + +def publications_match(locators: Iterable[ProjectionLocator]): + """Batch publication checks; callers must still validate each returned chunk's part. + + One record can match one requested part while another part is unavailable. + Organization/source scope is supplied by the caller. + """ + identities = set() + parts: dict[int, set[UUID]] = defaultdict(set) + for locator in locators: + identities.add( + (locator.record_id, locator.revision, locator.pipeline_version, locator.generation) + ) + parts[locator.part_index].add(locator.generation) return and_( - Entity.id == locator.record_id, - Entity.record_revision == locator.revision, - Entity.indexed_revision == locator.revision, - Entity.indexed_pipeline_version == locator.pipeline_version, - Sync.index_pipeline_version == locator.pipeline_version, - Entity.indexed_generation == locator.generation, + tuple_( + Entity.id, + Entity.record_revision, + Entity.indexed_pipeline_version, + Entity.indexed_generation, + ).in_(list(identities)), + Entity.indexed_revision == Entity.record_revision, + Entity.indexed_pipeline_version == Sync.index_pipeline_version, Entity.deleted_at.is_(None), content_is_available(), exists( - select(ProjectionGeneration.id).where( - ProjectionGeneration.id == locator.generation, + select(ProjectionGeneration.id) + .correlate(Entity) + .where( + ProjectionGeneration.id == Entity.indexed_generation, ProjectionGeneration.retired_at.is_(None), or_( ProjectionGeneration.extraction_coverage.is_(None), - ProjectionGeneration.extraction_coverage.contains( - {"parts": [{"part_index": locator.part_index, "outcome": "indexed"}]} + *( + and_( + ProjectionGeneration.id.in_(list(generations)), + ProjectionGeneration.extraction_coverage.contains( + {"parts": [{"part_index": part, "outcome": "indexed"}]} + ), + ) + for part, generations in parts.items() ), ), ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 32fa48c63f..ab091d3350 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -554,3 +554,53 @@ async def test_partial_extraction_http_hit_original_and_nonindexed_part_gate( await db.commit() response = await client.get(f"/sync/{fence.sync_id}/records/{locator.record_id}") assert response.json()["extraction"] is None + + +async def test_visibility_batches_exact_parts_generations_and_duplicate_chunks(database, indexed): + from airweave.domains.search.canonical_visibility import visible_results + from airweave.models.projection_generation import ProjectionGeneration + + fence, locator, connection = indexed + registry = SimpleNamespace( + get=lambda _: SimpleNamespace( + source_class_ref=SimpleNamespace(canonical_record_types=("event",)) + ) + ) + async with database() as db: + await db.execute( + update(ProjectionGeneration) + .where(ProjectionGeneration.id == locator.generation) + .values( + extraction_coverage={ + "parts": [ + {"part_index": 0, "key": "body", "kind": "body", "outcome": "indexed"}, + { + "part_index": 1, + "key": "file", + "kind": "file", + "outcome": "unsupported", + "reason": "unsupported_format", + }, + ], + } + ) + ) + await db.commit() + valid = hit(fence, locator.encode()) + inputs = [ + valid, + hit(fence, locator.model_copy(update={"part_index": 1}).encode()), + hit(fence, locator.model_copy(update={"generation": uuid4()}).encode()), + hit(fence, locator.model_copy(update={"pipeline_version": 1}).encode()), + valid, + valid.model_copy( + update={ + "airweave_system_metadata": valid.airweave_system_metadata.model_copy( + update={"sync_id": str(uuid4())} + ) + } + ), + ] + assert await visible_results( + db, fence.organization_id, connection.readable_collection_id, inputs, registry + ) == [valid, valid] diff --git a/backend/airweave/domains/search/canonical_visibility.py b/backend/airweave/domains/search/canonical_visibility.py index 5a02d5c00f..9670d79526 100644 --- a/backend/airweave/domains/search/canonical_visibility.py +++ b/backend/airweave/domains/search/canonical_visibility.py @@ -2,14 +2,16 @@ from uuid import UUID -from sqlalchemy import and_, or_, select +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.projection_store import publications_match 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 @@ -58,13 +60,10 @@ async def visible_results( 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: set[tuple[UUID, UUID, int, int, UUID]] = set() - candidate_values = list(candidates.values()) + 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): - predicates = [ - and_(Entity.sync_id == sync_id, publication_matches(locator)) - for sync_id, locator in candidate_values[start : start + 100] - ] + batch = candidate_values[start : start + 100] rows = await db.execute( select( Entity.sync_id, @@ -72,28 +71,46 @@ async def visible_results( 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, - or_(*predicates), + Entity.sync_id.in_({sync_id for sync_id, _ in batch}), + publications_match(locator for _, locator in batch), ) ) - eligible.update(tuple(row) for row in rows) + 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 ( - index in candidates - and ( - candidates[index][0], - candidates[index][1].record_id, - candidates[index][1].revision, - candidates[index][1].pipeline_version, - candidates[index][1].generation, - ) - in eligible - ) + 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/owned.py b/backend/airweave/domains/search/owned.py index 502e6a02a2..d2de517797 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -12,7 +12,7 @@ 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 publication_matches +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, @@ -243,7 +243,7 @@ async def _enrich( Entity.organization_id == ctx.organization.id, Sync.organization_id == ctx.organization.id, Entity.sync_id.in_(sync_ids), - or_(*(publication_matches(item) for item in locators)), + publications_match(locators), ) .execution_options(populate_existing=True) ) @@ -336,10 +336,6 @@ async def _final_publications( # SQL read boundary. No network I/O occurs after this authorization check. if not scores: return set() - identities = [ - (locator.record_id, locator.revision, locator.pipeline_version, locator.generation) - for _, locator in scores.values() - ] # One statement snapshots all candidates, including current authentication. # The 20*200 bound keeps these tuple parameters below PostgreSQL's limit. return set( @@ -364,16 +360,7 @@ async def _final_publications( SourceConnection.short_name, SourceConnection.readable_collection_id, ).in_(list(scope_snapshot)), - Entity.deleted_at.is_(None), - content_is_available(), - Entity.indexed_revision == Entity.record_revision, - Entity.indexed_pipeline_version == Sync.index_pipeline_version, - tuple_( - Entity.id, - Entity.record_revision, - Entity.indexed_pipeline_version, - Entity.indexed_generation, - ).in_(identities), + publications_match(locator for _, locator in scores.values()), ) ) ) From 147c472f2e908fe9b1005327b01ef162d7cb0f27 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 16:02:19 -0700 Subject: [PATCH 158/266] Assert Wispr recovery against the actual failed scope --- .../domains/entities/canonical/tests/test_wispr_recovery.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py index 2717365cb6..0598513304 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py @@ -213,15 +213,17 @@ async def execute(slug, arguments): 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 == [ - ("n2", 0), - ("n2", 3), + (failed_native, 0), + (failed_native, 3), ] # Bodies commit atomically, ranges do not checkpoint separately. async with database() as db: rows = list( From 0758d17b9dc9922fe69cf72ac6ddb03fa9e9b665 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 16:24:06 -0700 Subject: [PATCH 159/266] Reuse query embeddings within owned multi-collection searches --- .../canonical/tests/test_owned_search.py | 3 + backend/airweave/domains/search/executor.py | 43 ++++++++--- .../airweave/domains/search/fakes/executor.py | 15 +++- backend/airweave/domains/search/owned.py | 15 ++-- backend/airweave/domains/search/protocols.py | 6 ++ .../domains/search/tests/test_executor.py | 75 +++++++++++++++++++ .../domains/search/types/embeddings.py | 25 ++++++- docs/owned-sources/search-execution.md | 52 +++++++++++++ 8 files changed, 214 insertions(+), 20 deletions(-) create mode 100644 docs/owned-sources/search-execution.md diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index ab091d3350..cb5e3cbe43 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -342,10 +342,12 @@ async def test_final_gate_drops_early_hit_changed_during_later_collection( await db.commit() execute = executor.execute calls = 0 + prepared_queries = [] async def later_edit(**kwargs): nonlocal calls calls += 1 + prepared_queries.append(kwargs["prepared_query"]) if calls == 2: await capture(database, capture_service, fence, observation(payload={"updated": True})) return SearchResults(results=[]) @@ -373,6 +375,7 @@ async def ordered(*args): "mode": "keyword", }, ) + assert prepared_queries[0] is not None and prepared_queries[0] is prepared_queries[1] assert calls == 2 and response.status_code == 200 assert response.json()["items"] == [] and response.json()["excluded_candidates"] == 1 assert "Private original text" not in response.text diff --git a/backend/airweave/domains/search/executor.py b/backend/airweave/domains/search/executor.py index 2fd3b0416f..37b1ddd248 100644 --- a/backend/airweave/domains/search/executor.py +++ b/backend/airweave/domains/search/executor.py @@ -29,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, @@ -90,6 +91,7 @@ async def execute( 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 @@ -122,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 @@ -186,17 +190,8 @@ async def execute( excluded_candidates=excluded, ) - async def _execute_vector_search( - self, - plan: SearchPlan, - collection_id: str, - acl_principals: Optional[list[str]] = 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. - """ + 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 @@ -218,6 +213,30 @@ 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, diff --git a/backend/airweave/domains/search/fakes/executor.py b/backend/airweave/domains/search/fakes/executor.py index d3ad4588a5..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, @@ -42,6 +54,7 @@ async def execute( 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 index d2de517797..044bc415f3 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -68,14 +68,17 @@ async def search( scope_snapshot = self._scope_identity(scopes) hits, scores, 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(): results = await self._executor.execute( - plan=SearchPlan( - query=SearchQuery(primary=request.query), - limit=200, - offset=0, - retrieval_strategy=request.mode, - ), + plan=plan, + prepared_query=prepared_query, user_filter=[FilterGroup(conditions=self._prefilters(request, sync_ids))], collection_id=str(collection_id), db=db, diff --git a/backend/airweave/domains/search/protocols.py b/backend/airweave/domains/search/protocols.py index ad83ae8f63..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, @@ -43,6 +48,7 @@ async def execute( 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/tests/test_executor.py b/backend/airweave/domains/search/tests/test_executor.py index c7f72a8855..b58486a4e6 100644 --- a/backend/airweave/domains/search/tests/test_executor.py +++ b/backend/airweave/domains/search/tests/test_executor.py @@ -1215,3 +1215,78 @@ async def test_executor_returns_only_publication_gate_survivors(monkeypatch): 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 + 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/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/docs/owned-sources/search-execution.md b/docs/owned-sources/search-execution.md new file mode 100644 index 0000000000..4b22a2c622 --- /dev/null +++ b/docs/owned-sources/search-execution.md @@ -0,0 +1,52 @@ +# Owned search execution and performance evidence + +Owned search retrieves indexed originals. It never falls back to provider search. +Each request resolves authorized source connections, uses the unchanged 200-candidate +window per collection, checks exact publications before content leaves retrieval, +and rechecks source authentication and publications before returning results. + +Publication eligibility includes record/revision, pipeline version, generation, +non-retirement, parent visibility and indexed-part coverage. Batched tuple checks +share the visibility expression; an indexed body never authorizes an unsupported +attachment. Retained text reads use the same publication predicate. + +Query embeddings can be prepared once and reused across collections **within one +owned-search request**. Preparation belongs to the creating executor instance and +pins the primary query, ordered variations and retrieval strategy. A different +executor or changed query/strategy is rejected. Filters, ACL resolution, vector +queries and publication checks still run separately for each collection. Keyword +search never performs dense inference. There is no shared result, authorization, +query or embedding cache and no concurrent use of a database session. + +## Topology and limits + +Almanac's server-owned `CaptureTarget` chooses organization and collection for each +owner/provider/project. Provisioning does not require an organization or collection +per provider. Sources may share a collection, or use several collections in the +same organization. Organization is an authorization boundary; collection is an +index/deployment boundary. A provisioned source cannot silently change collection. + +The September 30 local pilot uses four isolated organizations, each with one +collection. Request-local embedding reuse does **not** reduce that pilot's four +embedding calls. The multi-collection SQL fixture qualifies reuse and preserves the +race fence; it is not evidence of deployed production topology or a new speedup +for the pilot. No corpus ownership or operator binding was remapped for this work. + +## Recorded local evidence + +Before/after batched publication checks, an actual-router Slack hybrid query for +`meeting` retained 200 candidates and fell from 2222 ms to 569 ms in two local +instrumented samples. Enrichment fell from 625 to 25 ms. Embedding and Vespa latency +varied between samples, so these are diagnostic samples, not a latency guarantee. + +A later full browser benchmark had one HTTP 503 whose body and duration were not +retained. The evaluation harness suppresses general logs to avoid private-content +logging, so the failing stage cannot be reconstructed. Both containers remained +running without OOM/restart. No cause-specific retry or timeout increase is justified +by that evidence. Safe status, exception type and elapsed-stage diagnostics are +being qualified separately in the Almanac adapter. + +The current request-owned database session also holds its connection across +embedding/index waits after the first SQL query. Separating read phases is a +follow-up requiring explicit auth/session ownership and retained final fences; +it is not included in embedding reuse and is not a proven cause of the 503. From 0f98717a00a82941a6c46dc65ffc657c86a200f5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 16:30:02 -0700 Subject: [PATCH 160/266] Record repeated search stage timings and follow-up boundaries --- docs/owned-sources/search-execution.md | 57 + docs/owned-sources/search-stage-samples.json | 2117 ++++++++++++++++++ 2 files changed, 2174 insertions(+) create mode 100644 docs/owned-sources/search-stage-samples.json diff --git a/docs/owned-sources/search-execution.md b/docs/owned-sources/search-execution.md index 4b22a2c622..32eb585679 100644 --- a/docs/owned-sources/search-execution.md +++ b/docs/owned-sources/search-execution.md @@ -50,3 +50,60 @@ The current request-owned database session also holds its connection across embedding/index waits after the first SQL query. Separating read phases is a follow-up requiring explicit auth/session ownership and retained final fences; it is not included in embedding reuse and is not a proven cause of the 503. + + +## Repeated actual API qualification after embedding reuse + +Six serial product requests each returned HTTP 200 and 25 results. Their wall times +were 1976, 836, 1408, 1202, 1210 and 3013 ms. Each fanned out to four organizations; +[24 sanitized stage samples](search-stage-samples.json) preserve measured evidence. +The slowest organization request per product request was: + +| Sample | Airweave total | Dense | Vespa adapter | Visibility | Enrichment | All SQL | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1176 | 410 | 452 | 35 | 230 | 273 | +| 2 | 703 | 163 | 274 | 54 | 176 | 212 | +| 3 | 1334 | 74 | 189 | 883 | 170 | 1031 | +| 4 | 1094 | 176 | 416 | 122 | 331 | 381 | +| 5 | 1019 | 175 | 378 | 70 | 346 | 389 | +| 6 | 2815 | 131 | 2609 | 23 | 19 | 36 | + +All durations are milliseconds, measured as nested wall intervals; SQL overlaps +other stages and these columns must not be summed. No errors were recorded. +The sixth sample's delay was in the Vespa adapter await, which includes thread +queueing, HTTP, native engine work and event-loop scheduling. It is **not** evidence +of a PostgreSQL plan switch. Sample 3 independently had an 831 ms SQL fingerprint +interval. Neither measurement reconstructs the previous unlogged 503/13-second tail. + +## Proposed next diagnostics and session ownership (not implemented) + +The current Vespa adapter invokes the synchronous SDK through `asyncio.to_thread`. +Measure worker start/end alongside the adapter await, and retain only numeric native +response timing fields if available. This separates thread queue delay from SDK wall +time without logging YQL, embeddings, results or identities. Do not change query +ranking, candidate counts, retries or database plans before attributing the tail. + +Database phase separation needs an explicit owner. The search route and its auth +dependency currently share FastAPI's cached `get_db` session. API context contains +Pydantic snapshots; Auth0 resolution nevertheless flushes a last-active update into +that session without committing it. Closing an arbitrary borrowed session inside +the executor would be an unsafe general contract. + +A bounded redesign would finish this route's authentication session explicitly, +then use existing session-factory read contexts: scopes/version to a typed detached +snapshot; embedding preparation without SQL; each collection's network retrieval +before its session performs visibility/enrichment SQL; finally fresh scopes, +coverage and exact publication checks in one read phase. Generic executor behavior +and other routes stay unchanged. Test auth side-effect semantics and revocation +between phases, including a constrained pool, before claiming reduced occupancy. +No extra service, table, result cache or concurrent `AsyncSession` use is proposed. + +The installed PyVespa `Vespa.query` creates a fresh `VespaSync` HTTP session per +call, then retains the full decoded response in `VespaQueryResponse.json`. This is +a source-level finding, not proof that connection setup causes the observed tail. +Vespa's [result-format documentation](https://docs.vespa.ai/en/reference/querying/default-result-format.html) +provides `querytime`, `summaryfetchtime` and `searchtime` in seconds when +`presentation.timing=true` is requested. A temporary diagnostic can copy the query +body and add only that presentation flag, retaining those three numeric values +alongside worker-entry/exit times. No query trace, response body or result content +needs to be logged. Pooling changes should wait for that decomposition. diff --git a/docs/owned-sources/search-stage-samples.json b/docs/owned-sources/search-stage-samples.json new file mode 100644 index 0000000000..ce0417144c --- /dev/null +++ b/docs/owned-sources/search-stage-samples.json @@ -0,0 +1,2117 @@ +{ + "qualification": "September30 local actual API with request-scoped stage instrumentation; 24organizationrequests from6serialproductqueries; durations are nested wall time not additive CPU; no queries, source IDs, content or SQLparameters recorded. Candidatewindow200 retained.", + "requests": [ + { + "sequence": 3, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 20.981332985684276 + }, + "dense": { + "count": 1, + "ms": 416.2748340750113 + }, + "sparse": { + "count": 1, + "ms": 6.506791920401156 + }, + "prepare_query": { + "count": 1, + "ms": 422.8015000699088 + }, + "vespa": { + "count": 1, + "ms": 382.25970801431686 + }, + "visibility": { + "count": 1, + "ms": 140.25079098064452 + }, + "execute": { + "count": 1, + "ms": 524.2698750225827 + }, + "_enrich": { + "count": 1, + "ms": 91.03287500329316 + }, + "_coverage": { + "count": 1, + "ms": 9.96904110070318 + }, + "_final_publications": { + "count": 1, + "ms": 11.461625108495355 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 17.354625975713134 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.983167021535337 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 70.8857080899179 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 14.206208055838943 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 61.02733395528048 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 22.471540956757963 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 5.890500033274293 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.7412910927087069 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 5.720125045627356 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 1081.774792051874 + }, + { + "sequence": 2, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 27.059999061748385 + }, + "dense": { + "count": 1, + "ms": 423.87624992989004 + }, + "sparse": { + "count": 1, + "ms": 1.0307920165359974 + }, + "prepare_query": { + "count": 1, + "ms": 424.9693750170991 + }, + "vespa": { + "count": 1, + "ms": 559.3767920508981 + }, + "visibility": { + "count": 1, + "ms": 40.0237909052521 + }, + "execute": { + "count": 1, + "ms": 601.858958019875 + }, + "_enrich": { + "count": 1, + "ms": 58.30166698433459 + }, + "_coverage": { + "count": 1, + "ms": 26.148000033572316 + }, + "_final_publications": { + "count": 1, + "ms": 11.744874995201826 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 23.133749957196414 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 7.161208079196513 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 20.701125031337142 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 42.683916981332004 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 1.3428329257294536 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 18.984375055879354 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.0869590332731605 + }, + "2499200562431a3a": { + "count": 1, + "ms": 8.137083961628377 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 1150.5849999375641 + }, + { + "sequence": 1, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 24.168250034563243 + }, + "dense": { + "count": 1, + "ms": 415.1032909285277 + }, + "sparse": { + "count": 1, + "ms": 3.470625029876828 + }, + "prepare_query": { + "count": 1, + "ms": 419.2308330675587 + }, + "vespa": { + "count": 1, + "ms": 455.41983388829976 + }, + "visibility": { + "count": 1, + "ms": 65.13550004456192 + }, + "execute": { + "count": 1, + "ms": 526.5297499718145 + }, + "_enrich": { + "count": 1, + "ms": 73.77662509679794 + }, + "_coverage": { + "count": 1, + "ms": 104.42025004886091 + }, + "_final_publications": { + "count": 1, + "ms": 14.168875059112906 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 17.26112491451204 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 4.835000028833747 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 28.40045909397304 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 14.269707957282662 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 58.46558301709592 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 7.127500022761524 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 18.680915934965014 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.4107919996604323 + }, + "4860037db3700175": { + "count": 1, + "ms": 41.594209033064544 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 1.4456250937655568 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 1.257791998796165 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 9.199041058309376 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 7.366082980297506 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 1164.2424170859158 + }, + { + "sequence": 4, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 24.722332949750125 + }, + "dense": { + "count": 1, + "ms": 410.0976249901578 + }, + "sparse": { + "count": 1, + "ms": 6.485457997769117 + }, + "prepare_query": { + "count": 1, + "ms": 416.62533301860094 + }, + "vespa": { + "count": 1, + "ms": 452.3796249413863 + }, + "visibility": { + "count": 1, + "ms": 35.135333077050745 + }, + "execute": { + "count": 1, + "ms": 490.4897080268711 + }, + "_enrich": { + "count": 1, + "ms": 230.23374995682389 + }, + "_coverage": { + "count": 1, + "ms": 5.5731250904500484 + }, + "_final_publications": { + "count": 1, + "ms": 7.313958019949496 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 22.142291883938015 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 18.409666954539716 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 8.498082985170186 + }, + "227345276c4b5446": { + "count": 1, + "ms": 217.60816604364663 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.8341670036315918 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 1.3499170308932662 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.5095419473946095 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 3.139832988381386 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 1175.6089580012485 + }, + { + "sequence": 7, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 9.415167034603655 + }, + "dense": { + "count": 1, + "ms": 150.34812502563 + }, + "sparse": { + "count": 1, + "ms": 7.051291991956532 + }, + "prepare_query": { + "count": 1, + "ms": 158.213583054021 + }, + "vespa": { + "count": 1, + "ms": 205.83049999549985 + }, + "visibility": { + "count": 1, + "ms": 89.61008302867413 + }, + "execute": { + "count": 1, + "ms": 301.0749579407275 + }, + "_enrich": { + "count": 1, + "ms": 104.27737492136657 + }, + "_coverage": { + "count": 1, + "ms": 11.745374999009073 + }, + "_final_publications": { + "count": 1, + "ms": 36.33583406917751 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 7.525542052462697 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.370500074699521 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 15.718917013145983 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 54.76729106158018 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 91.03541704826057 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 1.473416923545301 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 4.532459075562656 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.223042025230825 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 26.070374995470047 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 623.4907500911504 + }, + { + "sequence": 8, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 11.766041163355112 + }, + "dense": { + "count": 1, + "ms": 153.2676670467481 + }, + "sparse": { + "count": 1, + "ms": 6.886459072120488 + }, + "prepare_query": { + "count": 1, + "ms": 160.18770798109472 + }, + "vespa": { + "count": 1, + "ms": 231.21958295814693 + }, + "visibility": { + "count": 1, + "ms": 81.94658404681832 + }, + "execute": { + "count": 1, + "ms": 317.49387504532933 + }, + "_enrich": { + "count": 1, + "ms": 71.11679099034518 + }, + "_coverage": { + "count": 1, + "ms": 47.35345905646682 + }, + "_final_publications": { + "count": 1, + "ms": 19.881334039382637 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 9.379792027175426 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 4.599625011906028 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 23.095999960787594 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 22.25341706071049 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 40.53116694558412 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 9.81604098342359 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 11.780292028561234 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.5035420181229711 + }, + "4860037db3700175": { + "count": 1, + "ms": 0.6775000365450978 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 2.9274170519784093 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 2.296374994330108 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 10.991000104695559 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 9.165249997749925 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 629.3531249975786 + }, + { + "sequence": 6, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 5.98758296109736 + }, + "dense": { + "count": 1, + "ms": 152.33945904765278 + }, + "sparse": { + "count": 1, + "ms": 11.915208073332906 + }, + "prepare_query": { + "count": 1, + "ms": 164.3058330519125 + }, + "vespa": { + "count": 1, + "ms": 358.3093340275809 + }, + "visibility": { + "count": 1, + "ms": 12.147084111347795 + }, + "execute": { + "count": 1, + "ms": 374.7342500137165 + }, + "_enrich": { + "count": 1, + "ms": 99.2774999467656 + }, + "_coverage": { + "count": 1, + "ms": 10.457874974235892 + }, + "_final_publications": { + "count": 1, + "ms": 8.07679200079292 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.4282091073691845 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.4807909717783332 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 5.371040897443891 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 56.48758402094245 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 29.85670801717788 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 5.245625041425228 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.6650830619037151 + }, + "2499200562431a3a": { + "count": 1, + "ms": 2.0548749016597867 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 663.4803750785068 + }, + { + "sequence": 5, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 7.251292001456022 + }, + "dense": { + "count": 1, + "ms": 162.91516704950482 + }, + "sparse": { + "count": 1, + "ms": 5.476542050018907 + }, + "prepare_query": { + "count": 1, + "ms": 168.45150000881404 + }, + "vespa": { + "count": 1, + "ms": 274.40424996893853 + }, + "visibility": { + "count": 1, + "ms": 54.09208405762911 + }, + "execute": { + "count": 1, + "ms": 332.5495000462979 + }, + "_enrich": { + "count": 1, + "ms": 175.91229197569191 + }, + "_coverage": { + "count": 1, + "ms": 9.203209076076746 + }, + "_final_publications": { + "count": 1, + "ms": 8.18145798984915 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 5.3002911154180765 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 21.253084065392613 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 18.066833959892392 + }, + "227345276c4b5446": { + "count": 1, + "ms": 161.16387501824647 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.8148750057443976 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 0.8872910402715206 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.3558750031515956 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 4.315625061281025 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 702.7063750429079 + }, + { + "sequence": 12, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.792165942490101 + }, + "dense": { + "count": 1, + "ms": 66.94124999921769 + }, + "sparse": { + "count": 1, + "ms": 1.9606249406933784 + }, + "prepare_query": { + "count": 1, + "ms": 68.97516606841236 + }, + "vespa": { + "count": 1, + "ms": 155.2103339927271 + }, + "visibility": { + "count": 1, + "ms": 929.0852500125766 + }, + "execute": { + "count": 1, + "ms": 1086.1585000529885 + }, + "_enrich": { + "count": 1, + "ms": 70.99999999627471 + }, + "_coverage": { + "count": 1, + "ms": 12.965750065632164 + }, + "_final_publications": { + "count": 1, + "ms": 15.56666602846235 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.2450839644297957 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.278917072340846 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 26.266500004567206 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 73.34062503650784 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 59.933540993370116 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 1.5460000140592456 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 2.7932500233873725 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 6.924125016666949 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 8.849709061905742 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 1259.9269580096006 + }, + { + "sequence": 9, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 2.710792003199458 + }, + "dense": { + "count": 1, + "ms": 69.73824999295175 + }, + "sparse": { + "count": 1, + "ms": 0.48016710206866264 + }, + "prepare_query": { + "count": 1, + "ms": 70.28787501621991 + }, + "vespa": { + "count": 1, + "ms": 1096.9839169410989 + }, + "visibility": { + "count": 1, + "ms": 39.961625006981194 + }, + "execute": { + "count": 1, + "ms": 1138.2223330438137 + }, + "_enrich": { + "count": 1, + "ms": 48.959000036120415 + }, + "_coverage": { + "count": 1, + "ms": 6.951833958737552 + }, + "_final_publications": { + "count": 1, + "ms": 7.558958022855222 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.493541058152914 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 7.188707939349115 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 22.2346669761464 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 23.249167017638683 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 17.297792015597224 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 3.2182090217247605 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.46629097778350115 + }, + "2499200562431a3a": { + "count": 1, + "ms": 3.5706249764189124 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 1275.4977500298992 + }, + { + "sequence": 11, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.292208071798086 + }, + "dense": { + "count": 1, + "ms": 62.33175005763769 + }, + "sparse": { + "count": 1, + "ms": 0.2817499917000532 + }, + "prepare_query": { + "count": 1, + "ms": 62.67604196909815 + }, + "vespa": { + "count": 1, + "ms": 1031.9742910796776 + }, + "visibility": { + "count": 1, + "ms": 102.05487499479204 + }, + "execute": { + "count": 1, + "ms": 1148.3783340081573 + }, + "_enrich": { + "count": 1, + "ms": 25.6517919478938 + }, + "_coverage": { + "count": 1, + "ms": 55.74112490285188 + }, + "_final_publications": { + "count": 1, + "ms": 17.10937498137355 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.4157919688150287 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.5292500611394644 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 63.26654099393636 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 23.13979109749198 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 15.811291057616472 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 2.5095000164583325 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 10.878624976612628 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 6.8604579428210855 + }, + "4860037db3700175": { + "count": 1, + "ms": 22.120374953374267 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.6761249387636781 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 0.2955840900540352 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 2.76004197075963 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 10.070209042169154 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 1313.444250030443 + }, + { + "sequence": 10, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 2.544791088439524 + }, + "dense": { + "count": 1, + "ms": 74.00362507905811 + }, + "sparse": { + "count": 1, + "ms": 0.7005829829722643 + }, + "prepare_query": { + "count": 1, + "ms": 74.76600003428757 + }, + "vespa": { + "count": 1, + "ms": 188.86408396065235 + }, + "visibility": { + "count": 1, + "ms": 883.0980840139091 + }, + "execute": { + "count": 1, + "ms": 1073.9188329316676 + }, + "_enrich": { + "count": 1, + "ms": 170.41695804800838 + }, + "_coverage": { + "count": 1, + "ms": 5.164291942492127 + }, + "_final_publications": { + "count": 1, + "ms": 6.403250037692487 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.664999988861382 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 831.1943749431521 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 34.70541606657207 + }, + "227345276c4b5446": { + "count": 1, + "ms": 158.93658401910216 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.985874910838902 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 1.1015419149771333 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.17662497702986002 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 2.5326249888166785 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 1333.68475001771 + }, + { + "sequence": 15, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.080042893998325 + }, + "dense": { + "count": 1, + "ms": 165.07424996234477 + }, + "sparse": { + "count": 1, + "ms": 1.0408329544588923 + }, + "prepare_query": { + "count": 1, + "ms": 166.30387492477894 + }, + "vespa": { + "count": 1, + "ms": 440.48049999400973 + }, + "visibility": { + "count": 1, + "ms": 179.47479209396988 + }, + "execute": { + "count": 1, + "ms": 621.7332079540938 + }, + "_enrich": { + "count": 1, + "ms": 100.10841605253518 + }, + "_coverage": { + "count": 1, + "ms": 15.588499954901636 + }, + "_final_publications": { + "count": 1, + "ms": 36.27545898780227 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.9191669998690486 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 70.07212506141514 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 36.27354197669774 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 52.24266601726413 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 76.24662504531443 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 2.1802500123158097 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 9.423875017091632 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.6631249561905861 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 30.16150009352714 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 944.0113340970129 + }, + { + "sequence": 16, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 6.74495892599225 + }, + "dense": { + "count": 1, + "ms": 174.3142909836024 + }, + "sparse": { + "count": 1, + "ms": 0.6041250890120864 + }, + "prepare_query": { + "count": 1, + "ms": 174.9319169903174 + }, + "vespa": { + "count": 1, + "ms": 501.8192919669673 + }, + "visibility": { + "count": 1, + "ms": 64.14791697170585 + }, + "execute": { + "count": 1, + "ms": 577.1721249911934 + }, + "_enrich": { + "count": 1, + "ms": 110.81891704816371 + }, + "_coverage": { + "count": 1, + "ms": 81.40800008550286 + }, + "_final_publications": { + "count": 1, + "ms": 48.73295803554356 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 4.563168040476739 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.170082949101925 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 23.009750060737133 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 10.039665969088674 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 80.65083296969533 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 3.7403330206871033 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 20.194957964122295 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 3.3077499829232693 + }, + "4860037db3700175": { + "count": 1, + "ms": 0.883165979757905 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.7463330402970314 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 8.105917018838227 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 16.201584017835557 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 35.00854200683534 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 1001.1217090068385 + }, + { + "sequence": 13, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 24.764375062659383 + }, + "dense": { + "count": 1, + "ms": 175.05683400668204 + }, + "sparse": { + "count": 1, + "ms": 1.1266249930486083 + }, + "prepare_query": { + "count": 1, + "ms": 176.25720903743058 + }, + "vespa": { + "count": 1, + "ms": 607.9739999258891 + }, + "visibility": { + "count": 1, + "ms": 43.512500007636845 + }, + "execute": { + "count": 1, + "ms": 653.4734159940854 + }, + "_enrich": { + "count": 1, + "ms": 96.6626669978723 + }, + "_coverage": { + "count": 1, + "ms": 22.702542017214 + }, + "_final_publications": { + "count": 1, + "ms": 35.11658392380923 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 22.74199912790209 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 16.115250065922737 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 8.703375002369285 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 70.47658402007073 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 8.543459000065923 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 14.393875026144087 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 2.33279203530401 + }, + "2499200562431a3a": { + "count": 1, + "ms": 11.552125099115074 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 1009.7570419311523 + }, + { + "sequence": 14, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.7359159905463457 + }, + "dense": { + "count": 1, + "ms": 175.78433302696794 + }, + "sparse": { + "count": 1, + "ms": 0.9996249573305249 + }, + "prepare_query": { + "count": 1, + "ms": 176.8083330243826 + }, + "vespa": { + "count": 1, + "ms": 416.2487080320716 + }, + "visibility": { + "count": 1, + "ms": 122.3345419857651 + }, + "execute": { + "count": 1, + "ms": 540.0912909535691 + }, + "_enrich": { + "count": 1, + "ms": 331.29774988628924 + }, + "_coverage": { + "count": 1, + "ms": 15.635541989468038 + }, + "_final_publications": { + "count": 1, + "ms": 25.462583056651056 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.1389579633250833 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 9.843125008046627 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 47.68558300565928 + }, + "227345276c4b5446": { + "count": 1, + "ms": 309.52387501019984 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 3.208791953511536 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 1.7917499644681811 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.4837081069126725 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 6.455249967984855 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 1093.6309170210734 + }, + { + "sequence": 19, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 22.333709057420492 + }, + "dense": { + "count": 1, + "ms": 150.0156670808792 + }, + "sparse": { + "count": 1, + "ms": 11.587167042307556 + }, + "prepare_query": { + "count": 1, + "ms": 162.66179201193154 + }, + "vespa": { + "count": 1, + "ms": 357.17754205688834 + }, + "visibility": { + "count": 1, + "ms": 129.57099999766797 + }, + "execute": { + "count": 1, + "ms": 496.22041604015976 + }, + "_enrich": { + "count": 1, + "ms": 178.12020890414715 + }, + "_coverage": { + "count": 1, + "ms": 14.941166038624942 + }, + "_final_publications": { + "count": 1, + "ms": 35.05670896265656 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 19.898416940122843 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 21.18158305529505 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 17.832041950896382 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 48.60466602258384 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 139.66970809269696 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 23.935875040479004 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 7.616457995027304 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.9968749945983291 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 20.423874957486987 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 910.839542048052 + }, + { + "sequence": 18, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 10.032958001829684 + }, + "dense": { + "count": 1, + "ms": 178.3760410035029 + }, + "sparse": { + "count": 1, + "ms": 0.5773329176008701 + }, + "prepare_query": { + "count": 1, + "ms": 178.99479204788804 + }, + "vespa": { + "count": 1, + "ms": 454.5237079728395 + }, + "visibility": { + "count": 1, + "ms": 105.22020899225026 + }, + "execute": { + "count": 1, + "ms": 571.801791083999 + }, + "_enrich": { + "count": 1, + "ms": 61.847709002904594 + }, + "_coverage": { + "count": 1, + "ms": 100.27866705786437 + }, + "_final_publications": { + "count": 1, + "ms": 22.4263749551028 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 4.487082944251597 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 20.635874941945076 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 34.77570798713714 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 22.953292005695403 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 43.44141704495996 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 3.7646250566467643 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 15.19566704519093 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.352582941763103 + }, + "4860037db3700175": { + "count": 1, + "ms": 2.4471249198541045 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.8779580239206553 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 3.8167500169947743 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 21.30549994762987 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 10.47720795031637 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 946.5210839407519 + }, + { + "sequence": 17, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.191416966728866 + }, + "dense": { + "count": 1, + "ms": 181.75150000024587 + }, + "sparse": { + "count": 1, + "ms": 19.68583301641047 + }, + "prepare_query": { + "count": 1, + "ms": 201.46366697736084 + }, + "vespa": { + "count": 1, + "ms": 509.1138749849051 + }, + "visibility": { + "count": 1, + "ms": 35.40741698816419 + }, + "execute": { + "count": 1, + "ms": 565.984082990326 + }, + "_enrich": { + "count": 1, + "ms": 189.30870899930596 + }, + "_coverage": { + "count": 1, + "ms": 18.03158305119723 + }, + "_final_publications": { + "count": 1, + "ms": 8.639499894343317 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.113042002543807 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 6.577042047865689 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 12.31124997138977 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 155.43862502090633 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 14.549666084349155 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 7.87733297329396 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.5354579770937562 + }, + "2499200562431a3a": { + "count": 1, + "ms": 3.00191598944366 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 987.1257080230862 + }, + { + "sequence": 20, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 6.4297078642994165 + }, + "dense": { + "count": 1, + "ms": 175.3441250184551 + }, + "sparse": { + "count": 1, + "ms": 19.589458010159433 + }, + "prepare_query": { + "count": 1, + "ms": 194.98054194264114 + }, + "vespa": { + "count": 1, + "ms": 377.5235840585083 + }, + "visibility": { + "count": 1, + "ms": 70.01120899803936 + }, + "execute": { + "count": 1, + "ms": 449.16333397850394 + }, + "_enrich": { + "count": 1, + "ms": 345.7779579330236 + }, + "_coverage": { + "count": 1, + "ms": 6.964874919503927 + }, + "_final_publications": { + "count": 1, + "ms": 14.39250004477799 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 5.538332974538207 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 31.704834080301225 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 12.205415987409651 + }, + "227345276c4b5446": { + "count": 1, + "ms": 332.4839168926701 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.7360000163316727 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 2.4184579961001873 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.4141249228268862 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 3.7707919254899025 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 1019.4548750296235 + }, + { + "sequence": 24, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.785124910995364 + }, + "dense": { + "count": 1, + "ms": 138.52658297400922 + }, + "sparse": { + "count": 1, + "ms": 0.21062500309199095 + }, + "prepare_query": { + "count": 1, + "ms": 138.77249998040497 + }, + "vespa": { + "count": 1, + "ms": 1715.635541942902 + }, + "visibility": { + "count": 1, + "ms": 132.22104194574058 + }, + "execute": { + "count": 1, + "ms": 1849.270167062059 + }, + "_enrich": { + "count": 1, + "ms": 52.795750088989735 + }, + "_coverage": { + "count": 1, + "ms": 26.668832986615598 + }, + "_final_publications": { + "count": 1, + "ms": 10.617665946483612 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.2279999684542418 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 10.033750091679394 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 41.89637501258403 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 34.7755829570815 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 31.50499996263534 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 2.8527920367196202 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 21.01708296686411 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.4943750100210309 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 6.05637498665601 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 2083.523500012234 + }, + { + "sequence": 23, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.943749936297536 + }, + "dense": { + "count": 1, + "ms": 136.26425003167242 + }, + "sparse": { + "count": 1, + "ms": 3.064083051867783 + }, + "prepare_query": { + "count": 1, + "ms": 139.38750000670552 + }, + "vespa": { + "count": 1, + "ms": 1803.4203751012683 + }, + "visibility": { + "count": 1, + "ms": 73.67808290291578 + }, + "execute": { + "count": 1, + "ms": 1880.1395420450717 + }, + "_enrich": { + "count": 1, + "ms": 49.1235830122605 + }, + "_coverage": { + "count": 1, + "ms": 57.574499980546534 + }, + "_final_publications": { + "count": 1, + "ms": 12.067666975781322 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.7562910690903664 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 14.353249920532107 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 19.084625062532723 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 24.221083032898605 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 26.21775004081428 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 7.67866603564471 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 9.051040979102254 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.0478750336915255 + }, + "4860037db3700175": { + "count": 1, + "ms": 16.732458025217056 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 3.302415949292481 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 2.9357500607147813 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 6.882583140395582 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 8.687541936524212 + } + }, + "errors": [], + "candidates": 200, + "total_ms": 2142.9362499620765 + }, + { + "sequence": 21, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 10.725875967182219 + }, + "dense": { + "count": 1, + "ms": 163.01108396146446 + }, + "sparse": { + "count": 1, + "ms": 0.859708059579134 + }, + "prepare_query": { + "count": 1, + "ms": 164.00095797143877 + }, + "vespa": { + "count": 1, + "ms": 1751.5132919652387 + }, + "visibility": { + "count": 1, + "ms": 63.263374962843955 + }, + "execute": { + "count": 1, + "ms": 1816.254833014682 + }, + "_enrich": { + "count": 1, + "ms": 139.0849999152124 + }, + "_coverage": { + "count": 1, + "ms": 7.946375058963895 + }, + "_final_publications": { + "count": 1, + "ms": 10.63124998472631 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 6.748125073499978 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 5.06999995559454 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 30.09774989914149 + }, + "227345276c4b5446": { + "count": 1, + "ms": 124.29991597309709 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 6.123207975178957 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 3.836291958577931 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.8284159703180194 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 5.511333933100104 + } + }, + "errors": [], + "candidates": 196, + "total_ms": 2149.137625005096 + }, + { + "sequence": 22, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.825417068786919 + }, + "dense": { + "count": 1, + "ms": 130.95570902805775 + }, + "sparse": { + "count": 1, + "ms": 4.648332949727774 + }, + "prepare_query": { + "count": 1, + "ms": 135.95612498465925 + }, + "vespa": { + "count": 1, + "ms": 2608.8416669517756 + }, + "visibility": { + "count": 1, + "ms": 22.659041918814182 + }, + "execute": { + "count": 1, + "ms": 2633.3210000302643 + }, + "_enrich": { + "count": 1, + "ms": 18.839332973584533 + }, + "_coverage": { + "count": 1, + "ms": 10.801083990372717 + }, + "_final_publications": { + "count": 1, + "ms": 10.405333014205098 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.020083880983293 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.6238749483600259 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 14.410667004995048 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 8.465667022392154 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 0.7834999123588204 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 3.1355000101029873 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.3542909398674965 + }, + "2499200562431a3a": { + "count": 1, + "ms": 5.137000000104308 + } + }, + "errors": [], + "candidates": 198, + "total_ms": 2814.926875056699 + } + ] +} From e37d018de786d3b9446b31a100a8dcd1c4ae9009 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 16:46:07 -0700 Subject: [PATCH 161/266] Cancel owned index search when its HTTP client disconnects --- backend/airweave/api/v1/endpoints/records.py | 25 ++- .../api/test_owned_search_cancellation.py | 178 ++++++++++++++++++ 2 files changed, 201 insertions(+), 2 deletions(-) create mode 100644 backend/tests/unit/api/test_owned_search_cancellation.py diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index 4587df6b1e..8c088911ad 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -1,9 +1,10 @@ """Internal service API for committed source records; provider APIs are never read here.""" +from asyncio import FIRST_COMPLETED, create_task, gather, wait from typing import Literal from uuid import UUID -from fastapi import Depends, Path, Query, Request, Response +from fastapi import Depends, HTTPException, Path, Query, Request, Response from fastapi.responses import JSONResponse from sqlalchemy.ext.asyncio import AsyncSession @@ -31,15 +32,35 @@ router = TrailingSlashRouter() +async def _search_disconnected(request: Request) -> None: + # FastAPI has consumed this read-only search POST's body before entering the route. + while (await request.receive())["type"] != "http.disconnect": + pass + + @router.post("/search", response_model=OwnedSearchResponse) async def search_records( request: OwnedSearchRequest, + http_request: Request, db: AsyncSession = Depends(get_db), ctx: ApiContext = Depends(deps.get_context), container: Container = Depends(deps.get_container), ) -> OwnedSearchResponse: """Retrieve bounded indexed originals; no provider requests or agent execution.""" - return await container.owned_search.search(db, ctx, request) + work = create_task(container.owned_search.search(db, ctx, request)) + disconnected = create_task(_search_disconnected(http_request)) + try: + done, _ = await wait((work, disconnected), return_when=FIRST_COMPLETED) + if work in done: + return await work + await disconnected + raise HTTPException(499, "Search client disconnected") + finally: + # Join async work before request-scoped dependencies close. Already-running + # synchronous Vespa or inference work cannot be forcibly stopped here. + work.cancel() + disconnected.cancel() + await gather(work, disconnected, return_exceptions=True) async def record_error_response(request: Request, error: CanonicalStoreError) -> JSONResponse: diff --git a/backend/tests/unit/api/test_owned_search_cancellation.py b/backend/tests/unit/api/test_owned_search_cancellation.py new file mode 100644 index 0000000000..30bc2605a7 --- /dev/null +++ b/backend/tests/unit/api/test_owned_search_cancellation.py @@ -0,0 +1,178 @@ +"""Read-only search cancellation on real TCP, with explicit fixture dependencies.""" + +import asyncio +import socket +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 + +import httpx +import pytest +import uvicorn +from fastapi import FastAPI, Header, HTTPException + +from airweave.api import deps +from airweave.api.v1.endpoints.records import router +from airweave.db.session import get_db +from airweave.domains.search.owned_models import OwnedSearchResponse + +BODY = {"query": "synthetic", "sync_ids": [str(uuid4())], "mode": "keyword"} +HEADERS = {"X-API-Key": "fixture-only"} + + +def response(): + return OwnedSearchResponse( + items=(), + sources=(), + candidate_window_full=False, + engine_partial=False, + excluded_candidates=0, + postfilter_excluded=0, + retrieval_incomplete=False, + ) + + +def app_for(search, released): + app = FastAPI() + app.include_router(router, prefix="/sync") + + async def database(): + try: + yield object() + finally: + released.set() + + async def context(x_api_key: str = Header()): + if x_api_key != "fixture-only": + raise HTTPException(401, "Invalid fixture key") + return SimpleNamespace(organization=SimpleNamespace(id=uuid4())) + + app.dependency_overrides[get_db] = database + app.dependency_overrides[deps.get_context] = context + app.dependency_overrides[deps.get_container] = lambda: SimpleNamespace( + owned_search=SimpleNamespace(search=search) + ) + return app + + +@pytest.mark.asyncio +async def test_search_post_preserves_response_errors_and_auth_without_disconnect_wait(): + released = asyncio.Event() + search = AsyncMock(return_value=response()) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app_for(search, released)), base_url="http://test" + ) as client: + async with asyncio.timeout(3): + good = await client.post("/sync/search", json=BODY, headers=HEADERS) + assert good.status_code == 200 and good.json()["items"] == [] + assert released.is_set() + search.side_effect = HTTPException(409, {"code": "reindex_required"}) + failed = await client.post("/sync/search", json=BODY, headers=HEADERS) + assert failed.status_code == 409 + assert failed.json()["detail"]["code"] == "reindex_required" + before = search.await_count + denied = await client.post("/sync/search", json=BODY, headers={"X-API-Key": "wrong"}) + assert denied.status_code == 401 and search.await_count == before + + +@pytest.mark.asyncio +async def test_tcp_search_disconnect_joins_work_before_database_dependency_closes(): + entered, cancelled, released, unblock = (asyncio.Event() for _ in range(4)) + + async def search(db, ctx, request): + assert request.query == BODY["query"] # Parsed POST body survives the listener. + entered.set() + try: + await unblock.wait() + return response() + except asyncio.CancelledError: + assert not released.is_set() + cancelled.set() + raise + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(16) + server = uvicorn.Server( + uvicorn.Config(app_for(search, released), lifespan="off", log_config=None, access_log=False) + ) + serving = asyncio.create_task(server.serve(sockets=[listener])) + try: + async with asyncio.timeout(5): + while not server.started: + if serving.done(): + await serving + await asyncio.sleep(0.001) + async with httpx.AsyncClient(timeout=None, trust_env=False) as client: + task = asyncio.create_task( + client.post( + f"http://127.0.0.1:{listener.getsockname()[1]}/sync/search", + json=BODY, + headers=HEADERS, + ) + ) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout=2) + await asyncio.wait_for(released.wait(), timeout=2) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + finally: + unblock.set() + server.should_exit = True + try: + await asyncio.wait_for(serving, timeout=5) + finally: + listener.close() + + +@pytest.mark.asyncio +async def test_completed_search_error_wins_simultaneous_disconnect(): + from airweave.api.v1.endpoints.records import search_records + from airweave.domains.search.owned_models import OwnedSearchRequest + + error = HTTPException(409, {"code": "reindex_required"}) + container = SimpleNamespace(owned_search=SimpleNamespace(search=AsyncMock(side_effect=error))) + request = SimpleNamespace(receive=AsyncMock(return_value={"type": "http.disconnect"})) + with pytest.raises(HTTPException) as caught: + await search_records(OwnedSearchRequest(**BODY), request, object(), object(), container) + assert caught.value is error + + +@pytest.mark.asyncio +async def test_caller_cancellation_joins_search_and_disconnect_listener(): + from airweave.api.v1.endpoints.records import search_records + from airweave.domains.search.owned_models import OwnedSearchRequest + + entered = asyncio.Event() + listener_entered = asyncio.Event() + stopped = set() + + async def blocked(name, ready): + ready.set() + try: + await asyncio.Event().wait() + finally: + stopped.add(name) + + container = SimpleNamespace( + owned_search=SimpleNamespace(search=lambda *args: blocked("search", entered)) + ) + request = SimpleNamespace(receive=lambda: blocked("listener", listener_entered)) + task = asyncio.create_task( + search_records(OwnedSearchRequest(**BODY), request, object(), object(), container) + ) + try: + await asyncio.wait_for(entered.wait(), 2) + await asyncio.wait_for(listener_entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert stopped == {"search", "listener"} + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) From a539e40979871f5ba23428b802845b5f1329f6c9 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 16:55:34 -0700 Subject: [PATCH 162/266] Record native Vespa and worker latency decomposition --- docs/owned-sources/search-execution.md | 37 + docs/owned-sources/search-worker-samples.json | 2454 +++++++++++++++++ 2 files changed, 2491 insertions(+) create mode 100644 docs/owned-sources/search-worker-samples.json diff --git a/docs/owned-sources/search-execution.md b/docs/owned-sources/search-execution.md index 32eb585679..954751a5e4 100644 --- a/docs/owned-sources/search-execution.md +++ b/docs/owned-sources/search-execution.md @@ -107,3 +107,40 @@ provides `querytime`, `summaryfetchtime` and `searchtime` in seconds when body and add only that presentation flag, retaining those three numeric values alongside worker-entry/exit times. No query trace, response body or result content needs to be logged. Pooling changes should wait for that decomposition. + + +## Worker and native Vespa decomposition + +A second instrumented candidate ran immutable commit +`e37d018de786d3b9446b31a100a8dcd1c4ae9009` on loopback port 18084. +Authentication qualification returned 401 without a key and 200 for an authorized +retained-record read. Almanac then switched its private origin to this candidate; +all six product requests returned 200 with 25 records and identical result-identity +digests. The original 18081 process was stopped only after this success. The candidate +continues serving the same isolated corpus; no provider fetch, schema migration or +operator identity remapping occurred. + +Product times were 2439, 922, 1333, 565, 660 and 1171 ms. +[24 sanitized worker samples](search-worker-samples.json) record native Vespa timing +and worker boundaries. The critical organization request per product request was: + +| Sample | Airweave total | Dense | Worker queue | SDK worker | Resume delay | Native Vespa search | All SQL | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 1702 | 703 | 0.1 | 900 | 0.2 | 509 | 47 | +| 2 | 818 | 161 | 3.3 | 358 | 0.2 | 96 | 226 | +| 3 | 1252 | 78 | 1.2 | 196 | 0.4 | 69 | 946 | +| 4 | 487 | 94 | 0.7 | 181 | 0.4 | 84 | 178 | +| 5 | 535 | 128 | 0.7 | 359 | 0.9 | 172 | 20 | +| 6 | 900 | 248 | 0.2 | 321 | 0.3 | 122 | 286 | + +Milliseconds; nested intervals again must not be summed. On these critical requests, +thread queueing and event-loop resumption are small. Native Vespa processing explains +part of SDK wall time; the remainder includes transport/session setup and JSON decode, +which were not measured separately. Sample 3 again contains a SQL delay. This does not +justify blaming thread starvation, choosing a database plan, or changing HTTP pooling +without a targeted test. Six successful samples neither establish a latency percentile +nor reconstruct the earlier unlogged 503 and 13-second response. + +Instrumentation lives only in the private launcher, preserving general log suppression. +Its copied query body adds `presentation.timing` only; result selection, ranking, +authorization and the 200-candidate bound remain unchanged. diff --git a/docs/owned-sources/search-worker-samples.json b/docs/owned-sources/search-worker-samples.json new file mode 100644 index 0000000000..1953567584 --- /dev/null +++ b/docs/owned-sources/search-worker-samples.json @@ -0,0 +1,2454 @@ +{ + "immutable_revision": "e37d018de786d3b9446b31a100a8dcd1c4ae9009", + "qualification": "Actual API on loopback18084 from immutable archive, 6serialproductqueries/24organizationrequests. Numeric stage/native timings only; no queries, source identifiers, SQLparameters or resultcontent. Intervals are nested walltime, not additiveCPU.", + "requests": [ + { + "sequence": 4, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 14.39700007904321 + }, + "dense": { + "count": 1, + "ms": 697.7981670061126 + }, + "sparse": { + "count": 1, + "ms": 6.554250023327768 + }, + "prepare_query": { + "count": 1, + "ms": 704.3882500147447 + }, + "vespa": { + "count": 1, + "ms": 652.7667089831084 + }, + "visibility": { + "count": 1, + "ms": 92.36441599205136 + }, + "execute": { + "count": 1, + "ms": 748.8067910308018 + }, + "_enrich": { + "count": 1, + "ms": 42.51174989622086 + }, + "_coverage": { + "count": 1, + "ms": 27.212833054363728 + }, + "_final_publications": { + "count": 1, + "ms": 14.53358307480812 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 13.119000010192394 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 3.1570419669151306 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 54.84887503553182 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 14.02979192789644 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 13.554790988564491 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 21.54762507416308 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 14.26812494173646 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 9.441999951377511 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 6.412832997739315 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 714.5442909095436, + "thread_submit_ms": 714.5484999055043, + "worker_enter_ms": 719.8180829873309, + "worker_exit_ms": 1364.0451659448445, + "await_resume_ms": 1364.1532909823582, + "adapter_return_ms": 1367.3077079001814 + }, + "vespa_native_seconds": { + "querytime": 0.244, + "summaryfetchtime": 0.082, + "searchtime": 0.333 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 1552.3313329322264 + }, + { + "sequence": 1, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 20.058708963915706 + }, + "dense": { + "count": 1, + "ms": 691.7723750229925 + }, + "sparse": { + "count": 1, + "ms": 13.597415992990136 + }, + "prepare_query": { + "count": 1, + "ms": 705.7115830248222 + }, + "vespa": { + "count": 1, + "ms": 688.2408329984173 + }, + "visibility": { + "count": 1, + "ms": 29.30437494069338 + }, + "execute": { + "count": 1, + "ms": 721.8387499451637 + }, + "_enrich": { + "count": 1, + "ms": 102.58941701613367 + }, + "_coverage": { + "count": 1, + "ms": 13.002125080674887 + }, + "_final_publications": { + "count": 1, + "ms": 7.20037508290261 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 5.277417018078268 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 12.694583041593432 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 8.783917059190571 + }, + "227345276c4b5446": { + "count": 1, + "ms": 79.17404198087752 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 11.609292007051408 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 9.656250011175871 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.16062508802860975 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 4.063541069626808 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 731.3685000408441, + "thread_submit_ms": 731.857500039041, + "worker_enter_ms": 732.0340829901397, + "worker_exit_ms": 1416.5454170433804, + "await_resume_ms": 1416.772291995585, + "adapter_return_ms": 1419.607250019908 + }, + "vespa_native_seconds": { + "querytime": 0.23, + "summaryfetchtime": 0.10400000000000001, + "searchtime": 0.336 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 1573.6422500340268 + }, + { + "sequence": 3, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 8.467874024063349 + }, + "dense": { + "count": 1, + "ms": 699.1888750344515 + }, + "sparse": { + "count": 1, + "ms": 6.565375020727515 + }, + "prepare_query": { + "count": 1, + "ms": 705.7651670183986 + }, + "vespa": { + "count": 1, + "ms": 678.466665907763 + }, + "visibility": { + "count": 1, + "ms": 61.461207922548056 + }, + "execute": { + "count": 1, + "ms": 743.8884590519592 + }, + "_enrich": { + "count": 1, + "ms": 35.72583303321153 + }, + "_coverage": { + "count": 1, + "ms": 114.05433295294642 + }, + "_final_publications": { + "count": 1, + "ms": 21.25195797998458 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 6.0969998594373465 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 9.712875005789101 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 31.63708304055035 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 5.438291002064943 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 12.6727499300614 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 12.342790956608951 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 18.270791973918676 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 2.7705830289050937 + }, + "4860037db3700175": { + "count": 1, + "ms": 40.20195791963488 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 1.977749983780086 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 2.1070840302854776 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 4.079291946254671 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 15.472374972887337 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 715.1667090365663, + "thread_submit_ms": 715.1694999774918, + "worker_enter_ms": 720.6202920060605, + "worker_exit_ms": 1390.5022089602426, + "await_resume_ms": 1390.915249940008, + "adapter_return_ms": 1393.631541985087 + }, + "vespa_native_seconds": { + "querytime": 0.23800000000000002, + "summaryfetchtime": 0.101, + "searchtime": 0.341 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 1629.7802090412006 + }, + { + "sequence": 2, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 7.911792024970055 + }, + "dense": { + "count": 1, + "ms": 702.8400419512764 + }, + "sparse": { + "count": 1, + "ms": 12.906166957691312 + }, + "prepare_query": { + "count": 1, + "ms": 715.8160830149427 + }, + "vespa": { + "count": 1, + "ms": 905.7754579698667 + }, + "visibility": { + "count": 1, + "ms": 24.27075000014156 + }, + "execute": { + "count": 1, + "ms": 930.7643749052659 + }, + "_enrich": { + "count": 1, + "ms": 25.458791060373187 + }, + "_coverage": { + "count": 1, + "ms": 12.371707940474153 + }, + "_final_publications": { + "count": 1, + "ms": 8.654584060423076 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 6.06404198333621 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 4.465624922886491 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 10.673750075511634 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 17.301541985943913 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 1.156416954472661 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 4.220000002533197 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.218083081766963 + }, + "2499200562431a3a": { + "count": 1, + "ms": 2.209833008237183 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 722.4849170306697, + "thread_submit_ms": 722.493875073269, + "worker_enter_ms": 722.5975840119645, + "worker_exit_ms": 1623.0179590638727, + "await_resume_ms": 1623.1887920293957, + "adapter_return_ms": 1628.2583340071142 + }, + "vespa_native_seconds": { + "querytime": 0.257, + "summaryfetchtime": 0.249, + "searchtime": 0.509 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 1701.5393340261653 + }, + { + "sequence": 5, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 6.998000084422529 + }, + "dense": { + "count": 1, + "ms": 161.28495789598674 + }, + "sparse": { + "count": 1, + "ms": 3.144333022646606 + }, + "prepare_query": { + "count": 1, + "ms": 164.9724580347538 + }, + "vespa": { + "count": 1, + "ms": 342.02920901589096 + }, + "visibility": { + "count": 1, + "ms": 65.0430410169065 + }, + "execute": { + "count": 1, + "ms": 409.411457949318 + }, + "_enrich": { + "count": 1, + "ms": 89.88524996675551 + }, + "_coverage": { + "count": 1, + "ms": 13.895208016037941 + }, + "_final_publications": { + "count": 1, + "ms": 26.216291007585824 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 4.420625977218151 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 7.723249960690737 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 27.281332993879914 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 10.08245802950114 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 65.84966601803899 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 14.081792091019452 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 5.074832937680185 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.9725000020116568 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 13.51004105526954 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 172.04616591334343, + "thread_submit_ms": 172.42312501184642, + "worker_enter_ms": 173.29445795621723, + "worker_exit_ms": 509.4542499864474, + "await_resume_ms": 509.67304094228894, + "adapter_return_ms": 514.0712499851361 + }, + "vespa_native_seconds": { + "querytime": 0.076, + "summaryfetchtime": 0.039, + "searchtime": 0.116 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 712.3994579305872 + }, + { + "sequence": 8, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 15.251790988259017 + }, + "dense": { + "count": 1, + "ms": 154.03837501071393 + }, + "sparse": { + "count": 1, + "ms": 4.9700409872457385 + }, + "prepare_query": { + "count": 1, + "ms": 159.11116695497185 + }, + "vespa": { + "count": 1, + "ms": 448.9807500503957 + }, + "visibility": { + "count": 1, + "ms": 35.94641690142453 + }, + "execute": { + "count": 1, + "ms": 488.1273750215769 + }, + "_enrich": { + "count": 1, + "ms": 70.2489159302786 + }, + "_coverage": { + "count": 1, + "ms": 39.53712503425777 + }, + "_final_publications": { + "count": 1, + "ms": 14.14441701490432 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 12.769749970175326 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.978458090685308 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 5.436790990643203 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 16.408832976594567 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 58.04041598457843 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 2.911832998506725 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 7.438958971761167 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 10.581750073470175 + }, + "4860037db3700175": { + "count": 1, + "ms": 0.4955410258844495 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.39175001438707113 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 0.6935000419616699 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 2.578666084446013 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 7.482332992367446 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 169.43408292718232, + "thread_submit_ms": 169.43954199086875, + "worker_enter_ms": 172.78091702610254, + "worker_exit_ms": 608.765707933344, + "await_resume_ms": 612.2454999713227, + "adapter_return_ms": 618.4113329509273 + }, + "vespa_native_seconds": { + "querytime": 0.07200000000000001, + "summaryfetchtime": 0.036000000000000004, + "searchtime": 0.11 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 787.4784169252962 + }, + { + "sequence": 6, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.894750076346099 + }, + "dense": { + "count": 1, + "ms": 161.17279091849923 + }, + "sparse": { + "count": 1, + "ms": 6.957000005058944 + }, + "prepare_query": { + "count": 1, + "ms": 168.1465000147 + }, + "vespa": { + "count": 1, + "ms": 480.32024991698563 + }, + "visibility": { + "count": 1, + "ms": 36.98929201345891 + }, + "execute": { + "count": 1, + "ms": 527.5187079096213 + }, + "_enrich": { + "count": 1, + "ms": 42.28266607969999 + }, + "_coverage": { + "count": 1, + "ms": 24.55241698771715 + }, + "_final_publications": { + "count": 1, + "ms": 22.843292099423707 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.8234589155763388 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 11.09312498010695 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 20.848499960266054 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 31.272124964743853 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 1.2763330014422536 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 10.02533407881856 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 5.498540936969221 + }, + "2499200562431a3a": { + "count": 1, + "ms": 17.50499999616295 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 175.64520810265094, + "thread_submit_ms": 175.67504104226828, + "worker_enter_ms": 179.79354108683765, + "worker_exit_ms": 647.8493330068886, + "await_resume_ms": 650.1307500293478, + "adapter_return_ms": 655.9629580006003 + }, + "vespa_native_seconds": { + "querytime": 0.075, + "summaryfetchtime": 0.098, + "searchtime": 0.17500000000000002 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 791.4707080926746 + }, + { + "sequence": 7, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.7414170801639557 + }, + "dense": { + "count": 1, + "ms": 161.01354197598994 + }, + "sparse": { + "count": 1, + "ms": 4.646750050596893 + }, + "prepare_query": { + "count": 1, + "ms": 165.6768750399351 + }, + "vespa": { + "count": 1, + "ms": 386.19487499818206 + }, + "visibility": { + "count": 1, + "ms": 40.67708400543779 + }, + "execute": { + "count": 1, + "ms": 430.91362493578345 + }, + "_enrich": { + "count": 1, + "ms": 204.53729096334428 + }, + "_coverage": { + "count": 1, + "ms": 4.470500047318637 + }, + "_final_publications": { + "count": 1, + "ms": 8.46066593658179 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.8975419225171208 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 6.736875046044588 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 23.642083979211748 + }, + "227345276c4b5446": { + "count": 1, + "ms": 188.68970894254744 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.9001250145956874 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 0.9309579618275166 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.2186669735237956 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 2.6028340216726065 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 170.5259169684723, + "thread_submit_ms": 170.5344170331955, + "worker_enter_ms": 173.81800001021475, + "worker_exit_ms": 532.1142920292914, + "await_resume_ms": 532.2840840090066, + "adapter_return_ms": 556.7168339621276 + }, + "vespa_native_seconds": { + "querytime": 0.067, + "summaryfetchtime": 0.026000000000000002, + "searchtime": 0.096 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 818.2105419691652 + }, + { + "sequence": 12, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 8.565208176150918 + }, + "dense": { + "count": 1, + "ms": 76.52866700664163 + }, + "sparse": { + "count": 1, + "ms": 0.7226669695228338 + }, + "prepare_query": { + "count": 1, + "ms": 77.32437504455447 + }, + "vespa": { + "count": 1, + "ms": 174.11716596689075 + }, + "visibility": { + "count": 1, + "ms": 34.63616594672203 + }, + "execute": { + "count": 1, + "ms": 212.0330409379676 + }, + "_enrich": { + "count": 1, + "ms": 710.2265419671312 + }, + "_coverage": { + "count": 1, + "ms": 20.33891691826284 + }, + "_final_publications": { + "count": 1, + "ms": 33.103249967098236 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 4.8387921415269375 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.34675002098083496 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 5.06216601934284 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 20.76575008686632 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 687.4182500177994 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 14.593999949283898 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 10.954124969430268 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 4.585916991345584 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 23.892374942079186 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 85.20262502133846, + "thread_submit_ms": 85.29004093725234, + "worker_enter_ms": 89.04145797714591, + "worker_exit_ms": 255.55124995298684, + "await_resume_ms": 255.90416602790356, + "adapter_return_ms": 259.317290969193 + }, + "vespa_native_seconds": { + "querytime": 0.041, + "summaryfetchtime": 0.019, + "searchtime": 0.061 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 1062.9791660467163 + }, + { + "sequence": 10, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 21.33975096512586 + }, + "dense": { + "count": 1, + "ms": 75.94337500631809 + }, + "sparse": { + "count": 1, + "ms": 1.6167500289157033 + }, + "prepare_query": { + "count": 1, + "ms": 77.60191604029387 + }, + "vespa": { + "count": 1, + "ms": 194.8804589919746 + }, + "visibility": { + "count": 1, + "ms": 34.99312500935048 + }, + "execute": { + "count": 1, + "ms": 231.5204170299694 + }, + "_enrich": { + "count": 1, + "ms": 685.0154160056263 + }, + "_coverage": { + "count": 1, + "ms": 67.61808297596872 + }, + "_final_publications": { + "count": 1, + "ms": 71.81354099884629 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 17.993916990235448 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 6.8410830572247505 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 10.948791052214801 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 8.314082981087267 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 668.7928329920396 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 8.747541927732527 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 9.087041020393372 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 8.88729200232774 + }, + "4860037db3700175": { + "count": 1, + "ms": 3.888124949298799 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 2.902667038142681 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 10.116999968886375 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 5.833208095282316 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 16.19304099585861 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 85.22920799441636, + "thread_submit_ms": 85.25708294473588, + "worker_enter_ms": 85.50212497357279, + "worker_exit_ms": 273.845124989748, + "await_resume_ms": 277.19900000374764, + "adapter_return_ms": 280.1077499752864 + }, + "vespa_native_seconds": { + "querytime": 0.026000000000000002, + "summaryfetchtime": 0.035, + "searchtime": 0.063 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 1157.5952499406412 + }, + { + "sequence": 9, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 8.465541061013937 + }, + "dense": { + "count": 1, + "ms": 104.20608299318701 + }, + "sparse": { + "count": 1, + "ms": 0.9832920040935278 + }, + "prepare_query": { + "count": 1, + "ms": 105.32324993982911 + }, + "vespa": { + "count": 1, + "ms": 881.9307920057327 + }, + "visibility": { + "count": 1, + "ms": 29.598708963021636 + }, + "execute": { + "count": 1, + "ms": 912.5305000925437 + }, + "_enrich": { + "count": 1, + "ms": 181.9128750357777 + }, + "_coverage": { + "count": 1, + "ms": 16.407124930992723 + }, + "_final_publications": { + "count": 1, + "ms": 10.464500053785741 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 5.84354205057025 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 3.6404170095920563 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 17.10862503387034 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 165.6676250277087 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 4.445333033800125 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 10.916332947090268 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.9451670339331031 + }, + "2499200562431a3a": { + "count": 1, + "ms": 3.2458750065416098 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 107.84874996170402, + "thread_submit_ms": 107.86354192532599, + "worker_enter_ms": 107.92920901440084, + "worker_exit_ms": 987.0548750041053, + "await_resume_ms": 987.5456669833511, + "adapter_return_ms": 989.7763750050217 + }, + "vespa_native_seconds": { + "querytime": 0.022, + "summaryfetchtime": 0.078, + "searchtime": 0.101 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 1236.1256669973955 + }, + { + "sequence": 11, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 5.712416023015976 + }, + "dense": { + "count": 1, + "ms": 78.43829190824181 + }, + "sparse": { + "count": 1, + "ms": 0.6155420560389757 + }, + "prepare_query": { + "count": 1, + "ms": 79.08908405806869 + }, + "vespa": { + "count": 1, + "ms": 199.77108296006918 + }, + "visibility": { + "count": 1, + "ms": 24.06924997922033 + }, + "execute": { + "count": 1, + "ms": 224.90783303510398 + }, + "_enrich": { + "count": 1, + "ms": 929.6094160526991 + }, + "_coverage": { + "count": 1, + "ms": 4.354374948889017 + }, + "_final_publications": { + "count": 1, + "ms": 7.863500039093196 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.8265829207375646 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 4.591165925376117 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 10.602999944239855 + }, + "227345276c4b5446": { + "count": 1, + "ms": 922.780166962184 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.7919160416349769 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 1.0205419966951013 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.21600001491606236 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 1.944832969456911 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 84.82829201966524, + "thread_submit_ms": 84.83520802110434, + "worker_enter_ms": 86.03083295747638, + "worker_exit_ms": 282.12995803914964, + "await_resume_ms": 282.57870802190155, + "adapter_return_ms": 284.59716704674065 + }, + "vespa_native_seconds": { + "querytime": 0.034, + "summaryfetchtime": 0.034, + "searchtime": 0.069 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 1252.1563329501078 + }, + { + "sequence": 14, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 8.505958947353065 + }, + "dense": { + "count": 1, + "ms": 77.20608299132437 + }, + "sparse": { + "count": 1, + "ms": 1.3491669669747353 + }, + "prepare_query": { + "count": 1, + "ms": 78.60499992966652 + }, + "vespa": { + "count": 1, + "ms": 185.96895900554955 + }, + "visibility": { + "count": 1, + "ms": 78.58295796904713 + }, + "execute": { + "count": 1, + "ms": 266.189958085306 + }, + "_enrich": { + "count": 1, + "ms": 36.97695804294199 + }, + "_coverage": { + "count": 1, + "ms": 7.01791699975729 + }, + "_final_publications": { + "count": 1, + "ms": 19.737666007131338 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 7.61345901992172 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.6217500194907188 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 11.499374988488853 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 43.80200000014156 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 23.803792078979313 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 1.8334579654037952 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 3.497124998830259 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.4103750688955188 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 13.730291975662112 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 81.73179102595896, + "thread_submit_ms": 81.75112504977733, + "worker_enter_ms": 81.90645801369101, + "worker_exit_ms": 264.2532910685986, + "await_resume_ms": 264.5625410368666, + "adapter_return_ms": 267.6987500162795 + }, + "vespa_native_seconds": { + "querytime": 0.044, + "summaryfetchtime": 0.027, + "searchtime": 0.076 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 417.4483750248328 + }, + { + "sequence": 16, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.9521249709650874 + }, + "dense": { + "count": 1, + "ms": 79.67091700993478 + }, + "sparse": { + "count": 1, + "ms": 3.3190420363098383 + }, + "prepare_query": { + "count": 1, + "ms": 83.04604096338153 + }, + "vespa": { + "count": 1, + "ms": 213.26254203449935 + }, + "visibility": { + "count": 1, + "ms": 49.71575003582984 + }, + "execute": { + "count": 1, + "ms": 278.72954204212874 + }, + "_enrich": { + "count": 1, + "ms": 28.903541970066726 + }, + "_coverage": { + "count": 1, + "ms": 56.68974993750453 + }, + "_final_publications": { + "count": 1, + "ms": 19.680165918543935 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.7149999514222145 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.9695420050993562 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 16.130291973240674 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 4.977707983925939 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 18.96450005006045 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 2.739582909271121 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 6.958708050660789 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.9093750743195415 + }, + "4860037db3700175": { + "count": 1, + "ms": 32.64320897869766 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 1.2667080154642463 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 0.49120804760605097 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 2.533750026486814 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 16.540458076633513 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 93.54887495283037, + "thread_submit_ms": 93.56145793572068, + "worker_enter_ms": 94.34149996377528, + "worker_exit_ms": 297.8986669331789, + "await_resume_ms": 300.5399169633165, + "adapter_return_ms": 306.80899997241795 + }, + "vespa_native_seconds": { + "querytime": 0.043000000000000003, + "summaryfetchtime": 0.051000000000000004, + "searchtime": 0.095 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 471.67395800352097 + }, + { + "sequence": 13, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 5.820124875754118 + }, + "dense": { + "count": 1, + "ms": 93.94433302804828 + }, + "sparse": { + "count": 1, + "ms": 1.1335000162944198 + }, + "prepare_query": { + "count": 1, + "ms": 95.12208309024572 + }, + "vespa": { + "count": 1, + "ms": 319.22808301169425 + }, + "visibility": { + "count": 1, + "ms": 18.21220899000764 + }, + "execute": { + "count": 1, + "ms": 339.0349169494584 + }, + "_enrich": { + "count": 1, + "ms": 22.3297500051558 + }, + "_coverage": { + "count": 1, + "ms": 14.834166970103979 + }, + "_final_publications": { + "count": 1, + "ms": 9.104582946747541 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 4.30383300408721 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.797541044652462 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 2.975999959744513 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 12.258249917067587 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 3.5086669959127903 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 10.332042002119124 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.1515000369399786 + }, + "2499200562431a3a": { + "count": 1, + "ms": 6.353541975840926 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 98.91566692385823, + "thread_submit_ms": 98.921624943614, + "worker_enter_ms": 99.02224992401898, + "worker_exit_ms": 414.5970839308575, + "await_resume_ms": 414.8382089333609, + "adapter_return_ms": 418.14120893832296 + }, + "vespa_native_seconds": { + "querytime": 0.045, + "summaryfetchtime": 0.116, + "searchtime": 0.162 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 486.68208392336965 + }, + { + "sequence": 15, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 2.754500019364059 + }, + "dense": { + "count": 1, + "ms": 93.76249997876585 + }, + "sparse": { + "count": 1, + "ms": 1.439583022147417 + }, + "prepare_query": { + "count": 1, + "ms": 95.2114169485867 + }, + "vespa": { + "count": 1, + "ms": 186.33945798501372 + }, + "visibility": { + "count": 1, + "ms": 55.72645796928555 + }, + "execute": { + "count": 1, + "ms": 243.3704580180347 + }, + "_enrich": { + "count": 1, + "ms": 127.53029202576727 + }, + "_coverage": { + "count": 1, + "ms": 10.312499944120646 + }, + "_final_publications": { + "count": 1, + "ms": 7.153499987907708 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.6069590346887708 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 22.428249940276146 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 25.53637500386685 + }, + "227345276c4b5446": { + "count": 1, + "ms": 117.5829169806093 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.810332945547998 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 2.4559160228818655 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 4.415082978084683 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 2.6487080613151193 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 97.31341700535268, + "thread_submit_ms": 97.31720900163054, + "worker_enter_ms": 98.00841705873609, + "worker_exit_ms": 279.1433749953285, + "await_resume_ms": 279.5155419735238, + "adapter_return_ms": 283.6489169858396 + }, + "vespa_native_seconds": { + "querytime": 0.039, + "summaryfetchtime": 0.043000000000000003, + "searchtime": 0.084 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 486.83012498077005 + }, + { + "sequence": 20, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.2466669799759984 + }, + "dense": { + "count": 1, + "ms": 125.17383298836648 + }, + "sparse": { + "count": 1, + "ms": 3.4318750258535147 + }, + "prepare_query": { + "count": 1, + "ms": 128.62733309157193 + }, + "vespa": { + "count": 1, + "ms": 146.52225002646446 + }, + "visibility": { + "count": 1, + "ms": 21.02175005711615 + }, + "execute": { + "count": 1, + "ms": 168.73904201202095 + }, + "_enrich": { + "count": 1, + "ms": 126.72604201361537 + }, + "_coverage": { + "count": 1, + "ms": 12.922166963107884 + }, + "_final_publications": { + "count": 1, + "ms": 9.215540951117873 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.2911669220775366 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 0.4805000498890877 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 3.204040927812457 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 5.898417090065777 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 115.89466710574925 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 2.005874994210899 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 2.3886249400675297 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 7.568750064820051 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 6.003041984513402 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 131.62004202604294, + "thread_submit_ms": 131.6234590485692, + "worker_enter_ms": 131.65604206733406, + "worker_exit_ms": 274.65837507043034, + "await_resume_ms": 274.9237919924781, + "adapter_return_ms": 278.13983405940235 + }, + "vespa_native_seconds": { + "querytime": 0.034, + "summaryfetchtime": 0.023, + "searchtime": 0.058 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 449.8126250691712 + }, + { + "sequence": 18, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.7867088578641415 + }, + "dense": { + "count": 1, + "ms": 120.1005419716239 + }, + "sparse": { + "count": 1, + "ms": 4.211790976114571 + }, + "prepare_query": { + "count": 1, + "ms": 124.37412491999567 + }, + "vespa": { + "count": 1, + "ms": 257.70354201085865 + }, + "visibility": { + "count": 1, + "ms": 28.95637508481741 + }, + "execute": { + "count": 1, + "ms": 288.38570800144225 + }, + "_enrich": { + "count": 1, + "ms": 22.36049994826317 + }, + "_coverage": { + "count": 1, + "ms": 16.854791902005672 + }, + "_final_publications": { + "count": 1, + "ms": 7.475791964679956 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.691541962325573 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.6582499956712127 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 12.282041017897427 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 2.7418748941272497 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 12.086707982234657 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 0.8839999791234732 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 5.212290911003947 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.2736670430749655 + }, + "4860037db3700175": { + "count": 1, + "ms": 0.16220798715949059 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.21908397320657969 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 0.11995900422334671 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 2.5255829095840454 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 3.300667041912675 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 128.8688329514116, + "thread_submit_ms": 128.89004196040332, + "worker_enter_ms": 128.95845796447247, + "worker_exit_ms": 363.96941693965346, + "await_resume_ms": 374.1924579953775, + "adapter_return_ms": 386.5694999694824 + }, + "vespa_native_seconds": { + "querytime": 0.04, + "summaryfetchtime": 0.06, + "searchtime": 0.10200000000000001 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 464.9904170073569 + }, + { + "sequence": 19, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 3.1694999197497964 + }, + "dense": { + "count": 1, + "ms": 112.03283397480845 + }, + "sparse": { + "count": 1, + "ms": 7.914333022199571 + }, + "prepare_query": { + "count": 1, + "ms": 120.17295800615102 + }, + "vespa": { + "count": 1, + "ms": 219.4754589581862 + }, + "visibility": { + "count": 1, + "ms": 53.484917036257684 + }, + "execute": { + "count": 1, + "ms": 275.14170797076076 + }, + "_enrich": { + "count": 1, + "ms": 78.62829195801169 + }, + "_coverage": { + "count": 1, + "ms": 5.563125014305115 + }, + "_final_publications": { + "count": 1, + "ms": 9.205333073623478 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.083083032630384 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 9.56091599073261 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 23.958457983098924 + }, + "227345276c4b5446": { + "count": 1, + "ms": 67.18949996866286 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 0.3800831036642194 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 0.8849999867379665 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.18833298236131668 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 5.697582964785397 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 123.94775007851422, + "thread_submit_ms": 123.97029204294086, + "worker_enter_ms": 124.3567920755595, + "worker_exit_ms": 318.38562502525747, + "await_resume_ms": 319.638334098272, + "adapter_return_ms": 343.42062508221716 + }, + "vespa_native_seconds": { + "querytime": 0.037, + "summaryfetchtime": 0.06, + "searchtime": 0.098 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 492.43737501092255 + }, + { + "sequence": 17, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 2.748124999925494 + }, + "dense": { + "count": 1, + "ms": 128.40916693676263 + }, + "sparse": { + "count": 1, + "ms": 3.2955839997157454 + }, + "prepare_query": { + "count": 1, + "ms": 131.7740420345217 + }, + "vespa": { + "count": 1, + "ms": 362.80562495812774 + }, + "visibility": { + "count": 1, + "ms": 13.968167011626065 + }, + "execute": { + "count": 1, + "ms": 377.57695792242885 + }, + "_enrich": { + "count": 1, + "ms": 12.1610420756042 + }, + "_coverage": { + "count": 1, + "ms": 5.292040994390845 + }, + "_final_publications": { + "count": 1, + "ms": 4.601667053066194 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 1.487292000092566 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 1.5190830454230309 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 7.20541609916836 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 4.393625073134899 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 0.8197919232770801 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 2.3606669856235385 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.1662500435486436 + }, + "2499200562431a3a": { + "count": 1, + "ms": 1.6536659095436335 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 134.5634589670226, + "thread_submit_ms": 134.57520899828523, + "worker_enter_ms": 135.32129197847098, + "worker_exit_ms": 494.1168340155855, + "await_resume_ms": 495.0119169661775, + "adapter_return_ms": 497.36733397003263 + }, + "vespa_native_seconds": { + "querytime": 0.044, + "summaryfetchtime": 0.126, + "searchtime": 0.17200000000000001 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 534.515458974056 + }, + { + "sequence": 23, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 6.207042024470866 + }, + "dense": { + "count": 1, + "ms": 251.6138330101967 + }, + "sparse": { + "count": 1, + "ms": 5.3990420419722795 + }, + "prepare_query": { + "count": 1, + "ms": 257.07570801023394 + }, + "vespa": { + "count": 1, + "ms": 282.42316702380776 + }, + "visibility": { + "count": 1, + "ms": 67.78208294417709 + }, + "execute": { + "count": 1, + "ms": 351.3190840603784 + }, + "_enrich": { + "count": 1, + "ms": 29.396165977232158 + }, + "_coverage": { + "count": 1, + "ms": 61.066332971677184 + }, + "_final_publications": { + "count": 1, + "ms": 88.22549995966256 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 2.3950840113684535 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 13.396708061918616 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 21.271208999678493 + }, + "06f6532ebc76de6e": { + "count": 1, + "ms": 14.94274998549372 + }, + "e88a191d2772b4c3": { + "count": 1, + "ms": 17.872833996079862 + }, + "7ff2c92fa74aac62": { + "count": 1, + "ms": 2.6501669781282544 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 46.39487492386252 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.1212920071557164 + }, + "5ba1d0d81b5d4a1a": { + "count": 1, + "ms": 74.2000830359757 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 259.72420803736895, + "thread_submit_ms": 259.73125000018626, + "worker_enter_ms": 259.93562501389533, + "worker_exit_ms": 533.6556250695139, + "await_resume_ms": 533.8690420612693, + "adapter_return_ms": 542.1450000721961 + }, + "vespa_native_seconds": { + "querytime": 0.055, + "summaryfetchtime": 0.044, + "searchtime": 0.101 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 794.0028749871999 + }, + { + "sequence": 24, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 10.192250949330628 + }, + "dense": { + "count": 1, + "ms": 243.9227090217173 + }, + "sparse": { + "count": 1, + "ms": 5.815709009766579 + }, + "prepare_query": { + "count": 1, + "ms": 249.76879195310175 + }, + "vespa": { + "count": 1, + "ms": 348.05883292574435 + }, + "visibility": { + "count": 1, + "ms": 49.388541956432164 + }, + "execute": { + "count": 1, + "ms": 399.75908293854445 + }, + "_enrich": { + "count": 1, + "ms": 165.58324999641627 + }, + "_coverage": { + "count": 1, + "ms": 31.614792067557573 + }, + "_final_publications": { + "count": 1, + "ms": 12.574583990499377 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 9.072582935914397 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 11.464042006991804 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 18.014999921433628 + }, + "4ee445d5d52cef6a": { + "count": 1, + "ms": 3.193582990206778 + }, + "d0899a4cbc066a50": { + "count": 1, + "ms": 128.07166704442352 + }, + "d5b9cf8005c6b2d3": { + "count": 1, + "ms": 1.9443329656496644 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 7.557749981060624 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.7783749606460333 + }, + "4860037db3700175": { + "count": 1, + "ms": 2.628499991260469 + }, + "0e5689f21b6e2cb6": { + "count": 1, + "ms": 0.7797080324962735 + }, + "c86dd043af29ddfe": { + "count": 1, + "ms": 0.3227499546483159 + }, + "4759652f7e45b471": { + "count": 2, + "ms": 2.2662499686703086 + }, + "bc2564b15d26cfb8": { + "count": 1, + "ms": 9.172708028927445 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 260.2858330355957, + "thread_submit_ms": 260.3041250258684, + "worker_enter_ms": 260.57608297560364, + "worker_exit_ms": 603.818624978885, + "await_resume_ms": 604.1352079482749, + "adapter_return_ms": 608.341958024539 + }, + "vespa_native_seconds": { + "querytime": 0.064, + "summaryfetchtime": 0.075, + "searchtime": 0.14100000000000001 + }, + "vespa_status": 200, + "candidates": 200, + "total_ms": 870.1690000016242 + }, + { + "sequence": 22, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 4.9316249787807465 + }, + "dense": { + "count": 1, + "ms": 252.40749993827194 + }, + "sparse": { + "count": 1, + "ms": 5.0570000894367695 + }, + "prepare_query": { + "count": 1, + "ms": 257.47762504033744 + }, + "vespa": { + "count": 1, + "ms": 533.1936249276623 + }, + "visibility": { + "count": 1, + "ms": 33.718875027261674 + }, + "execute": { + "count": 1, + "ms": 569.8454580269754 + }, + "_enrich": { + "count": 1, + "ms": 38.22412504814565 + }, + "_coverage": { + "count": 1, + "ms": 12.29758292902261 + }, + "_final_publications": { + "count": 1, + "ms": 9.380374918691814 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 3.5796259762719274 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 4.136416013352573 + }, + "6948fc1e998d8c91": { + "count": 1, + "ms": 23.413957911543548 + }, + "2a3e378facdff592": { + "count": 1, + "ms": 28.427459066733718 + }, + "024351fdb57a17bd": { + "count": 1, + "ms": 1.3487499672919512 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 8.956332923844457 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 0.35566696897149086 + }, + "2499200562431a3a": { + "count": 1, + "ms": 5.193708930164576 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 260.61008300166577, + "thread_submit_ms": 260.61320793814957, + "worker_enter_ms": 261.1957909539342, + "worker_exit_ms": 781.5220409538597, + "await_resume_ms": 789.4895409699529, + "adapter_return_ms": 793.8009579665959 + }, + "vespa_native_seconds": { + "querytime": 0.067, + "summaryfetchtime": 0.14100000000000001, + "searchtime": 0.21 + }, + "vespa_status": 200, + "candidates": 198, + "total_ms": 892.7794579649344 + }, + { + "sequence": 21, + "stages": { + "_resolve_scopes": { + "count": 2, + "ms": 10.100001120008528 + }, + "dense": { + "count": 1, + "ms": 248.2455000281334 + }, + "sparse": { + "count": 1, + "ms": 4.537333967164159 + }, + "prepare_query": { + "count": 1, + "ms": 253.28354199882597 + }, + "vespa": { + "count": 1, + "ms": 332.133625051938 + }, + "visibility": { + "count": 1, + "ms": 27.713124989531934 + }, + "execute": { + "count": 1, + "ms": 362.4508329667151 + }, + "_enrich": { + "count": 1, + "ms": 257.4409160297364 + }, + "_coverage": { + "count": 1, + "ms": 6.178250070661306 + }, + "_final_publications": { + "count": 1, + "ms": 9.463750058785081 + } + }, + "sql": { + "3748346667f87dba": { + "count": 2, + "ms": 7.278832956217229 + }, + "8b73a5bccd0cd7ff": { + "count": 1, + "ms": 5.297499941661954 + }, + "456d4afd8fa45e37": { + "count": 1, + "ms": 14.771750080399215 + }, + "227345276c4b5446": { + "count": 1, + "ms": 243.41858306434005 + }, + "f3f9cbf3fee22c9f": { + "count": 1, + "ms": 7.21791700925678 + }, + "8afce4da4e764b48": { + "count": 1, + "ms": 1.972291967831552 + }, + "bea4ae99e32371a1": { + "count": 1, + "ms": 1.4450419694185257 + }, + "a0f7d10209698549": { + "count": 1, + "ms": 4.8278330359607935 + } + }, + "errors": [], + "vespa_timeline": { + "adapter_enter_ms": 260.712374933064, + "thread_submit_ms": 260.8172500040382, + "worker_enter_ms": 261.0166249796748, + "worker_exit_ms": 582.3783749947324, + "await_resume_ms": 582.7239169739187, + "adapter_return_ms": 592.8405419690534 + }, + "vespa_native_seconds": { + "querytime": 0.062, + "summaryfetchtime": 0.059000000000000004, + "searchtime": 0.122 + }, + "vespa_status": 200, + "candidates": 196, + "total_ms": 899.8248749412596 + } + ] +} From 71fc2f4cbc1e91b286daf6e81801d1736f7c75a7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:07:02 -0700 Subject: [PATCH 163/266] Project only retained search card fields during enrichment --- .../canonical/tests/test_owned_search.py | 50 ++++++++++---- backend/airweave/domains/search/owned.py | 68 +++++++++++++++---- docs/owned-sources/search-execution.md | 35 ++++++++++ 3 files changed, 123 insertions(+), 30 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index cb5e3cbe43..f66108770a 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -7,7 +7,7 @@ import pytest from fastapi import FastAPI from httpx import ASGITransport, AsyncClient -from sqlalchemy import select, update +from sqlalchemy import event, select, update from airweave.api import deps from airweave.api.v1.endpoints.records import router @@ -134,16 +134,36 @@ async def test_owned_http_deduplicates_and_preserves_partial_engine(database, in vector.seed_results( SearchResults(results=[first, second], engine_partial=True, engine_coverage_percent=75) ) - response = await client.post( - "/sync/search", - json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, - ) + transferred_columns = set() + + def columns_received(connection, cursor, statement, parameters, context, executemany): + if cursor.description: + transferred_columns.update(column[0] for column in cursor.description) + + async with database() as db: + engine = db.bind.sync_engine + event.listen(engine, "after_cursor_execute", columns_received) + try: + response = await client.post( + "/sync/search", + json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, + ) + finally: + event.remove(engine, "after_cursor_execute", columns_received) + assert not transferred_columns.intersection({"source_payload", "blob_references"}) assert response.status_code == 200, response.text body = response.json() assert len(body["items"]) == 1 item = body["items"][0] assert item["identity"] == {"record_type": "event", "native_id": "one", "container_id": None} assert item["source_connection_id"] == str(connection.id) + assert item["record_id"] == str(locator.record_id) + assert item["revision"] == locator.revision and item["sync_id"] == str(fence.sync_id) + assert item["provider"] == "gmail" and item["title"] == first.name + assert item["completeness"] == "complete" and item["observed_at"] + assert item["source_created_at"] is None and item["source_updated_at"] is None + assert item["email_thread_id"] is None + assert item["excerpts"] == [first.textual_representation, "Other evidence"] assert len(item["excerpts"]) == 2 and "payload" not in item assert body["engine_partial"] and body["retrieval_incomplete"] assert body["coverage"] == "bounded_candidates" and "next_cursor" not in body @@ -382,17 +402,19 @@ async def ordered(*args): @pytest.mark.parametrize( - "provider,thread_id,expected", + "provider,payload,expected", [ - ("gmail", "thread-1", "thread-1"), - ("slack", "thread-1", None), - ("gmail", None, None), - ("gmail", "bad/id", None), - ("gmail", {"bad": "value"}, None), + ("gmail", {"threadId": "thread-1"}, "thread-1"), + ("slack", {"threadId": "thread-1"}, None), + ("gmail", {}, None), + ("gmail", {"threadId": None}, None), + ("gmail", {"threadId": 12345}, None), + ("gmail", {"threadId": "bad/id"}, None), + ("gmail", {"threadId": {"bad": "value"}}, None), ], ) async def test_email_route_uses_visible_canonical_provider_payload( - database, indexed, http_search, provider, thread_id, expected + database, indexed, http_search, provider, payload, expected ): from airweave.models.entity import Entity @@ -401,9 +423,7 @@ async def test_email_route_uses_visible_canonical_provider_payload( async with database() as db: await db.execute(update(SourceConnection).values(short_name=provider)) await db.execute( - update(Entity).values( - entity_definition_short_name="message", source_payload={"threadId": thread_id} - ) + update(Entity).values(entity_definition_short_name="message", source_payload=payload) ) await db.commit() candidate = hit(fence, locator.encode()) diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 044bc415f3..a155974c5e 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -5,7 +5,8 @@ from uuid import UUID from fastapi import HTTPException -from sqlalchemy import and_, func, or_, select, tuple_ +from pydantic import AwareDatetime, BaseModel, ConfigDict +from sqlalchemy import and_, case, func, or_, select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession from airweave.api.context import ApiContext @@ -36,6 +37,25 @@ from airweave.models.sync import Sync +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 + observed_at: AwareDatetime + source_created_at: AwareDatetime | None + source_updated_at: AwareDatetime | None + completeness: str + email_thread_id: str | None + + class OwnedSearchService: """No provider lifecycle, planner, answer generator or reranker dependency.""" @@ -238,19 +258,37 @@ async def _enrich( hits, scores, exclusions, postfiltered = {}, {}, 0, 0 # Recheck current publication while obtaining native identity. Never enrich # from an unvalidated cached hit after a concurrent capture/permission change. - records = ( - await db.scalars( - select(Entity) - .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), - ) - .execution_options(populate_existing=True) + 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.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"), ) - ).all() + .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( @@ -307,7 +345,7 @@ async def _enrich( thread_id if scopes[row.sync_id].short_name == "gmail" and row.entity_definition_short_name == "message" - and isinstance(thread_id := (row.source_payload or {}).get("threadId"), str) + and (thread_id := row.email_thread_id) is not None and re.fullmatch(r"[A-Za-z0-9_-]{1,512}", thread_id) else None ), @@ -376,7 +414,7 @@ def _locator(result: SearchResult) -> ProjectionLocator | None: return None @staticmethod - def _matches(row: Entity, request: OwnedSearchRequest) -> bool: + 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 ( diff --git a/docs/owned-sources/search-execution.md b/docs/owned-sources/search-execution.md index 954751a5e4..e23e13e968 100644 --- a/docs/owned-sources/search-execution.md +++ b/docs/owned-sources/search-execution.md @@ -144,3 +144,38 @@ nor reconstruct the earlier unlogged 503 and 13-second response. Instrumentation lives only in the private launcher, preserving general log suppression. Its copied query body adds `presentation.timing` only; result selection, ranking, authorization and the 200-candidate bound remain unchanged. + +## Narrow enrichment projection + +Search enrichment now selects only the scalar identity/publication/date/completeness +fields used by a result card, plus the potential Gmail thread ID. A frozen typed row +replaces the live `Entity` ORM object. Native payloads and blob manifests are no longer +transferred or decoded into the search process merely to build cards. PostgreSQL may +still read the JSON value to extract the thread ID; this is not a claim of eliminating +all database JSON work or a measured latency improvement. + +The thread expression returns text only when the native JSON value is a string. +Numbers, objects, missing keys and JSON null remain unavailable. Existing provider, +record-kind and allowed-ID checks still apply. Exact publication predicates, indexed +part coverage, exclusions and the final fresh authorization check are unchanged. + +The phase-session proposal remains separate. Installed FastAPI 0.115.14 does not +provide an early function-scoped yield dependency. A dedicated non-yield owned-search +auth dependency could resolve the same `ContextResolver` inside an owned read context, +close that context, and return the detached `ApiContext`. Generic `get_context` remains +unchanged. The search service could then receive the existing session factory and +own its read contexts, carrying only typed source snapshots between them. Preserve +the current read route's Auth0 last-active rollback semantics explicitly; do not add +an implicit commit or close a caller's borrowed transaction. Qualify a one-connection +pool while embeddings/Vespa are paused, plus revocation and record changes between +phases, before adopting this separate proposal. + +Projection qualification: 20 owned-search/visibility PostgreSQL tests passed against +an isolated archive of `a539e40979871f5ba23428b802845b5f1329f6c9` with only the +search implementation and its test file overlaid. Existing card checks now assert +identity, dates, completeness and excerpt parity and inspect actual DBAPI result +columns to reject transfer of full payloads/blob manifests. Thread-ID cases cover +missing/null/numeric/object/invalid strings and non-Gmail records. Existing stale, +foreign-source, unsupported-part and later-collection mutation fences also passed. +This is scoped correctness/resource-boundary qualification, not a measured speedup +or a claim that the candidate runtime has been updated with this change. From bc262557ee3509cdfac68a42450959f95a3d3a3f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:15:44 -0700 Subject: [PATCH 164/266] Capture bounded native Sheets under existing Drive ownership --- .../auth_provider/providers/composio.py | 4 +- .../entities/canonical/projection_mappers.py | 24 +- .../domains/entities/canonical/query.py | 61 ++++ .../entities/canonical/query_models.py | 22 ++ .../canonical/tests/test_sheets_capture.py | 227 ++++++++++++++ .../entities/canonical/workspace_sheets.py | 147 +++++++++ backend/airweave/platform/configs/config.py | 5 + .../airweave/platform/sources/google_drive.py | 9 +- .../sources/records/google_drive_content.py | 34 ++- .../sources/records/google_sheets_content.py | 285 ++++++++++++++++++ .../sources/records/sheets_manifest.py | 110 +++++++ .../platform/sources/records/sheets_models.py | 176 +++++++++++ .../http_client/test_composio_transport.py | 10 +- .../sources/records/test_sheets_content.py | 249 +++++++++++++++ 14 files changed, 1344 insertions(+), 19 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py create mode 100644 backend/airweave/domains/entities/canonical/workspace_sheets.py create mode 100644 backend/airweave/platform/sources/records/google_sheets_content.py create mode 100644 backend/airweave/platform/sources/records/sheets_manifest.py create mode 100644 backend/airweave/platform/sources/records/sheets_models.py create mode 100644 backend/tests/unit/platform/sources/records/test_sheets_content.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 0ee9a96402..5ef08f1e9d 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -114,7 +114,9 @@ async def get_auth_result( hosts = { "gmail": frozenset({"gmail.googleapis.com"}), "google_calendar": frozenset({"www.googleapis.com"}), - "google_drive": frozenset({"www.googleapis.com", "docs.googleapis.com"}), + "google_drive": frozenset( + {"www.googleapis.com", "docs.googleapis.com", "sheets.googleapis.com"} + ), "slack": frozenset({"slack.com"}), "linear": frozenset({"api.linear.app", "uploads.linear.app"}), "github": frozenset({"api.github.com"}), diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 6034103105..54a6d79823 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -71,10 +71,26 @@ async def _drive( if any(blob.role == "representation_manifest" for blob in record.blobs): from airweave.domains.entities.canonical.workspace_docs import read_document - document = await read_document(record, storage) - content = document.text().encode("utf-8") - media_type = "text/plain" - suffix = ".txt" + 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) + 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: + document = await read_document(record, storage) + content = document.text().encode("utf-8") + media_type = "text/plain" + suffix = ".txt" else: # Historical Drive records retain their export-only contract. if len(record.blobs) != 1: diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py index 72e1eace76..e5c0be5404 100644 --- a/backend/airweave/domains/entities/canonical/query.py +++ b/backend/airweave/domains/entities/canonical/query.py @@ -17,12 +17,15 @@ RecordCursor, RecordListQuery, RecordPage, + SpreadsheetRead, ) from airweave.domains.entities.canonical.query_store import CanonicalQueryStore from airweave.domains.entities.canonical.store import CanonicalRecordStore, CanonicalStoreError from airweave.domains.entities.canonical.workspace_docs import read_document +from airweave.domains.entities.canonical.workspace_sheets import read_spreadsheet from airweave.domains.storage.exceptions import StorageException from airweave.domains.storage.protocols import StorageBackend +from airweave.platform.sources.records.sheets_manifest import GridBounds class InvalidRecordCursor(CanonicalStoreError): @@ -61,6 +64,12 @@ class DocumentUnavailable(CanonicalStoreError): code = "document_unavailable" +class SpreadsheetUnavailable(CanonicalStoreError): + """Native spreadsheet representation or selected sheet cannot be read at this revision.""" + + code = "spreadsheet_unavailable" + + class DocumentIncomplete(CanonicalStoreError): """Retained native document fields do not cover the requested document.""" @@ -268,6 +277,58 @@ async def document( document=captured.document, ) + async def spreadsheet( + self, + db: AsyncSession, + organization_id: UUID, + sync_id: UUID, + record_id: UUID, + revision: int, + storage: StorageBackend, + bounds: GridBounds | None = None, + ) -> SpreadsheetRead: + """Read a revision-pinned native grid; provider access is never a fallback.""" + if bounds is not None and ( + (bounds.end_row - bounds.start_row) * (bounds.end_column - bounds.start_column) > 10000 + ): + raise SpreadsheetUnavailable("Requested grid exceeds the 10000-cell read limit") + record = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(record, revision) + try: + captured = await read_spreadsheet(record, storage, bounds=bounds) + except (StorageException, BlobIntegrityError) as exc: + raise BlobUnavailable( + "Committed spreadsheet bytes are unavailable; retry later" + ) from exc + except ValueError as exc: + raise SpreadsheetUnavailable( + "No valid native spreadsheet is retained; refresh and reread" + ) from exc + try: + cells = captured.cells(bounds) if bounds is not None else () + except ValueError as exc: + raise SpreadsheetUnavailable( + "Requested grid is outside retained sheet dimensions" + ) from exc + db.expire_all() + current = await self.read(db, organization_id, sync_id, record_id) + self._check_blob_record(current, revision) + if record.blobs != current.blobs: + raise StaleRecordRevision("Spreadsheet changed; reread before requesting content") + return SpreadsheetRead( + id=current.id, + identity=current.identity, + revision=current.revision, + observed_at=current.observed_at, + completeness=current.completeness, + spreadsheet=captured.spreadsheet, + grid_status=captured.manifest.native.status, + captured=tuple(p.bounds for p in captured.manifest.native.parts), + missing=captured.manifest.native.missing, + requested=bounds, + cells=tuple(cells), + ) + async def blob( self, db: AsyncSession, diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index fa21dc5498..c3809eefb6 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -10,6 +10,8 @@ from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.platform.sources.records.sheets_manifest import GridBounds, GridGap +from airweave.platform.sources.records.sheets_models import SpreadsheetCell from airweave.platform.sources.records.workspace_manifest import WorkspaceManifestV1 @@ -32,6 +34,26 @@ class DocumentRead(BaseModel): document: dict[str, JsonValue] +class SpreadsheetRead(BaseModel): + """Stored workbook metadata and bounded native cells; unknown ranges stay unknown.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + id: UUID + identity: RecordIdentity + revision: int = Field(ge=1) + observed_at: AwareDatetime + completeness: Literal["complete", "partial", "metadata_only"] + spreadsheet: dict[str, JsonValue] + grid_status: Literal["complete", "partial"] + captured: tuple[GridBounds, ...] + missing: tuple[GridGap, ...] + requested: GridBounds | None = None + cells: tuple[SpreadsheetCell, ...] = () + comments: Literal["omitted"] = "omitted" + embedded_media: Literal["not_retained"] = "not_retained" + calculated_values: Literal["observed_at_capture"] = "observed_at_capture" + + class RecordFilters(BaseModel): """Supported exact filters, preserved in every continuation token.""" 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..d53169f1de --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py @@ -0,0 +1,227 @@ +"""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") + 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 + fallback = await _drive(newer, files.storage, tmp_path / "fallback") + 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/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/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 56b997ebc8..6eb8a76b8f 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -389,6 +389,11 @@ class GoogleDriveConfig(SourceConfig): 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", diff --git a/backend/airweave/platform/sources/google_drive.py b/backend/airweave/platform/sources/google_drive.py index 7a6b9d5792..0f0783dd14 100644 --- a/backend/airweave/platform/sources/google_drive.py +++ b/backend/airweave/platform/sources/google_drive.py @@ -113,6 +113,7 @@ class GoogleDriveSource(BaseSource): 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: @@ -135,7 +136,11 @@ def capture_cycle_configuration(self) -> CycleConfiguration: { "capture_version": 3, "permission_id": self.expected_permission_id, - "scope": "drive-all-accessible-docs-native", + "scope": ( + "drive-all-accessible-workspace-native-v2" + if self.capture_native_sheets + else "drive-all-accessible-docs-native" + ), }, sort_keys=True, ).encode() @@ -163,6 +168,7 @@ async def _capture_body(self, record: CaptureRecord, files: FileService) -> Capt client=self.http_client, auth=self.auth, logger=self.logger, + capture_native_sheets=self.capture_native_sheets, ) async def capture_page( @@ -213,6 +219,7 @@ async def create( 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 diff --git a/backend/airweave/platform/sources/records/google_drive_content.py b/backend/airweave/platform/sources/records/google_drive_content.py index 5d0feaa041..876d29cfc4 100644 --- a/backend/airweave/platform/sources/records/google_drive_content.py +++ b/backend/airweave/platform/sources/records/google_drive_content.py @@ -13,6 +13,8 @@ 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 @@ -24,6 +26,7 @@ async def capture_file_content( 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": @@ -37,18 +40,14 @@ async def capture_file_content( if isinstance(capabilities, dict) and capabilities.get("canDownload") is False: return record url = f"{BASE}/files/{quote(record.identity.native_id, safe='')}" - if mime.startswith("application/vnd.google-apps."): - export = GOOGLE_EXPORT_FORMATS.get(mime) - if export is None: - return record - media_type = export[0] - download_url = url + "/export?" + urlencode({"mimeType": media_type}) - else: - media_type = mime - download_url = url + "?alt=media&supportsAllDrives=true" + 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 ) + has_native_parts = False if mime == DOCS_MIME: record = await capture_document_parts( record, @@ -58,6 +57,12 @@ async def capture_file_content( 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"}) @@ -65,13 +70,22 @@ async def capture_file_content( raise ValueError( "Drive file changed during content capture; retry before advancing checkpoint" ) - if mime == DOCS_MIME: + if has_native_parts: return record return record.model_copy( update={"blobs": (blob,), "content_hash": blob.sha256, "completeness": "complete"} ) +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, 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/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py index 7d92552965..e12c618deb 100644 --- a/backend/tests/unit/platform/http_client/test_composio_transport.py +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -311,7 +311,7 @@ async def aclose(self): @pytest.mark.asyncio -async def test_drive_binding_allows_native_docs_only_with_same_verified_toolkit(monkeypatch): +async def test_drive_binding_allows_native_workspace_only_with_same_verified_toolkit(monkeypatch): from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider provider = await ComposioAuthProvider.create( @@ -330,8 +330,12 @@ async def test_drive_binding_allows_native_docs_only_with_same_verified_toolkit( ) result = await provider.get_auth_result("google_drive", ["access_token"]) assert result.managed_auth.connected_account_id == "ca_drive" - assert result.managed_auth.allowed_hosts == {"www.googleapis.com", "docs.googleapis.com"} - assert "sheets.googleapis.com" not in result.managed_auth.allowed_hosts + assert result.managed_auth.allowed_hosts == { + "www.googleapis.com", + "docs.googleapis.com", + "sheets.googleapis.com", + } + assert "untrusted.example" not in result.managed_auth.allowed_hosts @pytest.mark.asyncio diff --git a/backend/tests/unit/platform/sources/records/test_sheets_content.py b/backend/tests/unit/platform/sources/records/test_sheets_content.py new file mode 100644 index 0000000000..09502f7393 --- /dev/null +++ b/backend/tests/unit/platform/sources/records/test_sheets_content.py @@ -0,0 +1,249 @@ +"""Native structured capture preserves content and truthful bounded coverage.""" + +from copy import deepcopy +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest + +from airweave.adapters.storage.filesystem import FilesystemBackend +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 +from airweave.platform.sources.records.sheets_models import validate_grid + + +def metadata(rows=2, columns=2): + return { + "spreadsheetId": "book", + "properties": {"locale": "en_US", "timeZone": "UTC"}, + "sheets": [ + { + "properties": { + "sheetId": 0, + "title": "A'B", + "hidden": True, + "gridProperties": {"rowCount": rows, "columnCount": columns}, + }, + "merges": [ + { + "sheetId": 0, + "startRowIndex": 0, + "endRowIndex": 1, + "startColumnIndex": 0, + "endColumnIndex": 2, + } + ], + } + ], + } + + +def grid(): + value = metadata() + value["sheets"][0]["data"] = [ + { + "rowData": [ + { + "values": [ + { + "userEnteredValue": {"formulaValue": "=1+1"}, + "effectiveValue": {"numberValue": 2}, + "formattedValue": "2", + "note": "native note", + "effectiveFormat": {"numberFormat": {"type": "NUMBER"}}, + } + ] + } + ] + } + ] + return value + + +async def acquire(tmp_path, monkeypatch, *, partial=False, wrong=False): + monkeypatch.setattr( + "airweave.domains.storage.file_service.paths.temp_sync_dir", + lambda _: str(tmp_path / "temp"), + ) + if partial: + monkeypatch.setattr( + "airweave.platform.sources.records.google_sheets_content.MAX_GRID_REQUESTS", 0 + ) + storage = FilesystemBackend(tmp_path / "store") + files = FileService(uuid4(), storage, sync_id=uuid4()) + calls = [] + + def handler(request): + calls.append(request) + if request.url.path.endswith("/export"): + return httpx.Response(200, content=b"full-xlsx") + assert request.method == "GET" + if "ranges" in request.url.params: + assert request.url.params["ranges"] == "'A''B'!A1:B2" + result = grid() + if wrong: + result["sheets"][0]["properties"]["sheetId"] = 9 + return httpx.Response(200, json=result) + return httpx.Response(200, json=metadata()) + + async def latest(*args, **kwargs): + return {"version": "7"} + + original = file_record({"id": "book", "version": "7", "mimeType": SHEETS_MIME}) + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + captured = await capture_file_content( + original, + files=files, + get=latest, + client=client, + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + capture_native_sheets=True, + ) + return captured, storage, calls + + +@pytest.mark.asyncio +async def test_native_formula_values_and_sparse_grid_retained(tmp_path, monkeypatch): + captured, storage, calls = await acquire(tmp_path, monkeypatch) + assert captured.completeness == "complete" + assert len(captured.blobs) == 4 + assert len(calls) == 3 + + +@pytest.mark.asyncio +async def test_budget_preserves_export_and_explicit_unknown_ranges(tmp_path, monkeypatch): + captured, storage, calls = await acquire(tmp_path, monkeypatch, partial=True) + assert captured.completeness == "partial" + assert len(captured.blobs) == 3 + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_wrong_native_sheet_rejects_capture(tmp_path, monkeypatch): + with pytest.raises(ValueError, match="different sheets"): + await acquire(tmp_path, monkeypatch, wrong=True) + + +def test_sparse_grid_cannot_escape_requested_bounds(): + bounds = GridBounds(sheet_id=0, start_row=0, end_row=2, start_column=0, end_column=2) + validate_grid(grid(), "book", bounds) + bad = deepcopy(grid()) + bad["sheets"][0]["data"][0]["startRow"] = 2 + with pytest.raises(ValueError, match="outside"): + validate_grid(bad, "book", bounds) + + +@pytest.mark.asyncio +async def test_disabled_native_capture_keeps_export_only_behavior(tmp_path, monkeypatch): + monkeypatch.setattr( + "airweave.domains.storage.file_service.paths.temp_sync_dir", + lambda _: str(tmp_path / "temp"), + ) + files = FileService(uuid4(), FilesystemBackend(tmp_path / "store"), sync_id=uuid4()) + + async def latest(*args, **kwargs): + return {"version": "7"} + + def response(request): + assert request.url.host == "www.googleapis.com" + assert request.url.path.endswith("/export") + return httpx.Response(200, content=b"legacy-xlsx") + + async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client: + captured = await capture_file_content( + file_record({"id": "book", "version": "7", "mimeType": SHEETS_MIME}), + files=files, + get=latest, + client=client, + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + ) + assert captured.completeness == "complete" + assert len(captured.blobs) == 1 and captured.blobs[0].role is None + + +def test_overlapping_native_cells_are_rejected(): + value = grid() + value["sheets"][0]["data"].append(deepcopy(value["sheets"][0]["data"][0])) + with pytest.raises(ValueError, match="overlapping cells"): + validate_grid( + value, + "book", + GridBounds(sheet_id=0, start_row=0, end_row=2, start_column=0, end_column=2), + ) + + +def test_native_search_groups_sheet_titles_and_preserves_sparse_row_addresses(): + from airweave.domains.entities.canonical.workspace_sheets import ( + CapturedGrid, + CapturedSpreadsheet, + ) + from airweave.platform.sources.records.sheets_manifest import ( + GridPart, + SheetsState, + WorkspaceManifestV2, + ) + from airweave.platform.sources.records.workspace_manifest import ExportState + + parts = tuple( + GridPart( + bounds=GridBounds( + sheet_id=sheet, start_row=0, end_row=100, start_column=0, end_column=1 + ), + blob=str(sheet + 1) * 64, + ) + for sheet in (0, 1) + ) + sheets, grids = [], [] + for index, title in enumerate(("First sheet", "Second sheet")): + properties = { + "sheetId": index, + "title": title, + "gridProperties": {"rowCount": 100, "columnCount": 1}, + } + sheets.append({"properties": properties}) + grids.append( + CapturedGrid( + part=parts[index], + data={ + "spreadsheetId": "book", + "sheets": [ + { + "properties": properties, + "data": [ + { + "startRow": 48, + "rowData": [ + {"values": [{"formattedValue": f"unique-value-{index}"}]} + ], + } + ], + } + ], + }, + ) + ) + captured = CapturedSpreadsheet( + manifest=WorkspaceManifestV2( + file_id="book", + drive_version="7", + export=ExportState(status="unavailable", reason="export_size_limit"), + native=SheetsState(status="complete", metadata_blob="a" * 64, parts=parts), + ), + spreadsheet={"spreadsheetId": "book", "sheets": sheets}, + grids=tuple(grids), + ) + text = captured.text() + assert ( + text.index("First sheet") + < text.index("unique-value-0") + < text.index("Second sheet") + < text.index("unique-value-1") + ) + assert "A49: unique-value-0" in text + assert "'First sheet'!A1:A100" in text From 30a88f9667b4e67aaf127e8da5abba151b9c3271 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:17:19 -0700 Subject: [PATCH 165/266] Prove blank Sheets ranges preserve coverage with shared bytes --- .../sources/records/test_sheets_content.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/backend/tests/unit/platform/sources/records/test_sheets_content.py b/backend/tests/unit/platform/sources/records/test_sheets_content.py index 09502f7393..29e81b77fd 100644 --- a/backend/tests/unit/platform/sources/records/test_sheets_content.py +++ b/backend/tests/unit/platform/sources/records/test_sheets_content.py @@ -247,3 +247,50 @@ def test_native_search_groups_sheet_titles_and_preserves_sparse_row_addresses(): ) assert "A49: unique-value-0" in text assert "'First sheet'!A1:A100" in text + + +@pytest.mark.asyncio +async def test_distinct_blank_ranges_share_bytes_without_losing_coverage(tmp_path, monkeypatch): + from airweave.platform.sources.records.sheets_manifest import parse_sheet_manifest + from airweave.platform.sources.records.sheets_models import validate_partition + + monkeypatch.setattr( + "airweave.domains.storage.file_service.paths.temp_sync_dir", lambda _: str(tmp_path / "tmp") + ) + files = FileService(uuid4(), FilesystemBackend(tmp_path / "store"), sync_id=uuid4()) + empty = metadata(rows=200, columns=1) + requested = [] + + async def latest(*args, **kwargs): + return {"version": "7"} + + def response(request): + if request.url.path.endswith("/export"): + return httpx.Response(200, content=b"empty-workbook-export") + if "ranges" in request.url.params: + requested.append(request.url.params["ranges"]) + return httpx.Response(200, json=empty) + + async with httpx.AsyncClient(transport=httpx.MockTransport(response)) as client: + captured = await capture_file_content( + file_record({"id": "book", "version": "7", "mimeType": SHEETS_MIME}), + files=files, + get=latest, + client=client, + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + capture_native_sheets=True, + ) + assert requested == ["'A''B'!A1:A100", "'A''B'!A101:A200"] + assert len(captured.blobs) == 3 + marked = next(blob for blob in captured.blobs if blob.role == "representation_manifest") + manifest = parse_sheet_manifest( + await files.storage.read_file(marked.key), + file_id="book", + drive_version="7", + blobs=captured.blobs, + ) + assert len(manifest.native.parts) == 2 + assert {part.blob for part in manifest.native.parts} == {manifest.native.metadata_blob} + assert manifest.native.parts[0].bounds != manifest.native.parts[1].bounds + validate_partition(manifest, empty) From a32b61db18c034d5ec60aab53245a4f445d24b2a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:18:59 -0700 Subject: [PATCH 166/266] Release search database connections across remote waits --- backend/airweave/api/deps.py | 35 ++- .../canonical/tests/test_owned_search.py | 2 + .../tests/test_real_vespa_projection.py | 2 + .../canonical/tests/test_search_sessions.py | 265 ++++++++++++++++++ backend/airweave/domains/search/owned.py | 178 +++++++----- docs/owned-sources/search-execution.md | 79 +++--- 6 files changed, 456 insertions(+), 105 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_search_sessions.py diff --git a/backend/airweave/api/deps.py b/backend/airweave/api/deps.py index 92211f7edf..de395e8759 100644 --- a/backend/airweave/api/deps.py +++ b/backend/airweave/api/deps.py @@ -9,7 +9,7 @@ from fastapi import Depends, Header, HTTPException, Request from fastapi_auth0 import Auth0User -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from airweave import crud, schemas from airweave.api.auth import auth0 @@ -23,6 +23,7 @@ from airweave.core.protocols.cache import ContextCache from airweave.core.protocols.rate_limiter import RateLimiter from airweave.core.shared_models import AuthMethod +from airweave.db import session as db_session from airweave.db.session import get_db from airweave.domains.entities.canonical.query import CanonicalQueryService from airweave.domains.entities.canonical.query_store import CanonicalQueryStore @@ -109,6 +110,38 @@ async def get_context( return await resolver.resolve(request, db, auth0_user, x_api_key, x_organization_id) +def get_search_session_factory() -> async_sessionmaker[AsyncSession]: + """Resolve the configured factory dynamically, including isolated local deployments.""" + return db_session.AsyncSessionLocal + + +async def get_owned_search_context( + request: Request, + sessions: async_sessionmaker[AsyncSession] = Depends(get_search_session_factory), + x_api_key: Optional[str] = Header(None, alias="X-API-Key"), + x_organization_id: Optional[str] = Header(None, alias="X-Organization-ID"), + auth0_user: Optional[Auth0User] = Depends(auth0.get_user), + cache: ContextCache = Inject(ContextCache), + rate_limiter: RateLimiter = Inject(RateLimiter), +) -> ApiContext: + """Finish this read route's authentication transaction before search network waits. + + Reuse ordinary context resolution. Closing our own session preserves the previous + read route's rollback semantics, including Auth0's uncommitted last-active update. + Other routes retain their existing request-owned session. + """ + async with sessions() as db: + return await get_context( + request=request, + db=db, + x_api_key=x_api_key, + x_organization_id=x_organization_id, + auth0_user=auth0_user, + cache=cache, + rate_limiter=rate_limiter, + ) + + async def get_logger( context: ApiContext = Depends(get_context), ) -> ContextualLogger: diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index f66108770a..bcfc86300c 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -119,6 +119,8 @@ async def session(): ) 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=service) async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: yield client, vector, ctx, executor, dense 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 index 0e177a5d3e..0e4b1d29c9 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py @@ -137,6 +137,8 @@ async def session(): 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: 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..8894006e02 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py @@ -0,0 +1,265 @@ +"""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"}, + ) + ) + await asyncio.wait_for(entered.wait(), timeout=5) + if outcome == "cancellation": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(RuntimeError, match="synthetic enrichment failure"): + await task + 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/search/owned.py b/backend/airweave/domains/search/owned.py index a155974c5e..479d37dcd4 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -7,7 +7,7 @@ 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 +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from airweave.api.context import ApiContext from airweave.domains.entities.canonical.coverage import capture_coverage @@ -37,6 +37,18 @@ 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.""" @@ -65,26 +77,25 @@ def __init__(self, executor: SearchPlanExecutorProtocol, registry: SourceRegistr self._registry = registry async def search( - self, db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest + self, + sessions: async_sessionmaker[AsyncSession], + ctx: ApiContext, + request: OwnedSearchRequest, ) -> OwnedSearchResponse: """Resolve exact authorized scopes before any embedding/index request.""" - scopes, groups = await self._resolve_scopes(db, ctx, request) - if self._filtered(request): - versions = ( - await db.scalars( - select(Sync.index_pipeline_version).where( - Sync.organization_id == ctx.organization.id, Sync.id.in_(request.sync_ids) - ) - ) - ).all() - if any(version < SEARCH_METADATA_PIPELINE_VERSION for version in versions): - raise HTTPException( - 409, - { - "code": "reindex_required", - "message": "Selected sources need canonical search metadata re-projection", - }, - ) + 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, exclusions, postfiltered = {}, {}, 0, 0 engine_partial, full = False, False @@ -96,31 +107,35 @@ async def search( ) prepared_query = await self._executor.prepare_query(plan) if groups else None for (collection_id, readable_id), sync_ids in groups.items(): - 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, - ) + # 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, 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 - group_hits, group_scores, rejected, filtered = await self._enrich( - db, ctx, request, sync_ids, scopes, results - ) hits.update(group_hits) scores.update(group_scores) exclusions += rejected postfiltered += filtered - 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") - eligible = await self._final_publications(db, ctx, request, scores, scope_snapshot) + 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") + eligible = await self._final_publications(db, ctx, request, scores, scope_snapshot) exclusions += len(hits.keys() - eligible) ranked = sorted(hits.keys() & eligible, key=lambda key: (-scores[key][0], str(key))) return OwnedSearchResponse( @@ -196,40 +211,51 @@ def _prefilters(request: OwnedSearchRequest, sync_ids: list[UUID]) -> list[Filte async def _resolve_scopes( self, db: AsyncSession, ctx: ApiContext, request: OwnedSearchRequest - ) -> tuple[dict[UUID, SourceConnection], dict[tuple[UUID, str], list[UUID]]]: + ) -> tuple[dict[UUID, _SourceScope], dict[tuple[UUID, str], list[UUID]]]: rows = ( - await db.execute( - select(SourceConnection, Collection.id) - .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), + ( + 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), + ) ) - .execution_options(populate_existing=True) ) - ).all() - if len(rows) != len(request.sync_ids) or {row[0].sync_id for row in rows} != set( + .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[0].sync_id: row[0] for row in rows} + scopes = {row.sync_id: row for row in snapshots} allowed_types = set() groups = defaultdict(list) - for connection, collection_id in rows: + for connection in snapshots: types = getattr( self._registry.get(connection.short_name).source_class_ref, "canonical_record_types", @@ -238,7 +264,9 @@ async def _resolve_scopes( if not types: raise HTTPException(422, "Source does not support owned indexed records") allowed_types.update(types) - groups[(collection_id, connection.readable_collection_id)].append(connection.sync_id) + 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 @@ -249,7 +277,7 @@ async def _enrich( ctx: ApiContext, request: OwnedSearchRequest, sync_ids: list[UUID], - scopes: dict[UUID, SourceConnection], + scopes: dict[UUID, _SourceScope], results: SearchResults, ) -> tuple[dict[UUID, OwnedSearchHit], dict[UUID, tuple[float, ProjectionLocator]], int, int]: locators = [value for result in results.results if (value := self._locator(result))] @@ -358,9 +386,16 @@ async def _enrich( return hits, scores, exclusions, postfiltered @staticmethod - def _scope_identity(scopes: dict[UUID, SourceConnection]) -> set[tuple]: + def _scope_identity(scopes: dict[UUID, _SourceScope]) -> set[tuple]: return { - (sync, row.id, row.short_name, row.readable_collection_id) + ( + sync, + row.id, + row.short_name, + row.readable_collection_id, + row.collection_id, + row.index_pipeline_version, + ) for sync, row in scopes.items() } @@ -391,6 +426,13 @@ async def _final_publications( 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, @@ -400,6 +442,8 @@ async def _final_publications( 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()), ) diff --git a/docs/owned-sources/search-execution.md b/docs/owned-sources/search-execution.md index e23e13e968..2a542084a7 100644 --- a/docs/owned-sources/search-execution.md +++ b/docs/owned-sources/search-execution.md @@ -46,10 +46,10 @@ running without OOM/restart. No cause-specific retry or timeout increase is just by that evidence. Safe status, exception type and elapsed-stage diagnostics are being qualified separately in the Almanac adapter. -The current request-owned database session also holds its connection across -embedding/index waits after the first SQL query. Separating read phases is a -follow-up requiring explicit auth/session ownership and retained final fences; -it is not included in embedding reuse and is not a proven cause of the 503. +The runtime measured here still used a request-owned database session that held +its connection across embedding/index waits after the first SQL query. The later +phase-owned implementation below addresses that ownership separately; this was +not established as the cause of the 503. ## Repeated actual API qualification after embedding reuse @@ -75,28 +75,11 @@ queueing, HTTP, native engine work and event-loop scheduling. It is **not** evid of a PostgreSQL plan switch. Sample 3 independently had an 831 ms SQL fingerprint interval. Neither measurement reconstructs the previous unlogged 503/13-second tail. -## Proposed next diagnostics and session ownership (not implemented) - -The current Vespa adapter invokes the synchronous SDK through `asyncio.to_thread`. -Measure worker start/end alongside the adapter await, and retain only numeric native -response timing fields if available. This separates thread queue delay from SDK wall -time without logging YQL, embeddings, results or identities. Do not change query -ranking, candidate counts, retries or database plans before attributing the tail. - -Database phase separation needs an explicit owner. The search route and its auth -dependency currently share FastAPI's cached `get_db` session. API context contains -Pydantic snapshots; Auth0 resolution nevertheless flushes a last-active update into -that session without committing it. Closing an arbitrary borrowed session inside -the executor would be an unsafe general contract. - -A bounded redesign would finish this route's authentication session explicitly, -then use existing session-factory read contexts: scopes/version to a typed detached -snapshot; embedding preparation without SQL; each collection's network retrieval -before its session performs visibility/enrichment SQL; finally fresh scopes, -coverage and exact publication checks in one read phase. Generic executor behavior -and other routes stay unchanged. Test auth side-effect semantics and revocation -between phases, including a constrained pool, before claiming reduced occupancy. -No extra service, table, result cache or concurrent `AsyncSession` use is proposed. +## Diagnostic design + +The private instrumentation below measures worker start/end alongside the Vespa +adapter await and retains only numeric native response timing fields. It does not +log YQL, embeddings, results or identities or change ranking and candidate limits. The installed PyVespa `Vespa.query` creates a fresh `VespaSync` HTTP session per call, then retains the full decoded response in `VespaQueryResponse.json`. This is @@ -159,17 +142,6 @@ Numbers, objects, missing keys and JSON null remain unavailable. Existing provid record-kind and allowed-ID checks still apply. Exact publication predicates, indexed part coverage, exclusions and the final fresh authorization check are unchanged. -The phase-session proposal remains separate. Installed FastAPI 0.115.14 does not -provide an early function-scoped yield dependency. A dedicated non-yield owned-search -auth dependency could resolve the same `ContextResolver` inside an owned read context, -close that context, and return the detached `ApiContext`. Generic `get_context` remains -unchanged. The search service could then receive the existing session factory and -own its read contexts, carrying only typed source snapshots between them. Preserve -the current read route's Auth0 last-active rollback semantics explicitly; do not add -an implicit commit or close a caller's borrowed transaction. Qualify a one-connection -pool while embeddings/Vespa are paused, plus revocation and record changes between -phases, before adopting this separate proposal. - Projection qualification: 20 owned-search/visibility PostgreSQL tests passed against an isolated archive of `a539e40979871f5ba23428b802845b5f1329f6c9` with only the search implementation and its test file overlaid. Existing card checks now assert @@ -179,3 +151,36 @@ missing/null/numeric/object/invalid strings and non-Gmail records. Existing stal foreign-source, unsupported-part and later-collection mutation fences also passed. This is scoped correctness/resource-boundary qualification, not a measured speedup or a claim that the candidate runtime has been updated with this change. + + +## Phase-owned search reads + +Owned search now closes its own authentication and initial source-scope read +transactions before embedding preparation. Each collection gets a fresh session: +it performs no SQL until indexed retrieval returns, then validates publication and +builds cards. That session closes before the next collection. A final fresh read +rechecks authorized sources, coverage and every exact publication. Only frozen typed +source snapshots cross phases. Source collection UUID and pipeline version are part +of both the scope comparison and final SQL fence, alongside source/account identity. +A same-name collection replacement or pipeline change cannot reuse the old routing. + +The dedicated non-yield dependency calls the existing context resolver. API-key +checks, Auth0 authentication, rate limiting and caching are unchanged. Closing its +owned session preserves the existing read route's rollback of Auth0's flushed +last-active update; this change does not silently commit it. Generic auth dependencies +and other endpoints remain unchanged. No borrowed session is rolled back or closed. + +The session factory is resolved dynamically from the configured database module. +Local evaluation must keep its isolated factory installed (or explicitly override +the new dependency); overriding only `get_db` is insufficient. Test fixtures override +both new dependencies. Runtime activation requires checking this boundary again. + +Qualification uses real PostgreSQL with pool size one and no overflow: the sole +connection is available while sparse inference or Vespa is paused, after cancellation +there, and after failure/cancellation once enrichment has acquired the connection. +Actual API-key/Auth0 SQL authentication also releases that pool and leaves persisted +last-active unchanged. Scope revocation, collection replacement and pipeline changes +during retrieval reject the response; existing later-collection record mutation +checks still prevent returning an earlier stale hit. Search/embedding transports are +fakes in these tests. This is a pool-ownership/correctness proof, not a measured +latency gain or a claim that the local browser runtime has been activated. From 3af25e4cb3e1345693c018ffb21ec38c1acacfe0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:19:15 -0700 Subject: [PATCH 167/266] Wire retained spreadsheet reads and phase-owned search route --- backend/airweave/api/v1/endpoints/records.py | 53 ++++++++++++++++++-- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index 8c088911ad..b21a61853b 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -6,7 +6,7 @@ from fastapi import Depends, HTTPException, Path, Query, Request, Response from fastapi.responses import JSONResponse -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from airweave.api import deps from airweave.api.context import ApiContext @@ -24,10 +24,12 @@ RecordFilters, RecordListQuery, RecordPage, + SpreadsheetRead, ) from airweave.domains.entities.canonical.store import CanonicalStoreError from airweave.domains.entities.canonical.text_models import TextRead, TextRepresentationList from airweave.domains.search.owned_models import OwnedSearchRequest, OwnedSearchResponse +from airweave.platform.sources.records.sheets_manifest import GridBounds router = TrailingSlashRouter() @@ -42,12 +44,12 @@ async def _search_disconnected(request: Request) -> None: async def search_records( request: OwnedSearchRequest, http_request: Request, - db: AsyncSession = Depends(get_db), - ctx: ApiContext = Depends(deps.get_context), + sessions: async_sessionmaker[AsyncSession] = Depends(deps.get_search_session_factory), + ctx: ApiContext = Depends(deps.get_owned_search_context), container: Container = Depends(deps.get_container), ) -> OwnedSearchResponse: """Retrieve bounded indexed originals; no provider requests or agent execution.""" - work = create_task(container.owned_search.search(db, ctx, request)) + work = create_task(container.owned_search.search(sessions, ctx, request)) disconnected = create_task(_search_disconnected(http_request)) try: done, _ = await wait((work, disconnected), return_when=FIRST_COMPLETED) @@ -72,6 +74,7 @@ async def record_error_response(request: Request, error: CanonicalStoreError) -> "stale_record_revision": 409, "blob_unavailable": 503, "document_unavailable": 409, + "spreadsheet_unavailable": 409, "text_unavailable": 409, "document_incomplete": 409, "calendar_changed_restart": 409, @@ -134,6 +137,48 @@ async def read_record( return IndexedRecordRead(**record.model_dump(), extraction=extraction) +@router.get("/{sync_id}/records/{record_id}/spreadsheet", response_model=SpreadsheetRead) +async def read_stored_spreadsheet( + sync_id: UUID, + record_id: UUID, + response: Response, + revision: int = Query(ge=1), + sheet_id: int | None = Query(None, ge=0), + start_row: int | None = Query(None, ge=0), + end_row: int | None = Query(None, gt=0), + start_column: int | None = Query(None, ge=0), + end_column: int | None = Query(None, gt=0), + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(deps.get_context), + service: CanonicalQueryService = Depends(deps.get_canonical_query_service), + container: Container = Depends(deps.get_container), +) -> SpreadsheetRead: + """Read exact retained spreadsheet cells; no live provider fallback.""" + values = (sheet_id, start_row, end_row, start_column, end_column) + if any(value is not None for value in values) and any(value is None for value in values): + raise HTTPException(status_code=422, detail="Supply all grid bounds or none") + try: + bounds = ( + None + if sheet_id is None + else GridBounds( + sheet_id=sheet_id, + start_row=start_row, + end_row=end_row, + start_column=start_column, + end_column=end_column, + ) + ) + except ValueError: + raise HTTPException(status_code=422, detail="Grid bounds must be nonempty") from None + result = await service.spreadsheet( + db, ctx.organization.id, sync_id, record_id, revision, container.storage_backend, bounds + ) + response.headers["Cache-Control"] = "private, no-store" + response.headers["X-Content-Type-Options"] = "nosniff" + return result + + @router.get("/{sync_id}/records/{record_id}/document", response_model=DocumentRead) async def read_stored_document( sync_id: UUID, From 0bb78ebc553977f4d7604c7dafee0df205a684a0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 17:58:44 -0700 Subject: [PATCH 168/266] Compare repeated capture evidence without exposing originals --- backend/tests/live/capture_comparison.py | 114 ++++++++++++++++++ backend/tests/live/provider_lifecycle.py | 13 +- backend/tests/unit/test_capture_comparison.py | 69 +++++++++++ 3 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 backend/tests/live/capture_comparison.py create mode 100644 backend/tests/unit/test_capture_comparison.py diff --git a/backend/tests/live/capture_comparison.py b/backend/tests/live/capture_comparison.py new file mode 100644 index 0000000000..f5bb3ee998 --- /dev/null +++ b/backend/tests/live/capture_comparison.py @@ -0,0 +1,114 @@ +"""Aggregate retained observations before disposable trials remove private evidence.""" + +from collections import Counter +from collections.abc import Mapping +from uuid import UUID + +from pydantic import JsonValue + +from airweave.domains.entities.canonical.models import SourceRecord + +FieldPath = tuple[str | int, ...] + + +def _changed_paths(before: JsonValue, after: JsonValue, path: FieldPath = ()) -> list[FieldPath]: + if type(before) is not type(after): + return [path] + if isinstance(before, dict) and isinstance(after, dict): + paths = [] + for key in before.keys() | after.keys(): + if key not in before or key not in after: + paths.append((*path, key)) + else: + paths.extend(_changed_paths(before[key], after[key], (*path, key))) + return paths + if isinstance(before, list) and isinstance(after, list): + if len(before) != len(after): + return [path] + return [ + changed + for index, (left, right) in enumerate(zip(before, after, strict=True)) + for changed in _changed_paths(left, right, (*path, index)) + ] + return [] if before == after else [path] + + +def _category(provider: str, path: FieldPath) -> str: + # Never emit a native key/path: providers may put private text in JSON keys. + if provider == "gmail": + if path and path[0] == "labelIds": + return "gmail_labels" + if path == ("historyId",): + return "gmail_history_id" + if path[:1] == ("payload",) and path[-2:] == ("body", "attachmentId"): + middle = path[1:-2] + if len(middle) % 2 == 0 and all( + middle[index] == "parts" and type(middle[index + 1]) is int + for index in range(0, len(middle), 2) + ): + return "gmail_attachment_locator" + return "other_payload" + + +def compare_observations( + before: Mapping[UUID, SourceRecord], after: Mapping[UUID, SourceRecord], provider: str +) -> dict[str, int | dict[str, int]]: + """Return counts only; raw snapshots, identifiers and SHA values stay in memory. + + Full snapshots remain in memory only for the bounded disposable harness. + Manifest comparison aligns part paths, sizes and SHA; a changed path alone is + not evidence of changed bytes. It concerns external blobs, not all bodies. Inline + bytes remain part of raw payload comparison. No inference from equal byte totals. + """ + counts = Counter( + matched_records=0, + changed_revision_records=0, + changed_payload_records=0, + locator_only_payload_records=0, + unchanged_records=0, + changed_revision_without_payload_or_blob_change=0, + records_with_retained_blobs=0, + equal_blob_manifest_records=0, + different_blob_manifest_records=0, + equal_blob_size_different_sha_records=0, + ) + categories = Counter() + for identity in before.keys() & after.keys(): + left, right = before[identity], after[identity] + counts["matched_records"] += 1 + changed_revision = left.revision != right.revision + counts["changed_revision_records"] += changed_revision + counts["unchanged_records"] += not changed_revision + paths = _changed_paths(left.payload, right.payload) + record_categories = [_category(provider, path) for path in paths] + categories.update(record_categories) + counts["changed_payload_records"] += bool(paths) + counts["locator_only_payload_records"] += bool(paths) and set(record_categories) == { + "gmail_attachment_locator" + } + if changed_revision and not paths and left.blobs == right.blobs: + counts["changed_revision_without_payload_or_blob_change"] += 1 + if left.blobs or right.blobs: + counts["records_with_retained_blobs"] += 1 + + # Native part paths align bytes without persisting or publishing those paths. + def manifest(record: SourceRecord, include_sha: bool) -> list[tuple[str, int, str]]: + return sorted( + (blob.source_path or "", blob.size_bytes, blob.sha256 if include_sha else "") + for blob in record.blobs + ) + + same_sha = manifest(left, True) == manifest(right, True) + counts["equal_blob_manifest_records"] += same_sha + counts["different_blob_manifest_records"] += not same_sha + counts["equal_blob_size_different_sha_records"] += not same_sha and manifest( + left, False + ) == manifest(right, False) + return { + "schema_version": 1, + "baseline_records": len(before), + "added_records": len(after.keys() - before.keys()), + "removed_rows": len(before.keys() - after.keys()), + **counts, + "changed_field_categories": dict(sorted(categories.items())), + } diff --git a/backend/tests/live/provider_lifecycle.py b/backend/tests/live/provider_lifecycle.py index b588f1e9ec..9611bd62bf 100644 --- a/backend/tests/live/provider_lifecycle.py +++ b/backend/tests/live/provider_lifecycle.py @@ -23,6 +23,7 @@ import canonical_capture as harness # settings must precede application imports from calendar_lifecycle import count_calendar_request, event_checkpoint, verify_calendar_scopes +from capture_comparison import compare_observations from provider_sample import rest_source, wispr_source from sqlalchemy import select, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine @@ -38,7 +39,7 @@ ) from airweave.domains.entities.canonical.requests import CaptureRecord, CompletedScope, StartedScope from airweave.domains.entities.canonical.service import CanonicalCaptureService -from airweave.domains.entities.canonical.store import CanonicalRecordStore +from airweave.domains.entities.canonical.store import CanonicalRecordStore, source_record 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 @@ -426,6 +427,11 @@ async def db_context(): organization = await db.get(Organization, organization_id) sync = await db.get(Sync, sync_id) job = await prepare_job(db, organization_id, sync_id, job_id, attempt_number) + before = { + row.id: source_record(row) + for row in await db.scalars(select(Entity).where(Entity.sync_id == sync_id)) + if row.record_revision > 0 + } config = SyncConfig() config.behavior.skip_guardrails = True @@ -631,6 +637,11 @@ async def db_context(): "loaded_durable_checkpoint": loaded, "job_status": status, "payload_revision_digest": digest, + "observation_comparison": compare_observations( + before, + {row.id: source_record(row) for row in rows if row.record_revision > 0}, + name, + ), "mode": { "gmail": ("unfiltered_" + saved["canonical_cycle"]["mode"]) if manifest.get("gmail_unfiltered") diff --git a/backend/tests/unit/test_capture_comparison.py b/backend/tests/unit/test_capture_comparison.py new file mode 100644 index 0000000000..a999a8bb49 --- /dev/null +++ b/backend/tests/unit/test_capture_comparison.py @@ -0,0 +1,69 @@ +"""Disposable trial diagnostics distinguish metadata rotation from changed bytes.""" + +import json +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from uuid import uuid4 + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity + + +def test_locator_changes_and_equal_size_different_bytes_are_not_conflated(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[1] / "live")) + from capture_comparison import compare_observations + + def record(content): + digest = sha256(content).hexdigest() + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity(record_type="message", native_id="private-message"), + revision=1, + payload={"payload": {"parts": [{"body": {"attachmentId": "private-locator"}}]}}, + payload_schema_version=1, + capture_hash="private-capture-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=( + BlobReference( + key="private-key", + sha256=digest, + size_bytes=len(content), + source_path="/payload/parts/0/body", + ), + ), + indexed_revision=None, + indexed_pipeline_version=None, + ) + + rotated, changed = record(b"same"), record(b"old!") + new_rotated = rotated.model_copy( + update={ + "revision": 2, + "payload": {"payload": {"parts": [{"body": {"attachmentId": "other-locator"}}]}}, + } + ) + new_changed = changed.model_copy(update={"revision": 2, "blobs": record(b"new!").blobs}) + before = {row.id: row for row in (rotated, changed)} + after = {row.id: row for row in (new_rotated, new_changed)} + assert sum(blob.size_bytes for row in before.values() for blob in row.blobs) == sum( + blob.size_bytes for row in after.values() for blob in row.blobs + ) + result = compare_observations(before, after, "gmail") + assert result["changed_revision_records"] == 2 + assert result["changed_payload_records"] == result["locator_only_payload_records"] == 1 + assert result["changed_field_categories"] == {"gmail_attachment_locator": 1} + assert result["equal_blob_manifest_records"] == 1 + assert result["different_blob_manifest_records"] == 1 + assert result["equal_blob_size_different_sha_records"] == 1 + serialized = json.dumps(result) + assert all(str(identity) not in serialized for identity in before) + assert "private" not in serialized and "other-locator" not in serialized + assert all(blob.sha256 not in serialized for row in after.values() for blob in row.blobs) From 3002471888afe03533a29c98306c494ad945c0bd Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:10:38 -0700 Subject: [PATCH 169/266] Never refresh provider credentials for unauthenticated file downloads --- .../airweave/domains/storage/file_service.py | 13 +++++- .../storage/tests/test_file_service.py | 46 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/storage/file_service.py b/backend/airweave/domains/storage/file_service.py index f3786540ad..736e4d80c9 100644 --- a/backend/airweave/domains/storage/file_service.py +++ b/backend/airweave/domains/storage/file_service.py @@ -91,8 +91,13 @@ async def capture_canonical_url( try: await self._stream_download(client, url, headers, temp_path, logger) except httpx.HTTPStatusError as exc: - if exc.response.status_code != 401 or not auth.supports_refresh: + 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) async with aiofiles.open(temp_path, "rb") as downloaded: @@ -238,7 +243,11 @@ 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") headers = await authorization_headers(auth, refresh=True) try: diff --git a/backend/airweave/domains/storage/tests/test_file_service.py b/backend/airweave/domains/storage/tests/test_file_service.py index a2fc54b294..47ccba2741 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 @@ -159,3 +160,48 @@ async def test_storage_failure_does_not_publish_reference(self, tmp_path): 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"] From f8fa47a065206ed99de4a675c1159daf600521c8 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:14:04 -0700 Subject: [PATCH 170/266] Account for retained and unavailable Slack attachment parts --- .../entities/canonical/projection_mappers.py | 5 + .../entities/canonical/slack_projection.py | 91 +++++++++ .../canonical/tests/test_slack_projection.py | 173 ++++++++++++++++++ .../airweave/platform/entities/__init__.py | 3 +- backend/airweave/platform/entities/slack.py | 11 +- docs/owned-sources/representations.md | 29 +++ 6 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/slack_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_projection.py diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 54a6d79823..5ebf99d428 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -294,6 +294,11 @@ async def map_record( yield await map_gmail(record, storage, directory) return + elif source_name == "slack" and record.identity.record_type == "message": + from airweave.domains.entities.canonical.slack_projection import map_slack_files + + yield await map_slack_files(record, _slack(record)[0], storage, directory) + return else: entities = await _map_entities(record, source_name, storage, directory) parts = tuple(_projection_input(index, entity) for index, entity in enumerate(entities)) diff --git a/backend/airweave/domains/entities/canonical/slack_projection.py b/backend/airweave/domains/entities/canonical/slack_projection.py new file mode 100644 index 0000000000..351e2674af --- /dev/null +++ b/backend/airweave/domains/entities/canonical/slack_projection.py @@ -0,0 +1,91 @@ +"""Offline Slack file inputs from retained message JSON and account-owned blobs.""" + +import mimetypes +import re +from pathlib import Path + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from airweave.domains.entities.canonical.blob_materializer import read_blob, 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.storage.protocols import StorageBackend +from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields +from airweave.platform.entities._base import BaseEntity, Breadcrumb +from airweave.platform.entities.slack import SlackAttachmentEntity + + +class SlackFile(BaseModel): + """Only projection fields are parsed; the retained native payload stays unchanged.""" + + model_config = ConfigDict(extra="ignore", strict=True) + id: str = Field(pattern=r"^F[A-Z0-9]+$", max_length=128) + name: str | None = Field(default=None, max_length=4096) + mimetype: str | None = Field(default=None, max_length=256) + + +class SlackFiles(BaseModel): + """A missing file list means no attachments; malformed lists never mean empty.""" + + model_config = ConfigDict(extra="ignore", strict=True) + files: list[SlackFile] = Field(default_factory=list) + + +async def map_slack_files( + record: SourceRecord, body: BaseEntity, storage: StorageBackend, directory: Path +) -> ProjectionInputs: + """Retain one body and one expected part per file, including uncaptured originals.""" + try: + files = SlackFiles.model_validate(record.payload).files + except ValidationError: + raise ValueError("Slack file metadata is malformed") from None + if len({file.id for file in files}) != len(files): + raise ValueError("Slack message contains duplicate native file identities") + paths = {f"/files/{index}" for index in range(len(files))} + if any(blob.role is None and blob.source_path not in paths for blob in record.blobs): + raise ValueError("Slack blob does not identify a current message file") + populate_base_fields(body) + parts = [ + ProjectionInput(part=ExtractionPart(part_index=0, key="body", kind="body"), entity=body) + ] + for index, file in enumerate(files): + filename = file.name or file.id + suffix = Path(filename).suffix.lower() + if not re.fullmatch(r"\.[a-z0-9]{1,12}", suffix): + suffix = mimetypes.guess_extension(file.mimetype or "") or ".bin" + descriptor = ExtractionPart( + part_index=index + 1, + key=f"file:{file.id}", + kind="file", + media_type=file.mimetype, + extension=suffix, + ) + references = [blob for blob in record.blobs if blob.source_path == f"/files/{index}"] + if len(references) > 1: + raise ValueError("Slack file has ambiguous original blob references") + if not references: + if record.completeness == "complete": + raise ValueError("Complete Slack message lacks a retained file") + parts.append(ProjectionInput(part=descriptor, entity=None)) + continue + content = await read_blob(record, references[0], storage) + path = await write_blob(content, directory, suffix=suffix) + entity = SlackAttachmentEntity( + attachment_key=f"{record.identity.container_id}:{record.identity.native_id}:{file.id}", + filename=filename, + breadcrumbs=[ + Breadcrumb( + entity_id=body.entity_id, name="Slack message", entity_type=type(body).__name__ + ) + ], + # Projection never dereferences URLs; the original is identified by its parent. + url="", + size=len(content), + file_type=suffix.lstrip("."), + mime_type=file.mimetype, + local_path=str(path), + ) + populate_base_fields(entity) + parts.append(ProjectionInput(part=descriptor, entity=entity)) + return ProjectionInputs(parts=tuple(parts)) 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..b990ec4810 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py @@ -0,0 +1,173 @@ +"""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 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 + 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(), 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" diff --git a/backend/airweave/platform/entities/__init__.py b/backend/airweave/platform/entities/__init__.py index e16ae894f1..9d177da0db 100644 --- a/backend/airweave/platform/entities/__init__.py +++ b/backend/airweave/platform/entities/__init__.py @@ -286,7 +286,7 @@ SlabPostEntity, SlabTopicEntity, ) -from .slack import SlackChannelEntity, SlackMessageEntity +from .slack import SlackAttachmentEntity, SlackChannelEntity, SlackMessageEntity from .slite import SliteNoteEntity from .stripe import ( StripeBalanceEntity, @@ -650,6 +650,7 @@ SlabTopicEntity, ], "slack": [ + SlackAttachmentEntity, SlackChannelEntity, SlackMessageEntity, ], diff --git a/backend/airweave/platform/entities/slack.py b/backend/airweave/platform/entities/slack.py index 84b59c007b..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]: @@ -141,3 +141,12 @@ class SlackChannelEntity(BaseEntity): 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/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index c98a4b425b..89daeada41 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -105,3 +105,32 @@ failure: native size exceeded the initial capture limit and no bytes were stored | Native capability/resource | Product name | Access operation | Stored representation | Sync/change guarantee | Proof | Gap | | --- | --- | --- | --- | --- | --- | --- | | File metadata without captured body | Original / extraction availability | Existing record and extraction reads | Native metadata plus explicit unavailable-original coverage | Later captured bytes form a new revision; loss of body retires previous text eligibility | Real SQL metadata-only → retained text → metadata-only lifecycle; complete/missing-blob remains error | Actual recapture depends on provider access and configured limits; no converter success claimed | + +## Slack retained file projection + +Slack messages now enumerate their body and each native `files[]` identity as +separate extraction parts. A retained file blob must reference its exact +`/files/N` location and pass the existing sync ownership, SHA-256 and byte-count +checks. Supported retained files use the existing converter/chunker/embedding +pipeline; missing original bytes are explicit `unavailable_original` parts. The +body remains searchable when a file was deliberately not captured. Unknown or +unsupported conversion outcomes retain the shared strict/unsupported semantics. + +| Native resource | Product name | Access | Stored representation | Change guarantee | Proof | Gap | +| --- | --- | --- | --- | --- | --- | --- | +| Message text | Slack message | Stored read/search | Native JSON | Existing message capture | Offline mapper and SQL publication | Existing source freshness limits | +| Hosted file | Message attachment | Stored original/extraction when retained | Existing immutable blob, `/files/N` | Same message revision and visibility | Synthetic PDF conversion/publication; integrity tests | Source byte acquisition not implemented by this change | +| Uncaptured file | Unavailable attachment | Coverage only | Native metadata | No byte-completeness claim | Explicit missing part tests | No inferred download or extracted content | + +Native file IDs identify extraction parts; byte-identical files with distinct IDs +remain distinct parts. Duplicate file IDs or malformed metadata fail projection +rather than choosing an ambiguous original. Complete records missing declared +file bytes fail integrity checks; explicitly partial originals can publish partial +extraction. No provider URL is fetched during projection. Attachments remain +parts of their parent message: this adds no independent file authorization or +file-only search facet. + +Existing indexed Slack generations require explicit reprojection before this +coverage appears. No global pipeline version was bumped. The current Slack source +still retains file metadata only; this projection change does not establish live +`files:read` access, download support, whole-workspace completeness, or freshness. From 626abd042f3934ac4e4f040457aba8bbe9803a9e Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:22:22 -0700 Subject: [PATCH 171/266] Allow canonical file capture to reject native redirects --- .../auth_provider/providers/composio.py | 2 +- .../airweave/domains/storage/file_service.py | 14 +++-- .../storage/tests/test_file_service.py | 53 ++++++++++++++++++- .../http_client/test_composio_transport.py | 37 +++++++++++++ 4 files changed, 101 insertions(+), 5 deletions(-) diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 5ef08f1e9d..7d8c4434de 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -117,7 +117,7 @@ async def get_auth_result( "google_drive": frozenset( {"www.googleapis.com", "docs.googleapis.com", "sheets.googleapis.com"} ), - "slack": frozenset({"slack.com"}), + "slack": frozenset({"slack.com", "files.slack.com"}), "linear": frozenset({"api.linear.app", "uploads.linear.app"}), "github": frozenset({"api.github.com"}), "notion": frozenset({"api.notion.com"}), diff --git a/backend/airweave/domains/storage/file_service.py b/backend/airweave/domains/storage/file_service.py index 736e4d80c9..26f3b8a704 100644 --- a/backend/airweave/domains/storage/file_service.py +++ b/backend/airweave/domains/storage/file_service.py @@ -80,6 +80,8 @@ async def capture_canonical_url( auth: SourceAuthProvider, logger: ContextualLogger, media_type: str | None = None, + *, + follow_redirects: bool = True, ) -> BlobReference: """Capture original bytes with size/auth checks, independent of search file types.""" if self.sync_id is None: @@ -89,7 +91,9 @@ async def capture_canonical_url( temp_path = f"{self.base_temp_dir}/{uuid4()}-canonical" try: try: - await self._stream_download(client, url, headers, temp_path, logger) + await self._stream_download( + client, url, headers, temp_path, logger, follow_redirects=follow_redirects + ) except httpx.HTTPStatusError as exc: if ( exc.response.status_code != 401 @@ -99,7 +103,9 @@ async def capture_canonical_url( 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) + await self._stream_download( + client, url, headers, temp_path, logger, follow_redirects=follow_redirects + ) 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) @@ -170,13 +176,15 @@ async def _stream_download( headers: dict, temp_path: str, logger: ContextualLogger, + *, + follow_redirects: bool = True, ) -> 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() diff --git a/backend/airweave/domains/storage/tests/test_file_service.py b/backend/airweave/domains/storage/tests/test_file_service.py index 47ccba2741..435952b230 100644 --- a/backend/airweave/domains/storage/tests/test_file_service.py +++ b/backend/airweave/domains/storage/tests/test_file_service.py @@ -138,7 +138,7 @@ async def test_canonical_capture_keeps_original_bytes_and_cleans_temp(self, tmp_ svc.sync_id = uuid4() content = b"original unsupported format" - async def download(client, url, headers, destination, logger): + async def download(client, url, headers, destination, logger, *, follow_redirects=True): with open(destination, "wb") as output: output.write(content) @@ -205,3 +205,54 @@ async def download(request): 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" diff --git a/backend/tests/unit/platform/http_client/test_composio_transport.py b/backend/tests/unit/platform/http_client/test_composio_transport.py index e12c618deb..cb3bd11830 100644 --- a/backend/tests/unit/platform/http_client/test_composio_transport.py +++ b/backend/tests/unit/platform/http_client/test_composio_transport.py @@ -361,3 +361,40 @@ async def test_attio_binding_uses_only_native_api_host_and_verified_account(monk fetch.return_value = {"id": "ca_other", "toolkit": {"slug": "attio"}, "status": "ACTIVE"} with pytest.raises(AuthProviderConfigError, match="identity"): await provider.get_auth_result("attio", ["access_token"]) + + +@pytest.mark.asyncio +async def test_slack_binding_permits_exact_file_host_not_lookalike(monkeypatch): + from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider + + provider = await ComposioAuthProvider.create( + credentials={"api_key": "key"}, config={"account_id": "ca_slack"} + ) + monkeypatch.setattr( + provider, + "_get_with_auth", + AsyncMock( + return_value={"id": "ca_slack", "toolkit": {"slug": "slack"}, "status": "ACTIVE"} + ), + ) + result = await provider.get_auth_result("slack", ["access_token"]) + assert result.managed_auth.allowed_hosts == {"slack.com", "files.slack.com"} + calls = [] + + async def proxy(request): + calls.append(json.loads(request.content)["endpoint"]) + return httpx.Response(200, json={"status": 200, "data": "file", "headers": {}}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(proxy)) as upstream: + async with httpx.AsyncClient( + transport=ComposioTransport( + client=upstream, + api_key="key", + connected_account_id="ca_slack", + allowed_hosts=result.managed_auth.allowed_hosts, + ) + ) as client: + await client.get("https://files.slack.com/files-pri/original.pdf") + with pytest.raises(ComposioProxyError): + await client.get("https://files.slack.com.evil.example/original.pdf") + assert calls == ["https://files.slack.com/files-pri/original.pdf"] From a12eb51102328e1bc466d052a6230bb7fbbfb4dc Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:27:43 -0700 Subject: [PATCH 172/266] Define retained Slack file acquisition evidence --- .../domains/entities/canonical/slack_files.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 backend/airweave/domains/entities/canonical/slack_files.py diff --git a/backend/airweave/domains/entities/canonical/slack_files.py b/backend/airweave/domains/entities/canonical/slack_files.py new file mode 100644 index 0000000000..f38c87af82 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/slack_files.py @@ -0,0 +1,46 @@ +"""Retained acquisition evidence for files owned by a Slack message.""" + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator + +SlackFileReason = Literal[ + "access_denied", "unsupported_external", "oversized", "missing_metadata", "not_found" +] + + +class SlackFileOutcome(BaseModel): + """One native file position, without modifying the message's native JSON.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + index: int = Field(ge=0, lt=100) + native_id: str = Field(pattern=r"^F[A-Z0-9]+$", max_length=128) + outcome: Literal["captured", "unavailable"] + reason: SlackFileReason | None = None + file: dict[str, JsonValue] | None = None + + @model_validator(mode="after") + def consistent(self): + """Retained enrichment must attest the same native file.""" + if (self.outcome == "captured") != (self.reason is None): + raise ValueError("Slack file outcome and reason disagree") + if self.file is not None and self.file.get("id") != self.native_id: + raise ValueError("Slack file enrichment identity disagrees") + return self + + +class SlackFileManifest(BaseModel): + """Bounded, ordered per-message acquisition results.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + version: Literal[1] = 1 + files: tuple[SlackFileOutcome, ...] = Field(max_length=100) + + @model_validator(mode="after") + def positions(self): + """No duplicate identities or omitted native positions.""" + if [entry.index for entry in self.files] != list(range(len(self.files))): + raise ValueError("Slack file positions are not contiguous") + if len({entry.native_id for entry in self.files}) != len(self.files): + raise ValueError("Slack file identities are duplicated") + return self From 919001e046790e33e1bab17fbeef0dd4cac0b498 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:27:55 -0700 Subject: [PATCH 173/266] Validate Slack file acquisition manifests during projection --- .../entities/canonical/slack_projection.py | 37 ++++++++++++-- .../canonical/tests/test_slack_projection.py | 49 +++++++++++++++++++ docs/owned-sources/representations.md | 21 ++++++-- 3 files changed, 100 insertions(+), 7 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/slack_projection.py b/backend/airweave/domains/entities/canonical/slack_projection.py index 351e2674af..fd0e6ea76d 100644 --- a/backend/airweave/domains/entities/canonical/slack_projection.py +++ b/backend/airweave/domains/entities/canonical/slack_projection.py @@ -10,6 +10,7 @@ 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.slack_files import SlackFileManifest 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, Breadcrumb @@ -32,19 +33,43 @@ class SlackFiles(BaseModel): files: list[SlackFile] = Field(default_factory=list) -async def map_slack_files( - record: SourceRecord, body: BaseEntity, storage: StorageBackend, directory: Path -) -> ProjectionInputs: - """Retain one body and one expected part per file, including uncaptured originals.""" +async def _file_inventory( + record: SourceRecord, storage: StorageBackend +) -> tuple[list[SlackFile], SlackFileManifest | None]: + """Verify optional acquisition evidence against unchanged native file positions.""" try: files = SlackFiles.model_validate(record.payload).files except ValidationError: raise ValueError("Slack file metadata is malformed") from None if len({file.id for file in files}) != len(files): raise ValueError("Slack message contains duplicate native file identities") + manifests = [blob for blob in record.blobs if blob.role == "representation_manifest"] + if len(manifests) > 1: + raise ValueError("Slack message has multiple file acquisition manifests") + manifest = None + if manifests: + content = await read_blob(record, manifests[0], storage) + try: + manifest = SlackFileManifest.model_validate_json(content) + if [entry.native_id for entry in manifest.files] != [file.id for file in files]: + raise ValueError("Slack file manifest does not match native message identities") + files = [ + SlackFile.model_validate(entry.file) if entry.file is not None else file + for file, entry in zip(files, manifest.files, strict=True) + ] + except ValidationError: + raise ValueError("Slack file acquisition manifest is malformed") from None paths = {f"/files/{index}" for index in range(len(files))} if any(blob.role is None and blob.source_path not in paths for blob in record.blobs): raise ValueError("Slack blob does not identify a current message file") + return files, manifest + + +async def map_slack_files( + record: SourceRecord, body: BaseEntity, storage: StorageBackend, directory: Path +) -> ProjectionInputs: + """Retain one body and one expected part per file, including uncaptured originals.""" + files, manifest = await _file_inventory(record, storage) populate_base_fields(body) parts = [ ProjectionInput(part=ExtractionPart(part_index=0, key="body", kind="body"), entity=body) @@ -64,6 +89,10 @@ async def map_slack_files( references = [blob for blob in record.blobs if blob.source_path == f"/files/{index}"] if len(references) > 1: raise ValueError("Slack file has ambiguous original blob references") + if manifest is not None: + outcome = manifest.files[index] + if (outcome.outcome == "captured") != bool(references): + raise ValueError("Slack file manifest contradicts retained original blobs") if not references: if record.completeness == "complete": raise ValueError("Complete Slack message lacks a retained file") diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py index b990ec4810..7c0e9d6451 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py @@ -171,3 +171,52 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source ] assert coverage.parts[1].key == "file:F1" assert row.completeness == "partial" + + +@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) diff --git a/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index 89daeada41..6abdb13b80 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -131,6 +131,21 @@ parts of their parent message: this adds no independent file authorization or file-only search facet. Existing indexed Slack generations require explicit reprojection before this -coverage appears. No global pipeline version was bumped. The current Slack source -still retains file metadata only; this projection change does not establish live -`files:read` access, download support, whole-workspace completeness, or freshness. +coverage appears. No global pipeline version was bumped. Slack original acquisition remains opt-in (`capture_files=false` by default). +This projection change does not establish live `files:read` access, successful +provider downloads, whole-workspace completeness, or freshness. + +When present, the Slack file acquisition manifest is read from its verified +`representation_manifest` blob. Its ordered native IDs must match the unchanged +message file list. Retained `files.info` enrichment can supply filename/MIME type +for Slack Connect stubs, but cannot replace the file identity. A captured outcome +requires exactly one native-path blob; an unavailable outcome must have none. +Contradictory evidence fails projection. Legacy records without a manifest retain +the explicit missing-original behavior above. + +The companion capture contract requires an original URL, native size and MIME +type, resolving missing fields through `files.info`. Missing metadata and explicit +file denial remain partial acquisition evidence; missing OAuth scope fails the +configured capture. Downloaded byte-count mismatch fails the page. Capture still +commits whole provider pages: aggregate page download/retry cost has not been +qualified live, even though individual files and metadata counts are bounded. From 8c50d875be290a6762a45739ab182c8b66c748ec Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:29:41 -0700 Subject: [PATCH 174/266] Validate declared file response MIME before retaining bytes --- .../airweave/domains/storage/file_service.py | 30 +++++++++++++- .../storage/tests/test_file_service.py | 40 ++++++++++++++++++- 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/backend/airweave/domains/storage/file_service.py b/backend/airweave/domains/storage/file_service.py index 26f3b8a704..5500d0974c 100644 --- a/backend/airweave/domains/storage/file_service.py +++ b/backend/airweave/domains/storage/file_service.py @@ -82,6 +82,7 @@ async def capture_canonical_url( 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: @@ -92,7 +93,13 @@ async def capture_canonical_url( try: try: await self._stream_download( - client, url, headers, temp_path, logger, follow_redirects=follow_redirects + client, + url, + headers, + temp_path, + logger, + follow_redirects=follow_redirects, + expected_media_type=expected_media_type, ) except httpx.HTTPStatusError as exc: if ( @@ -104,7 +111,13 @@ async def capture_canonical_url( # 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 + 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) @@ -178,6 +191,7 @@ async def _stream_download( 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( @@ -188,6 +202,18 @@ async def _stream_download( 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: diff --git a/backend/airweave/domains/storage/tests/test_file_service.py b/backend/airweave/domains/storage/tests/test_file_service.py index 435952b230..3cbbfa6878 100644 --- a/backend/airweave/domains/storage/tests/test_file_service.py +++ b/backend/airweave/domains/storage/tests/test_file_service.py @@ -138,7 +138,16 @@ async def test_canonical_capture_keeps_original_bytes_and_cleans_temp(self, tmp_ svc.sync_id = uuid4() content = b"original unsupported format" - async def download(client, url, headers, destination, logger, *, follow_redirects=True): + async def download( + client, + url, + headers, + destination, + logger, + *, + follow_redirects=True, + expected_media_type=None, + ): with open(destination, "wb") as output: output.write(content) @@ -256,3 +265,32 @@ async def download(request): 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") From 7ecc97524c91840c28dfbadc4230d586a225ac01 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:30:00 -0700 Subject: [PATCH 175/266] Capture Slack originals before page progress commits --- backend/airweave/platform/configs/config.py | 6 + backend/airweave/platform/sources/slack.py | 26 +- .../platform/sources/slack_content.py | 179 ++++++++++++++ .../airweave/platform/sources/slack_errors.py | 10 + .../platform/sources/test_slack_content.py | 233 ++++++++++++++++++ 5 files changed, 442 insertions(+), 12 deletions(-) create mode 100644 backend/airweave/platform/sources/slack_content.py create mode 100644 backend/airweave/platform/sources/slack_errors.py create mode 100644 backend/tests/unit/platform/sources/test_slack_content.py diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 6eb8a76b8f..af6a14b195 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -803,6 +803,12 @@ def validate_host_ssrf(cls, v: str) -> str: class SlackConfig(SourceConfig): """Trusted native workspace/user pair for owned Slack capture.""" + capture_files: bool = Field( + default=False, + description="Opt in to original file acquisition before each complete page commits; " + "aggregate page byte and duration budgets remain unqualified.", + ) + 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+$") diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index b03ae0c9c0..4b5e6a51ea 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -58,6 +58,8 @@ ) 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 capture_slack_files +from airweave.platform.sources.slack_errors import SlackApiError from airweave.schemas.source_connection import AuthenticationMethod, OAuthType @@ -125,15 +127,6 @@ def _replace(m: re.Match[str]) -> str: return re.sub(r" +", " ", cleaned).strip() -class SlackApiError(ValueError): - """A Slack application error; unlike empty pages it cannot complete a scope.""" - - def __init__(self, code: str): - """Retain the safe provider code for explicit access-loss handling.""" - self.code = code - super().__init__(f"Slack request failed: {code}") - - class SlackPrincipal(BaseModel): """Native authorized workspace/user identity, with no display-name inference.""" @@ -223,9 +216,12 @@ def _require_principal(self) -> SlackPrincipal: 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"] = True fingerprint = hashlib.sha256( json.dumps( - {"version": 2, "team_id": principal.team_id, "user_id": principal.user_id}, + material, sort_keys=True, ).encode() ).hexdigest() @@ -295,7 +291,7 @@ async def _get(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict: @staticmethod def _capture_message(message: dict, channel_id: str) -> CaptureRecord: - """Keep the complete native response; file bytes are explicitly not captured yet.""" + """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( @@ -326,7 +322,13 @@ async def capture_page( """Fetch one page; record and nested reply progress are committed by the pipeline.""" self._require_principal() try: - return await self._validated_capture_page(scope, continuation) + page = await self._validated_capture_page(scope, continuation) + if self.slack_config.capture_files and scope.record_type == "message": + records = tuple( + [await capture_slack_files(self, record, files) for record in page.records] + ) + page = page.model_copy(update={"records": records}) + return page except ValidationError: # Validation diagnostics may contain private provider values. raise ValueError("Slack capture returned invalid page or continuation data") from None diff --git a/backend/airweave/platform/sources/slack_content.py b/backend/airweave/platform/sources/slack_content.py new file mode 100644 index 0000000000..6e544122c1 --- /dev/null +++ b/backend/airweave/platform/sources/slack_content.py @@ -0,0 +1,179 @@ +"""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) + 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_files( + source: "SlackSource", record: CaptureRecord, files: FileService +) -> CaptureRecord: + """Whole-page caller commits only after every message acquisition succeeds.""" + metadata = SlackMessageFiles.model_validate(record.payload).files + if not metadata: + return record + if len({file.id for file in metadata}) != len(metadata): + raise ValueError("Slack message contains duplicate file identities") + outcomes = [] + blobs = [] + for index, file in enumerate(metadata): + file, native, reason = await _enrich(source, file) + 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": f"/files/{index}"})) + outcomes.append( + SlackFileOutcome( + index=index, + native_id=file.id, + outcome="captured" if reason is None else "unavailable", + reason=reason, + file=native, + ) + ) + manifest = SlackFileManifest(files=tuple(outcomes)) + 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 any(item.reason for item in outcomes) 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, 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/tests/unit/platform/sources/test_slack_content.py b/backend/tests/unit/platform/sources/test_slack_content.py new file mode 100644 index 0000000000..8b5cfaef19 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_slack_content.py @@ -0,0 +1,233 @@ +"""Whole-page Slack originals: opt-in, immutable provenance and no failed-page progress.""" + +import hashlib +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from airweave.domains.entities.canonical.requests import BlobReference, CompletedScope +from airweave.domains.entities.canonical.scan_models import ScanContinuation +from airweave.domains.entities.canonical.slack_files import SlackFileManifest +from airweave.domains.sources.exceptions import SourceError +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.platform.configs.config import SlackConfig +from airweave.platform.sources.slack import SlackPrincipal, SlackSource +from airweave.platform.sources.slack_content import capture_slack_files + + +def connector(): + source = SlackSource( + auth=StaticTokenProvider("token"), logger=MagicMock(), http_client=MagicMock() + ) + source.slack_config = SlackConfig( + expected_team_id="T1", expected_user_id="U1", capture_files=True + ) + source._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") + return source + + +def storage(): + files = MagicMock(MAX_FILE_SIZE_BYTES=200 * 1024 * 1024) + saved = [] + + async def store(content, *, media_type=None): + saved.append(content) + digest = hashlib.sha256(content).hexdigest() + return BlobReference( + key=digest, sha256=digest, size_bytes=len(content), media_type=media_type + ) + + files.store_canonical_blob = AsyncMock(side_effect=store) + files.capture_canonical_url = AsyncMock( + return_value=BlobReference( + key="original", sha256="a" * 64, size_bytes=3, media_type="text/plain" + ) + ) + return files, saved + + +@pytest.mark.asyncio +async def test_capture_enriches_separately_and_retains_explicit_partial_outcomes(): + source = connector() + message = { + "ts": "1", + "files": [ + {"id": "F1"}, + {"id": "F2", "is_external": True}, + { + "id": "F3", + "size": 300 * 1024 * 1024, + "mimetype": "text/plain", + "url_private": "https://files.slack.com/large", + }, + ], + } + source._get = AsyncMock( + return_value={ + "ok": True, + "file": { + "id": "F1", + "name": "note.txt", + "mimetype": "text/plain", + "size": 3, + "url_private_download": "https://files.slack.com/original", + }, + } + ) + files, saved = storage() + record = await capture_slack_files(source, source._capture_message(message, "C1"), files) + assert record.payload == message and "name" not in record.payload["files"][0] + manifest = SlackFileManifest.model_validate_json(saved[-1]) + assert [(x.outcome, x.reason) for x in manifest.files] == [ + ("captured", None), + ("unavailable", "unsupported_external"), + ("unavailable", "oversized"), + ] + assert manifest.files[0].file["name"] == "note.txt" + assert record.completeness == "partial" + assert [blob.source_path for blob in record.blobs] == ["/files/0", None] + assert record.blobs[-1].role == "representation_manifest" + assert files.capture_canonical_url.call_args.kwargs["follow_redirects"] is False + assert files.capture_canonical_url.call_args.kwargs["expected_media_type"] == "text/plain" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [302, 401, 429, 503]) +async def test_late_download_failure_returns_no_page_or_new_continuation(status): + source = connector() + files, saved = storage() + source._get = AsyncMock( + return_value={ + "messages": [ + { + "ts": "1", + "files": [ + { + "id": "F1", + "size": 3, + "mimetype": "text/plain", + "url_private": "https://files.slack.com/one", + }, + { + "id": "F2", + "size": 3, + "mimetype": "text/plain", + "url_private": "https://files.slack.com/two", + }, + ], + } + ], + "response_metadata": {"next_cursor": "next"}, + } + ) + error = httpx.HTTPStatusError( + "private", + request=httpx.Request("GET", "https://files.slack.com/private"), + response=httpx.Response(status), + ) + files.capture_canonical_url.side_effect = [files.capture_canonical_url.return_value, error] + continuation = ScanContinuation() + with pytest.raises(SourceError): + await source.capture_page( + CompletedScope(record_type="message", container_id="C1"), continuation, files=files + ) + assert continuation.value == {} and not saved + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "http://files.slack.com/x", + "https://files.slack.com.evil/x", + "https://token@slack.com/x", + "https://slack.com:444/x", + ], +) +async def test_untrusted_original_never_downloads(url): + source = connector() + files, _ = storage() + record = source._capture_message({"ts": "1", "files": [{"id": "F1", "url_private": url}]}, "C1") + with pytest.raises(ValueError, match="permitted origin"): + await capture_slack_files(source, record, files) + files.capture_canonical_url.assert_not_called() + + +def test_capture_files_is_opt_in_and_changes_only_enabled_fingerprint(): + source = connector() + enabled = source.capture_cycle_configuration.fingerprint + source.slack_config.capture_files = False + legacy = source.capture_cycle_configuration.fingerprint + import json + + expected = hashlib.sha256( + json.dumps({"version": 2, "team_id": "T1", "user_id": "U1"}, sort_keys=True).encode() + ).hexdigest() + assert legacy == expected and enabled != legacy + assert SlackConfig().capture_files is False + + +@pytest.mark.asyncio +async def test_size_race_and_metadata_identity_mismatch_do_not_publish_manifest(): + source = connector() + files, saved = storage() + record = source._capture_message( + { + "ts": "1", + "files": [ + { + "id": "F1", + "size": 4, + "mimetype": "text/plain", + "url_private": "https://files.slack.com/original", + } + ], + }, + "C1", + ) + with pytest.raises(ValueError, match="size changed"): + await capture_slack_files(source, record, files) + assert not saved + source._get = AsyncMock(return_value={"file": {"id": "F2"}}) + record = source._capture_message({"ts": "1", "files": [{"id": "F1"}]}, "C1") + with pytest.raises(ValueError, match="identity changed"): + await capture_slack_files(source, record, files) + assert not saved + + +@pytest.mark.asyncio +async def test_explicit_file_denial_is_partial_but_missing_scope_fails_page(): + from airweave.platform.sources.slack import SlackApiError + + source = connector() + files, saved = storage() + record = source._capture_message({"ts": "1", "files": [{"id": "F1"}]}, "C1") + source._get = AsyncMock(side_effect=SlackApiError("not_visible")) + result = await capture_slack_files(source, record, files) + assert result.completeness == "partial" + assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "access_denied" + source._get = AsyncMock(side_effect=SlackApiError("missing_scope")) + files.store_canonical_blob.reset_mock() + with pytest.raises(SlackApiError): + await capture_slack_files(source, record, files) + files.store_canonical_blob.assert_not_called() + + +@pytest.mark.asyncio +async def test_missing_metadata_remains_partial_and_never_uses_thumbnail(): + source = connector() + files, saved = storage() + record = source._capture_message( + {"ts": "1", "files": [{"id": "F1", "url_private": "https://files.slack.com/original"}]}, + "C1", + ) + source._get = AsyncMock( + return_value={ + "file": {"id": "F1", "thumb_360": "https://files.slack.com/thumbnail", "size": 3} + } + ) + result = await capture_slack_files(source, record, files) + assert result.completeness == "partial" + assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "missing_metadata" + files.capture_canonical_url.assert_not_called() From f2fe8ba44ea03710c9d425ee6e3223ea22cf4106 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:43:19 -0700 Subject: [PATCH 176/266] Attest full child scope parent epochs before acquisition --- .../domains/entities/canonical/scan_store.py | 6 +- .../entities/canonical/tests/test_cycles.py | 8 ++- .../tests/test_discovered_records.py | 3 + .../tests/test_full_parent_admission.py | 71 +++++++++++++++++++ .../domains/sync_pipeline/canonical_scan.py | 6 +- 5 files changed, 86 insertions(+), 8 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_full_parent_admission.py diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index b67eff7d39..2178ed2810 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -207,9 +207,9 @@ def validate_scope_changes( def attest_planned_owner(cycle: CaptureCycle, request: BeginScan, parent: Entity | None) -> None: """The post-I/O plan must still describe the captured owner used to select it.""" - if cycle.mode == "mixed" and ( - request.expected_parent_epoch != (parent.visibility_epoch if parent else None) - or request.expected_parent_revision != (parent.record_revision if parent else None) + if request.expected_parent_epoch != (parent.visibility_epoch if parent else None) or ( + cycle.mode == "mixed" + and request.expected_parent_revision != (parent.record_revision if parent else None) ): raise ScanConflict("Scope owner changed during plan selection") diff --git a/backend/airweave/domains/entities/canonical/tests/test_cycles.py b/backend/airweave/domains/entities/canonical/tests/test_cycles.py index cc148f2a91..ae2e18062c 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_cycles.py +++ b/backend/airweave/domains/entities/canonical/tests/test_cycles.py @@ -13,7 +13,11 @@ CompleteCycle, CycleConfiguration, ) -from airweave.domains.entities.canonical.cycle_store import CycleConflict, complete_cycle +from airweave.domains.entities.canonical.cycle_store import ( + CycleConflict, + complete_cycle, + scope_owner, +) from airweave.domains.entities.canonical.requests import ( CaptureRecord, CompletedScope, @@ -57,6 +61,8 @@ async def cycle(database, service, fence, **kwargs): async def scan(database, service, fence, active, scope=ROOT, **kwargs): async with database() as db: + parent = await scope_owner(db, fence, active, scope) + kwargs.setdefault("expected_parent_epoch", parent.visibility_epoch if parent else None) return await service.begin_scan( db, BeginScan( diff --git a/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py b/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py index 6625910a9b..4f8cf010df 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py +++ b/backend/airweave/domains/entities/canonical/tests/test_discovered_records.py @@ -7,6 +7,7 @@ from sqlalchemy import select from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration +from airweave.domains.entities.canonical.cycle_store import scope_owner from airweave.domains.entities.canonical.requests import CompletedScope, RecordIdentity from airweave.domains.entities.canonical.scan_models import ( BeginScan, @@ -27,6 +28,7 @@ def original(kind, native): async def begin(database, service, fence, cycle, scope): async with database() as db: + parent = await scope_owner(db, fence, cycle, scope) return await service.begin_scan( db, BeginScan( @@ -34,6 +36,7 @@ async def begin(database, service, fence, cycle, scope): cycle_id=cycle.version.cycle_id, scope=scope, fingerprint=cycle.configuration.fingerprint, + expected_parent_epoch=parent.visibility_epoch if parent else None, ), ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_full_parent_admission.py b/backend/airweave/domains/entities/canonical/tests/test_full_parent_admission.py new file mode 100644 index 0000000000..47ba5f6e24 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_full_parent_admission.py @@ -0,0 +1,71 @@ +"""Full child acquisition must admit the exact parent inventory snapshot.""" + +import pytest + +from airweave.domains.entities.canonical.scan_models import BeginScan +from airweave.domains.entities.canonical.scan_store import ScanConflict +from airweave.domains.entities.canonical.tests.helpers import capture +from airweave.domains.entities.canonical.tests.test_cycles import ( + CHILD, + CONFIG, + cycle, + finish, + page, + record, + scan, +) + + +async def test_full_parent_inventory_epoch_is_required_at_begin_resume_and_restart( + database, source +): + service, fence = source + active = await cycle(database, service, fence) + root = await scan(database, service, fence, active) + parent = record().model_copy( + update={"payload": {"files": ["F1"]}, "descendant_visibility_fields": ("files",)} + ) + root = await page(database, service, fence, root, parent, final=True) + await finish(database, service, fence, root) + async with database() as db: + work = await service.next_scope_work(db, fence, active.version.cycle_id) + request = BeginScan( + fence=fence, + scope=CHILD, + cycle_id=active.version.cycle_id, + fingerprint=CONFIG.fingerprint, + expected_parent_epoch=work.parent_visibility_epoch, + ) + # A source may retain this parent while another admitted operation updates it. + await capture( + database, service, fence, parent.model_copy(update={"payload": {"files": ["F2"]}}) + ) + async with database() as db: + with pytest.raises(ScanConflict, match="owner changed"): + await service.begin_scan(db, request) + async with database() as db: + current = await service.next_scope_work(db, fence, active.version.cycle_id) + request = request.model_copy( + update={"expected_parent_epoch": current.parent_visibility_epoch} + ) + admitted = await service.begin_scan(db, request) + # Full acquisition depends on attachment inventory, not unrelated text revision. + await capture( + database, + service, + fence, + parent.model_copy(update={"payload": {"files": ["F2"], "text": "edited"}}), + ) + async with database() as db: + resumed = await service.begin_scan(db, request) + assert resumed.version == admitted.version + await capture( + database, service, fence, parent.model_copy(update={"payload": {"files": ["F3"]}}) + ) + for restart in (False, True): + async with database() as db: + with pytest.raises(ScanConflict, match="owner changed"): + await service.begin_scan( + db, + request.model_copy(update={"expected": admitted.version, "restart": restart}), + ) diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index c0409ef7fb..14baf7bc3d 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -300,7 +300,7 @@ async def begin_scope( restart=restart, continuation=initial, plan=plan, - expected_parent_epoch=parent_epoch if cycle.mode == "mixed" else None, + expected_parent_epoch=parent_epoch, expected_parent_revision=parent.revision if parent and cycle.mode == "mixed" else None, @@ -338,9 +338,7 @@ async def restart_scope( restart=True, continuation=initial, plan=plan, - expected_parent_epoch=state.parent_visibility_epoch - if cycle.mode == "mixed" - else None, + expected_parent_epoch=state.parent_visibility_epoch, expected_parent_revision=parent.revision if parent and cycle.mode == "mixed" else None, From 3d57a1c734d57b44a87eed6b191f7bf1f93a5425 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:48:18 -0700 Subject: [PATCH 177/266] List authorized direct children by canonical parent record ID --- backend/airweave/api/v1/endpoints/records.py | 8 +- .../domains/entities/canonical/query.py | 13 ++- .../entities/canonical/query_models.py | 1 + .../domains/entities/canonical/query_store.py | 10 ++ .../entities/canonical/tests/test_query.py | 103 ++++++++++++++++++ 5 files changed, 133 insertions(+), 2 deletions(-) diff --git a/backend/airweave/api/v1/endpoints/records.py b/backend/airweave/api/v1/endpoints/records.py index b21a61853b..075855df55 100644 --- a/backend/airweave/api/v1/endpoints/records.py +++ b/backend/airweave/api/v1/endpoints/records.py @@ -92,6 +92,7 @@ async def list_records( sync_id: UUID, record_type: str | None = None, container_id: str | None = None, + parent_record_id: UUID | None = None, state: Literal["active", "deleted", "all"] = "active", cursor: str | None = None, limit: int = Query(100, ge=1, le=500), @@ -101,7 +102,12 @@ async def list_records( ) -> RecordPage: """List exact stored records in stable ID order; continuation is a live traversal.""" query = RecordListQuery( - filters=RecordFilters(record_type=record_type, container_id=container_id, state=state), + filters=RecordFilters( + record_type=record_type, + container_id=container_id, + parent_record_id=parent_record_id, + state=state, + ), cursor=cursor, limit=limit, ) diff --git a/backend/airweave/domains/entities/canonical/query.py b/backend/airweave/domains/entities/canonical/query.py index e5c0be5404..634f737c8e 100644 --- a/backend/airweave/domains/entities/canonical/query.py +++ b/backend/airweave/domains/entities/canonical/query.py @@ -124,8 +124,19 @@ async def list_records( if cursor.filters != query.filters: raise InvalidRecordCursor("Preserve the original filters when continuing a list") after_id = cursor.after_id + parent = None + if query.filters.parent_record_id is not None: + parent = await self.read(db, organization_id, sync_id, query.filters.parent_record_id) + if parent.deleted_at is not None or parent.content_access != "available": + raise RecordNotFound("Record not found in this source") rows = await self.queries.list_records( - db, organization_id, sync_id, query.filters, after_id=after_id, limit=query.limit + db, + organization_id, + sync_id, + query.filters, + after_id=after_id, + limit=query.limit, + parent=parent.identity if parent else None, ) more = len(rows) > query.limit page = rows[: query.limit] diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index c3809eefb6..e9360a7d55 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -60,6 +60,7 @@ class RecordFilters(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) record_type: str | None = None container_id: str | None = None + parent_record_id: UUID | None = None state: Literal["active", "deleted", "all"] = "active" diff --git a/backend/airweave/domains/entities/canonical/query_store.py b/backend/airweave/domains/entities/canonical/query_store.py index d6bbf7d4cc..fa6f309b8f 100644 --- a/backend/airweave/domains/entities/canonical/query_store.py +++ b/backend/airweave/domains/entities/canonical/query_store.py @@ -8,6 +8,7 @@ from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.query_models import RecordFilters +from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.entities.canonical.store import ( SourceNotFound, content_is_available, @@ -51,6 +52,7 @@ async def list_records( *, after_id: UUID | None, limit: int, + parent: RecordIdentity | None = None, ) -> tuple[SourceRecord, ...]: """Fetch one extra row so continuation does not require a count query.""" scope = await db.scalar( @@ -64,6 +66,14 @@ async def list_records( Entity.record_revision > 0, parent_is_visible(), ) + if filters.parent_record_id is not None and parent is None: + raise ValueError("Parent filter requires an authorized parent identity") + if parent is not None: + statement = statement.where( + Entity.parent_record_type == parent.record_type, + Entity.parent_native_id == parent.native_id, + Entity.parent_container_id.is_not_distinct_from(parent.container_id), + ) if filters.record_type is not None: statement = statement.where(Entity.entity_definition_short_name == filters.record_type) if filters.container_id is not None: diff --git a/backend/airweave/domains/entities/canonical/tests/test_query.py b/backend/airweave/domains/entities/canonical/tests/test_query.py index 4f02991db3..af7396b574 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_query.py +++ b/backend/airweave/domains/entities/canonical/tests/test_query.py @@ -120,3 +120,106 @@ async def test_deleted_record_is_readable_but_not_in_active_list(database, sourc 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), + ) From db2159b3f717fb609adcd51b717cae893a79931a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:48:58 -0700 Subject: [PATCH 178/266] Project durable Slack file children under explicit capture topology --- .../entities/canonical/projection_mappers.py | 5 + .../entities/canonical/slack_projection.py | 117 ++++++++++++-- .../canonical/tests/test_slack_projection.py | 152 ++++++++++++++++++ docs/owned-sources/representations.md | 25 +++ 4 files changed, 282 insertions(+), 17 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 5ebf99d428..45c7066b1a 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -294,6 +294,11 @@ async def map_record( yield await map_gmail(record, storage, directory) return + elif source_name == "slack" and record.identity.record_type == "file": + from airweave.domains.entities.canonical.slack_projection import map_slack_file + + yield await map_slack_file(record, storage, directory) + return elif source_name == "slack" and record.identity.record_type == "message": from airweave.domains.entities.canonical.slack_projection import map_slack_files diff --git a/backend/airweave/domains/entities/canonical/slack_projection.py b/backend/airweave/domains/entities/canonical/slack_projection.py index fd0e6ea76d..dafdd95ab2 100644 --- a/backend/airweave/domains/entities/canonical/slack_projection.py +++ b/backend/airweave/domains/entities/canonical/slack_projection.py @@ -10,6 +10,7 @@ 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.requests import BlobReference, parent_container_key from airweave.domains.entities.canonical.slack_files import SlackFileManifest from airweave.domains.storage.protocols import StorageBackend from airweave.domains.sync_pipeline.processors.entity_fields import populate_base_fields @@ -37,6 +38,10 @@ async def _file_inventory( record: SourceRecord, storage: StorageBackend ) -> tuple[list[SlackFile], SlackFileManifest | None]: """Verify optional acquisition evidence against unchanged native file positions.""" + if record.payload_schema_version not in (1, 2): + raise ValueError("Unsupported Slack message capture schema") + if record.payload_schema_version == 2 and record.blobs: + raise ValueError("Child-owned Slack message cannot retain inline file blobs") try: files = SlackFiles.model_validate(record.payload).files except ValidationError: @@ -74,6 +79,8 @@ async def map_slack_files( parts = [ ProjectionInput(part=ExtractionPart(part_index=0, key="body", kind="body"), entity=body) ] + if record.payload_schema_version == 2: + return ProjectionInputs(parts=tuple(parts)) for index, file in enumerate(files): filename = file.name or file.id suffix = Path(filename).suffix.lower() @@ -98,23 +105,99 @@ async def map_slack_files( raise ValueError("Complete Slack message lacks a retained file") parts.append(ProjectionInput(part=descriptor, entity=None)) continue - content = await read_blob(record, references[0], storage) - path = await write_blob(content, directory, suffix=suffix) - entity = SlackAttachmentEntity( - attachment_key=f"{record.identity.container_id}:{record.identity.native_id}:{file.id}", - filename=filename, - breadcrumbs=[ - Breadcrumb( - entity_id=body.entity_id, name="Slack message", entity_type=type(body).__name__ - ) - ], - # Projection never dereferences URLs; the original is identified by its parent. - url="", - size=len(content), - file_type=suffix.lstrip("."), - mime_type=file.mimetype, - local_path=str(path), + entity = await _materialize_file( + record, + file, + references[0], + storage, + directory, + suffix, + parent_id=body.entity_id, ) - populate_base_fields(entity) parts.append(ProjectionInput(part=descriptor, entity=entity)) return ProjectionInputs(parts=tuple(parts)) + + +async def _materialize_file( + record: SourceRecord, + file: SlackFile, + reference: BlobReference, + storage: StorageBackend, + directory: Path, + suffix: str, + *, + parent_id: str, +) -> SlackAttachmentEntity: + """Verify only this record's bytes and materialize a disposable converter input.""" + content = await read_blob(record, reference, storage) + path = await write_blob(content, directory, suffix=suffix) + entity = SlackAttachmentEntity( + attachment_key=f"{record.identity.container_id}:{record.identity.native_id}:{file.id}", + filename=file.name or file.id, + breadcrumbs=[ + Breadcrumb(entity_id=parent_id, name="Slack message", entity_type="SlackMessageEntity") + ], + url="", + size=len(content), + file_type=suffix.lstrip("."), + mime_type=file.mimetype, + local_path=str(path), + ) + populate_base_fields(entity) + return entity + + +async def map_slack_file( + record: SourceRecord, storage: StorageBackend, directory: Path +) -> ProjectionInputs: + """A durable child is the only extraction owner of its message/file occurrence.""" + parent = record.parent + if ( + record.payload_schema_version != 1 + or parent is None + or parent.record_type != "message" + or not parent.container_id + or record.identity.container_id != parent_container_key(parent) + ): + raise ValueError("Slack file requires an exact message parent") + try: + file = SlackFile.model_validate(record.payload) + except ValidationError: + raise ValueError("Slack child file metadata is malformed") from None + if file.id != record.identity.native_id: + raise ValueError("Slack child native file identity disagrees") + manifests = [blob for blob in record.blobs if blob.role == "representation_manifest"] + references = [blob for blob in record.blobs if blob.role is None] + if ( + len(manifests) != 1 + or len(references) > 1 + or any(blob.source_path != "" for blob in references) + ): + raise ValueError("Slack child requires one acquisition manifest and a root original") + content = await read_blob(record, manifests[0], storage) + try: + manifest = SlackFileManifest.model_validate_json(content) + if len(manifest.files) != 1 or manifest.files[0].native_id != file.id: + raise ValueError("Slack child acquisition manifest identity disagrees") + outcome = manifest.files[0] + if outcome.file is not None: + file = SlackFile.model_validate(outcome.file) + except ValidationError: + raise ValueError("Slack child acquisition manifest is malformed") from None + captured = outcome.outcome == "captured" + if captured != bool(references) or (not captured and record.completeness == "complete"): + raise ValueError("Slack child acquisition evidence contradicts retained bytes") + suffix = Path(file.name or "").suffix.lower() + if not re.fullmatch(r"\.[a-z0-9]{1,12}", suffix): + suffix = mimetypes.guess_extension(file.mimetype or "") or ".bin" + descriptor = ExtractionPart( + part_index=0, key=f"file:{file.id}", kind="file", media_type=file.mimetype, extension=suffix + ) + entity = ( + await _materialize_file( + record, file, references[0], storage, directory, suffix, parent_id=parent.native_id + ) + if captured + else None + ) + return ProjectionInputs(parts=(ProjectionInput(part=descriptor, entity=entity),)) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py index 7c0e9d6451..2192bff390 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py @@ -171,6 +171,78 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source ] 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(), 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 @@ -220,3 +292,83 @@ async def run(value, *, include_original=True): 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/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index 6abdb13b80..797aeccd9c 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -149,3 +149,28 @@ file denial remain partial acquisition evidence; missing OAuth scope fails the configured capture. Downloaded byte-count mismatch fails the page. Capture still commits whole provider pages: aggregate page download/retry cost has not been qualified live, even though individual files and metadata counts are bounded. + +### Durable Slack file children + +The child-owned topology is explicit: message payload schema **2** projects only +its body and must contain no inline file blobs. Its unchanged native `files` list +is inventory, not evidence that downloads completed. Schema **1** messages retain +the legacy inline behavior described above; missing blobs never select topology. +Unknown message schemas fail projection. + +A child `file` record must have its exact message parent, the canonical full-parent +container key, and the matching native file ID in its original payload. The one +original blob uses the empty JSON Pointer (`source_path=""`) because that record's +payload is the native file object itself. A singleton acquisition manifest supplies +captured/unavailable evidence and optional native metadata enrichment. Only the +child owns that file's extraction part. Generic ancestry continues to authorize +its reads, downloads and search; a locator never grants independent file access. + +Switching requires recapture under the new source fingerprint, schema-2 parent +revisions and child acquisition. Parent revision validation hides the old inline +publication before the replacement body-only generation and child generations +publish. No blob is copied back to the parent and no global pipeline version is +changed. Deployment and product reader migration remain separate activation gates. +Tests exercise a real retained PDF transitioning from an inline publication to one +body-only parent, one indexed file child and one unavailable child. Provider HTTP, +real embeddings and live Vespa are not part of this synthetic SQL proof. From e02c9407beb8fa7c8dba3a3588bad29b018b3b22 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 18:52:16 -0700 Subject: [PATCH 179/266] Capture Slack files as resumable message-owned children --- .../tests/test_slack_file_recovery.py | 114 ++++++++ backend/airweave/platform/configs/config.py | 4 +- backend/airweave/platform/sources/slack.py | 124 ++++++++- .../platform/sources/slack_content.py | 49 ++-- .../platform/sources/test_slack_content.py | 259 ++++++++---------- 5 files changed, 375 insertions(+), 175 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py 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..2bbe09a8f2 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py @@ -0,0 +1,114 @@ +"""Actual child frontier commits/retry with synthetic Slack HTTP and original bytes.""" + +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.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.entity import Entity +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): + 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") + connector._get = AsyncMock(side_effect=[ROOT, {"messages": [MESSAGE]}]) + 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" diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index af6a14b195..34f1114763 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -805,8 +805,8 @@ class SlackConfig(SourceConfig): capture_files: bool = Field( default=False, - description="Opt in to original file acquisition before each complete page commits; " - "aggregate page byte and duration budgets remain unqualified.", + 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+$") diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 4b5e6a51ea..09f0c24098 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -8,6 +8,7 @@ from datetime import datetime, timedelta, timezone from functools import partial from typing import Any, AsyncGenerator, Dict, List, Optional +from uuid import UUID from pydantic import ( BaseModel, @@ -36,6 +37,7 @@ CaptureRecord, CompletedScope, RecordIdentity, + parent_container_key, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation from airweave.domains.sources.exceptions import SourceAuthError, SourceError, SourceRateLimitError @@ -58,7 +60,7 @@ ) 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 capture_slack_files +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 @@ -176,6 +178,15 @@ class SlackMessageContinuation(BaseModel): 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", @@ -195,9 +206,28 @@ class SlackMessageContinuation(BaseModel): class SlackSource(BaseSource): """Capture all accessible conversation history and replies without search truncation.""" - canonical_record_types = ("channel", "message") - canonical_container_parents = {"message": "channel"} - slack_config: SlackConfig | None = None + # 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"} + ) + _verified_principal: SlackPrincipal | None = None def _require_principal(self) -> SlackPrincipal: @@ -218,7 +248,7 @@ def capture_cycle_configuration(self) -> CycleConfiguration: 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"] = True + material["capture_files"] = "message_owned_children_v1" fingerprint = hashlib.sha256( json.dumps( material, @@ -307,8 +337,14 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: """Keep the established native container identity for this flat provider.""" + if self.canonical_container_parents.get(record_type) != parent.identity.record_type: + raise ValueError("Unsupported Slack child scope") return CompletedScope( - record_type=record_type, container_id=parent.identity.native_id, parent=parent.identity + record_type=record_type, + container_id=parent_container_key(parent.identity) + if record_type == "file" + else parent.identity.native_id, + parent=parent.identity, ) async def capture_page( @@ -322,17 +358,87 @@ async def capture_page( """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( - [await capture_slack_files(self, record, files) for record in page.records] + for record in page.records: + self._file_inventory(record.payload) + page = page.model_copy( + update={ + "records": tuple( + record.model_copy( + update={ + "payload_schema_version": 2, + "descendant_visibility_fields": ("files",), + "completeness": "complete", + } + ) + for record in page.records + ) + } ) - page = page.model_copy(update={"records": records}) return page except ValidationError: # Validation diagnostics may contain private provider values. raise ValueError("Slack capture returned invalid page or continuation 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: diff --git a/backend/airweave/platform/sources/slack_content.py b/backend/airweave/platform/sources/slack_content.py index 6e544122c1..f49d83c728 100644 --- a/backend/airweave/platform/sources/slack_content.py +++ b/backend/airweave/platform/sources/slack_content.py @@ -67,40 +67,37 @@ def _original_url(file: SlackFileMetadata) -> str | None: return url -async def capture_slack_files( +async def capture_slack_file( source: "SlackSource", record: CaptureRecord, files: FileService ) -> CaptureRecord: - """Whole-page caller commits only after every message acquisition succeeds.""" - metadata = SlackMessageFiles.model_validate(record.payload).files - if not metadata: - return record - if len({file.id for file in metadata}) != len(metadata): - raise ValueError("Slack message contains duplicate file identities") - outcomes = [] + """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 = [] - for index, file in enumerate(metadata): - file, native, reason = await _enrich(source, file) - 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": f"/files/{index}"})) - outcomes.append( + 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=index, + index=0, native_id=file.id, outcome="captured" if reason is None else "unavailable", reason=reason, file=native, - ) + ), ) - manifest = SlackFileManifest(files=tuple(outcomes)) + ) reference = await files.store_canonical_blob( manifest.model_dump_json().encode(), media_type="application/json" ) @@ -108,7 +105,7 @@ async def capture_slack_files( return record.model_copy( update={ "blobs": tuple(blobs), - "completeness": "partial" if any(item.reason for item in outcomes) else "complete", + "completeness": "partial" if reason else "complete", } ) diff --git a/backend/tests/unit/platform/sources/test_slack_content.py b/backend/tests/unit/platform/sources/test_slack_content.py index 8b5cfaef19..6f9ffbb501 100644 --- a/backend/tests/unit/platform/sources/test_slack_content.py +++ b/backend/tests/unit/platform/sources/test_slack_content.py @@ -1,30 +1,58 @@ -"""Whole-page Slack originals: opt-in, immutable provenance and no failed-page progress.""" +"""One-file capture interpretation and retained-inventory continuation boundaries.""" import hashlib +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 import httpx import pytest -from airweave.domains.entities.canonical.requests import BlobReference, CompletedScope +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.page_source import InvalidScanContinuation +from airweave.domains.entities.canonical.requests import ( + BlobReference, + CaptureRecord, + RecordIdentity, +) from airweave.domains.entities.canonical.scan_models import ScanContinuation from airweave.domains.entities.canonical.slack_files import SlackFileManifest from airweave.domains.sources.exceptions import SourceError from airweave.domains.sources.token_providers.static import StaticTokenProvider from airweave.platform.configs.config import SlackConfig -from airweave.platform.sources.slack import SlackPrincipal, SlackSource -from airweave.platform.sources.slack_content import capture_slack_files +from airweave.platform.sources.slack import SlackApiError, SlackPrincipal, SlackSource +from airweave.platform.sources.slack_content import capture_slack_file def connector(): - source = SlackSource( - auth=StaticTokenProvider("token"), logger=MagicMock(), http_client=MagicMock() + result = SlackSource( + auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock() ) - source.slack_config = SlackConfig( + result.slack_config = SlackConfig( expected_team_id="T1", expected_user_id="U1", capture_files=True ) - source._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") - return source + result._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") + return result + + +def native(identity="F1", **changes): + return { + "id": identity, + "size": 3, + "mimetype": "text/plain", + "url_private": f"https://files.slack.com/{identity}", + **changes, + } + + +def original(payload): + return CaptureRecord( + identity=RecordIdentity( + record_type="file", native_id=payload["id"], container_id="parent:fixture" + ), + payload=payload, + observed_at=datetime.now(timezone.utc), + ) def storage(): @@ -40,99 +68,40 @@ async def store(content, *, media_type=None): files.store_canonical_blob = AsyncMock(side_effect=store) files.capture_canonical_url = AsyncMock( - return_value=BlobReference( - key="original", sha256="a" * 64, size_bytes=3, media_type="text/plain" - ) + return_value=BlobReference(key="original", sha256="a" * 64, size_bytes=3) ) return files, saved @pytest.mark.asyncio -async def test_capture_enriches_separately_and_retains_explicit_partial_outcomes(): +async def test_file_enrichment_preserves_raw_child_and_root_blob(): source = connector() - message = { - "ts": "1", - "files": [ - {"id": "F1"}, - {"id": "F2", "is_external": True}, - { - "id": "F3", - "size": 300 * 1024 * 1024, - "mimetype": "text/plain", - "url_private": "https://files.slack.com/large", - }, - ], - } - source._get = AsyncMock( - return_value={ - "ok": True, - "file": { - "id": "F1", - "name": "note.txt", - "mimetype": "text/plain", - "size": 3, - "url_private_download": "https://files.slack.com/original", - }, - } - ) files, saved = storage() - record = await capture_slack_files(source, source._capture_message(message, "C1"), files) - assert record.payload == message and "name" not in record.payload["files"][0] + source._get = AsyncMock(return_value={"file": native(name="note.txt")}) + item = original({"id": "F1"}) + result = await capture_slack_file(source, item, files) + assert result.payload == {"id": "F1"} and result.completeness == "complete" + assert [b.source_path for b in result.blobs] == ["", None] manifest = SlackFileManifest.model_validate_json(saved[-1]) - assert [(x.outcome, x.reason) for x in manifest.files] == [ - ("captured", None), - ("unavailable", "unsupported_external"), - ("unavailable", "oversized"), - ] assert manifest.files[0].file["name"] == "note.txt" - assert record.completeness == "partial" - assert [blob.source_path for blob in record.blobs] == ["/files/0", None] - assert record.blobs[-1].role == "representation_manifest" + assert manifest.files[0].outcome == "captured" assert files.capture_canonical_url.call_args.kwargs["follow_redirects"] is False assert files.capture_canonical_url.call_args.kwargs["expected_media_type"] == "text/plain" @pytest.mark.asyncio @pytest.mark.parametrize("status", [302, 401, 429, 503]) -async def test_late_download_failure_returns_no_page_or_new_continuation(status): +async def test_nonterminal_download_failure_cannot_create_manifest(status): source = connector() files, saved = storage() - source._get = AsyncMock( - return_value={ - "messages": [ - { - "ts": "1", - "files": [ - { - "id": "F1", - "size": 3, - "mimetype": "text/plain", - "url_private": "https://files.slack.com/one", - }, - { - "id": "F2", - "size": 3, - "mimetype": "text/plain", - "url_private": "https://files.slack.com/two", - }, - ], - } - ], - "response_metadata": {"next_cursor": "next"}, - } - ) - error = httpx.HTTPStatusError( + files.capture_canonical_url.side_effect = httpx.HTTPStatusError( "private", request=httpx.Request("GET", "https://files.slack.com/private"), response=httpx.Response(status), ) - files.capture_canonical_url.side_effect = [files.capture_canonical_url.return_value, error] - continuation = ScanContinuation() with pytest.raises(SourceError): - await source.capture_page( - CompletedScope(record_type="message", container_id="C1"), continuation, files=files - ) - assert continuation.value == {} and not saved + await capture_slack_file(source, original(native()), files) + assert not saved @pytest.mark.asyncio @@ -148,86 +117,100 @@ async def test_late_download_failure_returns_no_page_or_new_continuation(status) async def test_untrusted_original_never_downloads(url): source = connector() files, _ = storage() - record = source._capture_message({"ts": "1", "files": [{"id": "F1", "url_private": url}]}, "C1") with pytest.raises(ValueError, match="permitted origin"): - await capture_slack_files(source, record, files) + await capture_slack_file(source, original(native(url_private=url)), files) files.capture_canonical_url.assert_not_called() -def test_capture_files_is_opt_in_and_changes_only_enabled_fingerprint(): +@pytest.mark.asyncio +async def test_explicit_missing_and_denied_metadata_vs_configuration_failure(): source = connector() - enabled = source.capture_cycle_configuration.fingerprint - source.slack_config.capture_files = False - legacy = source.capture_cycle_configuration.fingerprint - import json - - expected = hashlib.sha256( - json.dumps({"version": 2, "team_id": "T1", "user_id": "U1"}, sort_keys=True).encode() - ).hexdigest() - assert legacy == expected and enabled != legacy - assert SlackConfig().capture_files is False + files, saved = storage() + item = original({"id": "F1"}) + source._get = AsyncMock(side_effect=SlackApiError("not_visible")) + result = await capture_slack_file(source, item, files) + assert result.completeness == "partial" + assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "access_denied" + source._get = AsyncMock( + return_value={"file": {"id": "F1", "thumb_360": "https://files.slack.com/thumb"}} + ) + await capture_slack_file(source, item, files) + assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "missing_metadata" + files.capture_canonical_url.assert_not_called() + source._get = AsyncMock(side_effect=SlackApiError("missing_scope")) + with pytest.raises(SlackApiError): + await capture_slack_file(source, item, files) @pytest.mark.asyncio -async def test_size_race_and_metadata_identity_mismatch_do_not_publish_manifest(): +async def test_file_size_and_enriched_identity_must_match(): source = connector() files, saved = storage() - record = source._capture_message( - { - "ts": "1", - "files": [ - { - "id": "F1", - "size": 4, - "mimetype": "text/plain", - "url_private": "https://files.slack.com/original", - } - ], - }, - "C1", - ) with pytest.raises(ValueError, match="size changed"): - await capture_slack_files(source, record, files) - assert not saved - source._get = AsyncMock(return_value={"file": {"id": "F2"}}) - record = source._capture_message({"ts": "1", "files": [{"id": "F1"}]}, "C1") + await capture_slack_file(source, original(native(size=4)), files) + source._get = AsyncMock(return_value={"file": native("F2")}) with pytest.raises(ValueError, match="identity changed"): - await capture_slack_files(source, record, files) + await capture_slack_file(source, original({"id": "F1"}), files) assert not saved @pytest.mark.asyncio -async def test_explicit_file_denial_is_partial_but_missing_scope_fails_page(): - from airweave.platform.sources.slack import SlackApiError +async def test_child_continuation_uses_retained_inventory_and_exact_parent(): + source = connector() + files, _ = storage() + parent = SourceRecord.model_construct( + id=uuid4(), + identity=RecordIdentity(record_type="message", native_id="1", container_id="C1"), + payload_schema_version=2, + payload={"files": [native(), native("F2")]}, + content_access="available", + deleted_at=None, + ) + scope = source.child_scope(parent, "file") + first = await source.capture_page(scope, ScanContinuation(), files=files, parent=parent) + assert not first.final and first.records[0].parent == parent.identity + second = await source.capture_page(scope, first.continuation, files=files, parent=parent) + assert second.final and second.records[0].identity.native_id == "F2" + changed = parent.model_copy(update={"payload": {"files": [native("F2")]}}) + with pytest.raises(InvalidScanContinuation): + await source.capture_page(scope, first.continuation, files=files, parent=changed) + other = parent.model_copy( + update={"identity": parent.identity.model_copy(update={"container_id": "C2"})} + ) + with pytest.raises(ValueError, match="current retained message"): + await source.capture_page(scope, ScanContinuation(), files=files, parent=other) + assert files.capture_canonical_url.await_count == 2 + + +def test_disabled_topology_and_fingerprint_stay_legacy(): + import json source = connector() - files, saved = storage() - record = source._capture_message({"ts": "1", "files": [{"id": "F1"}]}, "C1") - source._get = AsyncMock(side_effect=SlackApiError("not_visible")) - result = await capture_slack_files(source, record, files) - assert result.completeness == "partial" - assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "access_denied" - source._get = AsyncMock(side_effect=SlackApiError("missing_scope")) - files.store_canonical_blob.reset_mock() - with pytest.raises(SlackApiError): - await capture_slack_files(source, record, files) - files.store_canonical_blob.assert_not_called() + assert source.canonical_container_parents["file"] == "message" + enabled = source.capture_cycle_configuration.fingerprint + source.slack_config = source.slack_config.model_copy(update={"capture_files": False}) + expected = hashlib.sha256( + json.dumps({"version": 2, "team_id": "T1", "user_id": "U1"}, sort_keys=True).encode() + ).hexdigest() + assert source.capture_cycle_configuration.fingerprint == expected != enabled + assert source.canonical_record_types == ("channel", "message") + assert SlackConfig().capture_files is False + assert "file" in SlackSource.canonical_record_types @pytest.mark.asyncio -async def test_missing_metadata_remains_partial_and_never_uses_thumbnail(): +@pytest.mark.parametrize( + "metadata,reason", + [ + (native(is_external=True), "unsupported_external"), + (native(size=201 * 1024 * 1024), "oversized"), + ], +) +async def test_terminal_file_omission_is_explicit_child_state(metadata, reason): source = connector() files, saved = storage() - record = source._capture_message( - {"ts": "1", "files": [{"id": "F1", "url_private": "https://files.slack.com/original"}]}, - "C1", - ) - source._get = AsyncMock( - return_value={ - "file": {"id": "F1", "thumb_360": "https://files.slack.com/thumbnail", "size": 3} - } - ) - result = await capture_slack_files(source, record, files) + result = await capture_slack_file(source, original(metadata), files) assert result.completeness == "partial" - assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == "missing_metadata" + assert len(result.blobs) == 1 and result.blobs[0].role == "representation_manifest" + assert SlackFileManifest.model_validate_json(saved[-1]).files[0].reason == reason files.capture_canonical_url.assert_not_called() From 11d6c83cb885a857da5965f356dc3e8ed72b5c1c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 21:43:12 -0700 Subject: [PATCH 180/266] Verify Slack membership with configured capture topology --- .../entities/canonical/tests/test_slack_membership.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py index 09ed529aba..5efad75762 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_membership.py @@ -17,8 +17,9 @@ @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 + database, source, listing_fails, capture_files ): service, fence = source parent = RecordIdentity(record_type="channel", native_id="lost-channel") @@ -40,7 +41,9 @@ async def test_first_completed_membership_hides_prior_partial_channel( connector = SlackSource( auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=MagicMock() ) - connector.slack_config = SlackConfig(expected_team_id="T1", expected_user_id="U1") + 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( @@ -54,9 +57,9 @@ async def test_first_completed_membership_hides_prior_partial_channel( service, database, bus, - SlackSource.canonical_record_types, + connector.canonical_record_types, CaptureAttempt(id=fence.attempt_id, number=fence.attempt_number), - SlackSource.canonical_container_parents, + connector.canonical_container_parents, page_source=connector, files=MagicMock(), ) From 495c6b721891bdd125fbb65ef0807c2502c3862f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Wed, 30 Sep 2026 21:50:18 -0700 Subject: [PATCH 181/266] Preserve native Slack bodies in shared retained text projection --- .../entities/canonical/projection_inputs.py | 2 + .../entities/canonical/projection_mappers.py | 4 + .../domains/entities/canonical/projector.py | 12 +- .../entities/canonical/slack_projection.py | 10 +- .../canonical/tests/test_native_text.py | 156 ++++++++++++++++++ .../canonical/tests/test_projection.py | 2 +- .../entities/canonical/text_artifacts.py | 1 + .../domains/entities/canonical/text_models.py | 20 ++- .../domains/entities/canonical/text_query.py | 2 +- .../sync_pipeline/pipeline/text_builder.py | 61 ++++++- .../sync_pipeline/pipeline/text_models.py | 11 ++ .../sync_pipeline/processors/chunk_embed.py | 8 +- docs/owned-sources/representations.md | 29 +++- 13 files changed, 303 insertions(+), 15 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_native_text.py diff --git a/backend/airweave/domains/entities/canonical/projection_inputs.py b/backend/airweave/domains/entities/canonical/projection_inputs.py index 0d335ecb77..8747b22aa6 100644 --- a/backend/airweave/domains/entities/canonical/projection_inputs.py +++ b/backend/airweave/domains/entities/canonical/projection_inputs.py @@ -3,6 +3,7 @@ from pydantic import BaseModel, ConfigDict 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 @@ -12,6 +13,7 @@ class ProjectionInput(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) part: ExtractionPart entity: BaseEntity | None + native_body: NativeTextBody | None = None class ProjectionInputs(BaseModel): diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 45c7066b1a..885f6aeccd 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -156,7 +156,11 @@ def _slack(record: SourceRecord) -> tuple[BaseEntity, ...]: 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"): diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 765e09dfe5..6e8d0bc4fc 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -107,7 +107,17 @@ async def project_one( ) context = ProjectionContext(logger, source_name) runtime = ProjectionRuntime(StrictProjectionTracker()) - built = await self._processor.build_text(selected, context, runtime) + selected_ids = {entity.entity_id for entity in selected} + native_bodies = { + item.entity.entity_id: item.native_body + for item in mapped.parts + if item.native_body is not None + and item.entity is not None + and item.entity.entity_id in selected_ids + } + built = await self._processor.build_text( + selected, context, runtime, native_bodies=native_bodies + ) artifacts = prepare_text(built.representations, generation) chunks = await self._processor.process_built_text( built.entities, diff --git a/backend/airweave/domains/entities/canonical/slack_projection.py b/backend/airweave/domains/entities/canonical/slack_projection.py index dafdd95ab2..311fab3e54 100644 --- a/backend/airweave/domains/entities/canonical/slack_projection.py +++ b/backend/airweave/domains/entities/canonical/slack_projection.py @@ -13,6 +13,7 @@ from airweave.domains.entities.canonical.requests import BlobReference, parent_container_key from airweave.domains.entities.canonical.slack_files import SlackFileManifest from airweave.domains.storage.protocols import StorageBackend +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._base import BaseEntity, Breadcrumb from airweave.platform.entities.slack import SlackAttachmentEntity @@ -76,8 +77,15 @@ async def map_slack_files( """Retain one body and one expected part per file, including uncaptured originals.""" files, manifest = await _file_inventory(record, storage) populate_base_fields(body) + native_body = None + if "text" in record.payload: + native_body = NativeTextBody(text=record.payload["text"], metadata_fields=("text",)) parts = [ - ProjectionInput(part=ExtractionPart(part_index=0, key="body", kind="body"), entity=body) + ProjectionInput( + part=ExtractionPart(part_index=0, key="body", kind="body"), + entity=body, + native_body=native_body, + ) ] if record.payload_schema_version == 2: return ProjectionInputs(parts=tuple(parts)) 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..9e9caf9ac2 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_native_text.py @@ -0,0 +1,156 @@ +"""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* \n# Content\nend"] +) +async def test_native_body_exact_and_excluded_from_metadata(schema, text): + record = message([]).model_copy( + update={ + "payload": {"ts": "1.000001", "text": text, "blocks": [{"type": "unrendered"}]}, + "payload_schema_version": schema, + } + ) + async with map_record(record, "slack", AsyncMock()) as mapped: + item = mapped.parts[0] + assert item.entity.text == text + batch = await TextualRepresentationBuilder(MagicMock()).build_with_text( + [item.entity], + SimpleNamespace(source_short_name="slack", logger=MagicMock()), + SimpleNamespace(entity_tracker=SimpleNamespace(record_skipped=AsyncMock())), + native_bodies={item.entity.entity_id: item.native_body}, + ) + built = batch.representations[0] + assert built.kind == "native_text" + assert built.text[built.content_start :] == text + assert "**Text**:" not in built.text[: built.content_start] + assert built.text == batch.entities[0].textual_representation + assert "unrendered" not in built.text + + +async def test_missing_text_does_not_invent_native_body_from_blocks_or_title(): + record = message([]).model_copy( + update={ + "payload": {"ts": "1.000001", "blocks": [{"type": "rich_text"}]}, + } + ) + async with map_record(record, "slack", AsyncMock()) as mapped: + assert mapped.parts[0].native_body is None + assert mapped.parts[0].entity.text == "" + + +@pytest.mark.parametrize("text", [None, 123, {}, []]) +async def test_malformed_native_body_fails(text): + record = message([]).model_copy(update={"payload": {"ts": "1.000001", "text": text}}) + with pytest.raises(ValueError, match="Slack message text must be a string"): + async with map_record(record, "slack", AsyncMock()): + pass + + +@pytest.mark.parametrize("start,kind", [(None, "generated_text"), (0, "extracted_text")]) +def test_old_manifest_provenance_is_preserved(start, kind): + values = { + "id": uuid4(), + "part_index": 0, + "sha256": "0" * 64, + "size_bytes": 0, + "characters": 0, + "content_start": start, + } + assert TextArtifact.model_validate(values).kind == kind + for incorrect in ["native_text", "extracted_text"] if start is None else ["generated_text"]: + with pytest.raises(ValidationError, match="provenance"): + TextArtifact.model_validate({**values, "kind": incorrect}) + + +@pytest.mark.parametrize("body", [None, "", "Native 👋 text\n# Content\n<@U1>"]) +async def test_slack_native_retained_read_matches_text_used_for_index( + database, source, tmp_path, body +): + from sqlalchemy import update + + from airweave.adapters.storage.filesystem import FilesystemBackend + from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore + 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 capture, observation + from airweave.domains.entities.canonical.tests.test_extraction_coverage import ( + destination, + projector, + ) + from airweave.domains.entities.canonical.text_query import CanonicalTextReader, TextUnavailable + from airweave.models.sync import Sync + + service, fence = source + original = observation( + identity=RecordIdentity(record_type="message", native_id="1.000001", container_id="C1"), + payload={"ts": "1.000001", **({"text": body} if body is not None else {})}, + payload_schema_version=2, + ) + await capture(database, service, fence, original) + 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, "slack", destination(), MagicMock()) + reader = CanonicalTextReader( + CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "key"), + storage, + ) + async with database() as db: + ref = ( + await reader.list(db, fence.organization_id, fence.sync_id, work.record.id, 1) + ).representations[0] + if body is None: + assert ref.kind == "generated_text" and ref.content_characters is None + else: + assert ref.kind == "native_text" and ref.content_characters == len(body) + args = (fence.organization_id, fence.sync_id, work.record.id, 1, ref.generation, ref.id) + if body is None: + with pytest.raises(TextUnavailable, match="generated search text"): + await reader.read(db, *args) + else: + content = await reader.read(db, *args) + assert content.text == body and content.total_characters == len(body) + index = await reader.read(db, *args, view="index") + assert index.text == indexed[0] + if body is not None: + assert index.text.endswith(body) + if body: + first = await reader.read(db, *args, limit=8) + rest = await reader.read(db, *args, offset=first.next_offset) + assert first.text + rest.text == body + await db.execute( + update(Sync).where(Sync.id == fence.sync_id).values(index_pipeline_version=2) + ) + await db.commit() + with pytest.raises(TextUnavailable): + await reader.read(db, *args) diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py index 692d97f715..29d24a8859 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py @@ -141,7 +141,7 @@ async def process(entities, context, runtime, *, strict, expected_ids): from airweave.domains.sync_pipeline.pipeline.text_models import BuiltText, BuiltTextBatch - async def build_text(entities, context, runtime): + async def build_text(entities, context, runtime, *, native_bodies=None): return BuiltTextBatch( entities=entities, representations=tuple( diff --git a/backend/airweave/domains/entities/canonical/text_artifacts.py b/backend/airweave/domains/entities/canonical/text_artifacts.py index 3d4622beae..0d4055f0e8 100644 --- a/backend/airweave/domains/entities/canonical/text_artifacts.py +++ b/backend/airweave/domains/entities/canonical/text_artifacts.py @@ -29,6 +29,7 @@ def prepare_text( size_bytes=len(content), characters=len(item.text), content_start=item.content_start, + kind=item.kind, ) artifacts.append((artifact, content)) return tuple(artifacts) diff --git a/backend/airweave/domains/entities/canonical/text_models.py b/backend/airweave/domains/entities/canonical/text_models.py index 429afd1bae..2a9afa2617 100644 --- a/backend/airweave/domains/entities/canonical/text_models.py +++ b/backend/airweave/domains/entities/canonical/text_models.py @@ -16,10 +16,26 @@ class TextArtifact(BaseModel): 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 converter boundary must lie inside the retained indexing text.""" + """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 @@ -46,7 +62,7 @@ class TextRepresentation(BaseModel): generation: UUID pipeline_version: int part_key: str - kind: Literal["extracted_text", "generated_text"] + kind: Literal["native_text", "extracted_text", "generated_text"] media_type: Literal["text/markdown"] = "text/markdown" content_characters: int | None index_characters: int diff --git a/backend/airweave/domains/entities/canonical/text_query.py b/backend/airweave/domains/entities/canonical/text_query.py index 44f2b156af..76d406ce2e 100644 --- a/backend/airweave/domains/entities/canonical/text_query.py +++ b/backend/airweave/domains/entities/canonical/text_query.py @@ -128,7 +128,7 @@ def _describe(row: ProjectionGeneration, artifact: TextArtifact) -> TextRepresen generation=row.id, pipeline_version=row.pipeline_version, part_key=part.key, - kind="extracted_text" if artifact.content_start is not None else "generated_text", + kind=artifact.kind, content_characters=( artifact.characters - artifact.content_start if artifact.content_start is not None diff --git a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py index 9c3440105a..dc801b0bea 100644 --- a/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py +++ b/backend/airweave/domains/sync_pipeline/pipeline/text_builder.py @@ -9,7 +9,11 @@ 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 +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: @@ -70,13 +74,29 @@ async def build_with_text( 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) @@ -97,6 +117,13 @@ async def build_with_text( 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 ), @@ -106,26 +133,43 @@ async def build_with_text( # 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 @@ -136,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 @@ -159,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)) diff --git a/backend/airweave/domains/sync_pipeline/pipeline/text_models.py b/backend/airweave/domains/sync_pipeline/pipeline/text_models.py index 45e5cc3b09..3f2e0d3c5e 100644 --- a/backend/airweave/domains/sync_pipeline/pipeline/text_models.py +++ b/backend/airweave/domains/sync_pipeline/pipeline/text_models.py @@ -1,10 +1,20 @@ """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.""" @@ -12,6 +22,7 @@ class BuiltText(BaseModel): 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): diff --git a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py index 024f88f5aa..af0f96ac56 100644 --- a/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py +++ b/backend/airweave/domains/sync_pipeline/processors/chunk_embed.py @@ -19,7 +19,7 @@ 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 +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 @@ -77,9 +77,13 @@ async def build_text( 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) + return await self._text_builder.build_with_text( + entities, sync_context, runtime, native_bodies=native_bodies + ) async def process_built_text( self, diff --git a/docs/owned-sources/representations.md b/docs/owned-sources/representations.md index 797aeccd9c..74b0c63eda 100644 --- a/docs/owned-sources/representations.md +++ b/docs/owned-sources/representations.md @@ -36,7 +36,7 @@ Synthetic provider/PDF and migrated SQL tests are separate from live qualificati returns `record_id`, `revision`, `status` (`available` or `unavailable`) and `representations`. Each descriptor contains an opaque `id`, parent `record_id`, `revision`, `generation`, `pipeline_version`, source-local `part_key`, -`kind` (`extracted_text` or `generated_text`), `media_type: text/markdown`, +`kind` (`native_text`, `extracted_text`, or `generated_text`), `media_type: text/markdown`, `content_characters` (null for generated-only text), `index_characters`, and `source_anchors: unavailable`. Storage keys and internal part ordinals are omitted. @@ -174,3 +174,30 @@ changed. Deployment and product reader migration remain separate activation gate Tests exercise a real retained PDF transitioning from an inline publication to one body-only parent, one indexed file child and one unavailable child. Provider HTTP, real embeddings and live Vespa are not part of this synthetic SQL proof. + +### Native Slack message text + +Slack message projections select the captured `/text` string as `native_text`. +The same builder appends that string to generated search context and records its +exact character boundary. Retention and indexing consume that single result; +content reads return the string verbatim, without the generated context. An +explicit empty string is available empty content. A missing field has no native +content view; a present non-string field fails projection. Generated display +names are never substituted for native content. + +This is the `/text` representation, not a rendering of Slack blocks or legacy +attachments. Native formatting markers remain intact; block-only messages do not +claim extracted block text. Files keep their independent extraction parts (capture +schema 1) or child records (schema 2). Other providers require their own explicit +body selection; field names alone do not establish native provenance. + +Artifacts now record native, extracted, or generated text provenance in their +existing generation JSON. Older descriptors without provenance retain their +prior interpretation: content boundary means extracted text, otherwise generated +text. They are not retroactively reclassified as native. + +Deploying this code does not rebuild existing publications. Use the existing +scoped reprojection plan/apply operation to increment the selected sync pipeline +version and replay retained originals. The version fence hides old publications +until their replacement is published. No provider recapture is required, and no +global runtime activation or version mutation is included in this change. From 0a09f43294957d10b2cc9be1c1397f950d2bfea4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 01:11:10 -0700 Subject: [PATCH 182/266] Add typed Cohere embedding adapter with explicit query intent --- backend/airweave/core/container/factory.py | 8 ++ .../domains/embedders/dense/cohere.py | 107 ++++++++++++++++++ .../airweave/domains/embedders/dense/local.py | 6 +- .../domains/embedders/dense/mistral.py | 6 +- .../domains/embedders/dense/openai.py | 6 +- .../embedders/dense/tests/test_cohere.py | 99 ++++++++++++++++ .../domains/embedders/fakes/embedder.py | 6 +- .../airweave/domains/embedders/protocols.py | 7 +- .../domains/embedders/registry_data.py | 22 ++++ backend/airweave/domains/embedders/types.py | 4 + backend/airweave/domains/search/executor.py | 2 +- .../domains/search/tests/test_executor.py | 4 + .../airweave/search/operations/embed_query.py | 2 +- docs/owned-sources/cohere-embeddings.md | 32 ++++++ 14 files changed, 299 insertions(+), 12 deletions(-) create mode 100644 backend/airweave/domains/embedders/dense/cohere.py create mode 100644 backend/airweave/domains/embedders/dense/tests/test_cohere.py create mode 100644 docs/owned-sources/cohere-embeddings.md diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 21cf55c676..fedda80239 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -803,12 +803,20 @@ def _create_dense_embedder( Uses the domain config constants (DENSE_EMBEDDER, EMBEDDING_DIMENSIONS) and the registry to look up the spec and construct the correct embedder. """ + from airweave.domains.embedders.dense.cohere import CohereDenseEmbedder from airweave.domains.embedders.dense.local import LocalDenseEmbedder from airweave.domains.embedders.dense.mistral import MistralDenseEmbedder from airweave.domains.embedders.dense.openai import OpenAIDenseEmbedder spec = registry.get(DENSE_EMBEDDER) + if spec.embedder_class_ref is CohereDenseEmbedder: + return CohereDenseEmbedder( + api_key=settings.COHERE_API_KEY, + model=spec.api_model_name, + dimensions=EMBEDDING_DIMENSIONS, + ) + if spec.embedder_class_ref is OpenAIDenseEmbedder: return OpenAIDenseEmbedder( api_key=settings.OPENAI_API_KEY, diff --git a/backend/airweave/domains/embedders/dense/cohere.py b/backend/airweave/domains/embedders/dense/cohere.py new file mode 100644 index 0000000000..9c16da6898 --- /dev/null +++ b/backend/airweave/domains/embedders/dense/cohere.py @@ -0,0 +1,107 @@ +"""Cohere text embeddings with explicit retrieval intent and bounded requests.""" + +import math + +import cohere +import httpx +from cohere.core.api_error import ApiError +from cohere.types import EmbedByTypeResponse + +from airweave.domains.embedders.exceptions import ( + EmbedderAuthError, + EmbedderConfigError, + EmbedderConnectionError, + EmbedderDimensionError, + EmbedderInputError, + EmbedderProviderError, + EmbedderRateLimitError, + EmbedderResponseError, + EmbedderTimeoutError, +) +from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose + + +class CohereDenseEmbedder: + """Use the existing SDK; callers own retries and model migration policy.""" + + def __init__(self, *, api_key: str, model: str, dimensions: int) -> None: + """Fail configuration before creating network resources.""" + allowed = { + "embed-v4.0": {256, 512, 1024, 1536}, + "embed-v5.0-pro": {256, 512, 768, 1024, 1536, 2048}, + "embed-v5.0-fast": {256, 512, 768, 1024, 1536, 2048}, + } + if not api_key or model not in allowed or dimensions not in allowed[model]: + raise EmbedderConfigError("Cohere requires an API key, supported model and dimension") + self._model = model + self._dimensions = dimensions + self._http = httpx.AsyncClient(timeout=60) + self._client = cohere.AsyncClientV2(api_key=api_key, httpx_client=self._http) + + @property + def model_name(self) -> str: + """Persist this identity alongside vector dimensions.""" + return self._model + + @property + def dimensions(self) -> int: + """Configured output width.""" + return self._dimensions + + async def embed(self, text: str) -> DenseEmbedding: + """Embed a document; query callers must use the explicit batch purpose.""" + return (await self.embed_many([text]))[0] + + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: + """Keep order, reject silent truncation, and avoid unbounded request fan-out.""" + if purpose not in ("document", "query") or any(not text.strip() for text in texts): + raise EmbedderInputError("Cohere requires nonblank text and a valid embedding purpose") + results = [] + for offset in range(0, len(texts), 96): + batch = texts[offset : offset + 96] + response = await self._request(batch, purpose) + vectors = response.embeddings.float_ + if vectors is None or len(vectors) != len(batch): + raise EmbedderResponseError("Cohere returned an unexpected vector count") + for vector in vectors: + if len(vector) != self._dimensions: + raise EmbedderDimensionError(expected=self._dimensions, actual=len(vector)) + if not all(math.isfinite(value) for value in vector): + raise EmbedderResponseError("Cohere returned a non-finite vector") + results.append(DenseEmbedding(vector=vector)) + return results + + async def _request( + self, batch: list[str], purpose: EmbeddingPurpose + ) -> EmbedByTypeResponse: + """Translate provider errors without exposing request text or credentials.""" + try: + return await self._client.embed( + model=self._model, + texts=batch, + input_type="search_query" if purpose == "query" else "search_document", + output_dimension=self._dimensions, + embedding_types=["float"], + truncate="NONE", + request_options={"max_retries": 0}, + ) + except ApiError as exc: + if exc.status_code in (401, 403): + raise EmbedderAuthError("Cohere authentication failed", provider="cohere") from exc + if exc.status_code == 429: + raise EmbedderRateLimitError("Cohere rate limited", provider="cohere") from exc + raise EmbedderProviderError( + "Cohere embedding request failed", + provider="cohere", + retryable=exc.status_code is not None and exc.status_code >= 500, + ) from exc + except httpx.TimeoutException as exc: + raise EmbedderTimeoutError(provider="cohere") from exc + except httpx.RequestError as exc: + raise EmbedderConnectionError(provider="cohere") from exc + + async def close(self) -> None: + """Release the transport owned by this adapter.""" + await self._http.aclose() diff --git a/backend/airweave/domains/embedders/dense/local.py b/backend/airweave/domains/embedders/dense/local.py index 2acd61ee02..d66b96bb87 100644 --- a/backend/airweave/domains/embedders/dense/local.py +++ b/backend/airweave/domains/embedders/dense/local.py @@ -18,7 +18,7 @@ EmbedderResponseError, EmbedderTimeoutError, ) -from airweave.domains.embedders.types import DenseEmbedding +from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose _PROVIDER = "local" @@ -75,7 +75,9 @@ async def embed(self, text: str) -> DenseEmbedding: results = await self.embed_many([text]) return results[0] - async def embed_many(self, texts: list[str]) -> list[DenseEmbedding]: + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: """Embed a batch of texts into dense vectors.""" if not texts: return [] diff --git a/backend/airweave/domains/embedders/dense/mistral.py b/backend/airweave/domains/embedders/dense/mistral.py index 8bde0a785d..76ade773ac 100644 --- a/backend/airweave/domains/embedders/dense/mistral.py +++ b/backend/airweave/domains/embedders/dense/mistral.py @@ -21,7 +21,7 @@ EmbedderTimeoutError, ) from airweave.domains.embedders.protocols import DenseEmbedderProtocol -from airweave.domains.embedders.types import DenseEmbedding +from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose from airweave.platform.tokenizers import get_tokenizer _PROVIDER = "mistral" @@ -85,7 +85,9 @@ async def embed(self, text: str) -> DenseEmbedding: results = await self.embed_many([text]) return results[0] - async def embed_many(self, texts: list[str]) -> list[DenseEmbedding]: + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: """Embed a batch of texts into dense vectors.""" if not texts: return [] diff --git a/backend/airweave/domains/embedders/dense/openai.py b/backend/airweave/domains/embedders/dense/openai.py index 0fd574fe7c..355f2c0476 100644 --- a/backend/airweave/domains/embedders/dense/openai.py +++ b/backend/airweave/domains/embedders/dense/openai.py @@ -21,7 +21,7 @@ EmbedderTimeoutError, ) from airweave.domains.embedders.protocols import DenseEmbedderProtocol -from airweave.domains.embedders.types import DenseEmbedding +from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose _PROVIDER = "openai" @@ -84,7 +84,9 @@ async def embed(self, text: str) -> DenseEmbedding: results = await self.embed_many([text]) return results[0] - async def embed_many(self, texts: list[str]) -> list[DenseEmbedding]: + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: """Embed a batch of texts into dense vectors.""" if not texts: return [] diff --git a/backend/airweave/domains/embedders/dense/tests/test_cohere.py b/backend/airweave/domains/embedders/dense/tests/test_cohere.py new file mode 100644 index 0000000000..a40283cd79 --- /dev/null +++ b/backend/airweave/domains/embedders/dense/tests/test_cohere.py @@ -0,0 +1,99 @@ +"""Cohere boundary: real SDK serialization over a mocked HTTP transport.""" + +import httpx +import pytest + +from airweave.domains.embedders.dense.cohere import CohereDenseEmbedder +from airweave.domains.embedders.exceptions import ( + EmbedderAuthError, + EmbedderDimensionError, + EmbedderRateLimitError, + EmbedderResponseError, +) + + +@pytest.fixture +def build(monkeypatch): + original = httpx.AsyncClient + + def factory(handler): + monkeypatch.setattr( + httpx, + "AsyncClient", + lambda **kwargs: original(transport=httpx.MockTransport(handler), **kwargs), + ) + return CohereDenseEmbedder(api_key="fixture", model="embed-v5.0-pro", dimensions=256) + + return factory + + +def response(vectors): + return httpx.Response( + 200, + json={ + "id": "test", + "texts": [], + "embeddings": {"float": vectors}, + "meta": {"api_version": {"version": "2"}}, + }, + ) + + +async def test_intent_batch_order_and_no_truncation(build): + import json + + calls = [] + + def handler(request): + body = json.loads(request.content) + calls.append(body) + return response([[float(text.split(":")[0])] * 256 for text in body["texts"]]) + + embedder = build(handler) + try: + texts = [f"{i}:नमस्ते hello" for i in range(97)] + vectors = await embedder.embed_many(texts) + await embedder.embed_many(["0:नमस्ते"], purpose="query") + assert [v.vector[0] for v in vectors] == list(range(97)) + assert [len(call["texts"]) for call in calls] == [96, 1, 1] + assert [call["input_type"] for call in calls] == [ + "search_document", + "search_document", + "search_query", + ] + assert all(call["truncate"] == "NONE" for call in calls) + assert calls[0]["texts"][0] == texts[0] + assert await embedder.embed_many([]) == [] + assert len(calls) == 3 + finally: + await embedder.close() + + +@pytest.mark.parametrize("status,error", [(401, EmbedderAuthError), (429, EmbedderRateLimitError)]) +async def test_safe_errors_and_no_hidden_retries(build, status, error): + calls = [] + + def handler(request): + calls.append(request) + return httpx.Response(status, json={"message": "private provider detail"}) + + embedder = build(handler) + try: + with pytest.raises(error) as exc: + await embedder.embed("private document") + assert "private" not in str(exc.value) + assert len(calls) == 1 + finally: + await embedder.close() + + +@pytest.mark.parametrize( + "vectors,error", [([], EmbedderResponseError), ([[1.0]], EmbedderDimensionError)] +) +async def test_reject_corrupt_response(build, vectors, error): + embedder = build(lambda request: response(vectors)) + try: + with pytest.raises(error): + await embedder.embed("content") + finally: + await embedder.close() diff --git a/backend/airweave/domains/embedders/fakes/embedder.py b/backend/airweave/domains/embedders/fakes/embedder.py index 20e72b743d..2599f91373 100644 --- a/backend/airweave/domains/embedders/fakes/embedder.py +++ b/backend/airweave/domains/embedders/fakes/embedder.py @@ -1,6 +1,6 @@ """Fake embedder implementations for testing.""" -from airweave.domains.embedders.types import DenseEmbedding, SparseEmbedding +from airweave.domains.embedders.types import DenseEmbedding, EmbeddingPurpose, SparseEmbedding class FakeDenseEmbedder: @@ -44,7 +44,9 @@ async def embed(self, text: str) -> DenseEmbedding: self._check_error() return DenseEmbedding(vector=[0.0] * self._dimensions) - async def embed_many(self, texts: list[str]) -> list[DenseEmbedding]: + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: """Return zero-vectors for each text.""" self._check_error() return [DenseEmbedding(vector=[0.0] * self._dimensions) for _ in texts] diff --git a/backend/airweave/domains/embedders/protocols.py b/backend/airweave/domains/embedders/protocols.py index 7a71aa3db6..becb303b2b 100644 --- a/backend/airweave/domains/embedders/protocols.py +++ b/backend/airweave/domains/embedders/protocols.py @@ -6,6 +6,7 @@ from airweave.domains.embedders.types import ( DenseEmbedderEntry, DenseEmbedding, + EmbeddingPurpose, SparseEmbedderEntry, SparseEmbedding, ) @@ -32,8 +33,10 @@ async def embed(self, text: str) -> DenseEmbedding: """Embed a single text into a dense vector.""" ... - async def embed_many(self, texts: list[str]) -> list[DenseEmbedding]: - """Embed a batch of texts into dense vectors.""" + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: + """Embed ordered texts; asymmetric models distinguish queries from documents.""" ... async def close(self) -> None: diff --git a/backend/airweave/domains/embedders/registry_data.py b/backend/airweave/domains/embedders/registry_data.py index 5ca56de30e..cbcb5a0d61 100644 --- a/backend/airweave/domains/embedders/registry_data.py +++ b/backend/airweave/domains/embedders/registry_data.py @@ -6,6 +6,7 @@ from dataclasses import dataclass +from airweave.domains.embedders.dense.cohere import CohereDenseEmbedder from airweave.domains.embedders.dense.local import LocalDenseEmbedder from airweave.domains.embedders.dense.mistral import MistralDenseEmbedder from airweave.domains.embedders.dense.openai import OpenAIDenseEmbedder @@ -92,6 +93,27 @@ class SparseEmbedderSpec: ), ] +# Text-only registrations; image processing requires a separate qualified input path. +DENSE_EMBEDDERS.extend( + DenseEmbedderSpec( + short_name=short_name, + name=f"Cohere {model}", + description="Multilingual embeddings with explicit query/document purpose", + provider="cohere", + api_model_name=model, + max_dimensions=dimensions, + max_tokens=128000, + supports_matryoshka=True, + embedder_class=CohereDenseEmbedder, + required_setting="COHERE_API_KEY", + ) + for short_name, model, dimensions in ( + ("cohere_embed_v4", "embed-v4.0", 1536), + ("cohere_embed_v5_pro", "embed-v5.0-pro", 2048), + ("cohere_embed_v5_fast", "embed-v5.0-fast", 2048), + ) +) + SPARSE_EMBEDDERS: list[SparseEmbedderSpec] = [ SparseEmbedderSpec( short_name="fastembed_bm25", diff --git a/backend/airweave/domains/embedders/types.py b/backend/airweave/domains/embedders/types.py index bd16b262f8..5ec49c877b 100644 --- a/backend/airweave/domains/embedders/types.py +++ b/backend/airweave/domains/embedders/types.py @@ -1,9 +1,13 @@ """Types for the embedders domain.""" +from typing import Literal + from pydantic import BaseModel, Field from airweave.core.protocols.registry import BaseRegistryEntry +EmbeddingPurpose = Literal["document", "query"] + # --------------------------------------------------------------------------- # Embedding value types # --------------------------------------------------------------------------- diff --git a/backend/airweave/domains/search/executor.py b/backend/airweave/domains/search/executor.py index 37b1ddd248..f94d05847a 100644 --- a/backend/airweave/domains/search/executor.py +++ b/backend/airweave/domains/search/executor.py @@ -200,7 +200,7 @@ async def prepare_query(self, plan: SearchPlan) -> PreparedQueryEmbeddings: 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, diff --git a/backend/airweave/domains/search/tests/test_executor.py b/backend/airweave/domains/search/tests/test_executor.py index b58486a4e6..b9d521f747 100644 --- a/backend/airweave/domains/search/tests/test_executor.py +++ b/backend/airweave/domains/search/tests/test_executor.py @@ -1256,6 +1256,10 @@ async def test_prepared_query_reuses_only_embeddings_across_collection_scopes(mo ) 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"] 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/docs/owned-sources/cohere-embeddings.md b/docs/owned-sources/cohere-embeddings.md new file mode 100644 index 0000000000..c67cfc2753 --- /dev/null +++ b/docs/owned-sources/cohere-embeddings.md @@ -0,0 +1,32 @@ +# Cohere embedding adapter + +2026-10-01. Implemented; SDK/HTTP transport tests only. No live inference, +quality benchmark, deployment change or existing-index migration claimed. + +The existing dense embedder port now accepts an explicit batch purpose: +`document` (default for indexing) or `query`. Both current and legacy query +execution paths pass `query`; symmetric existing providers retain their behavior. +Cohere maps purpose to search_document/search_query. It uses the installed SDK, +requests float vectors, rejects truncation and validates count, dimensions and +finite values. Requests are sequential batches of at most96 texts with a60-second +transport timeout and SDK retries disabled; caller orchestration owns retries. +Provider error text is not copied into public exceptions. + +Registry choices: cohere_embed_v4, cohere_embed_v5_pro, cohere_embed_v5_fast. +Configure COHERE_API_KEY, DENSE_EMBEDDER and EMBEDDING_DIMENSIONS explicitly. +The adapter validates the discrete dimensions supported by each model. These are +text registrations, not image-processing or cross-model Pro/Fast query support. +Use the existing deployment metadata mismatch guard; do not overwrite a MiniLM +index with Cohere vectors. Create an isolated model-qualified evaluation index +before changing runtime configuration. Retain the old runtime for comparison. + +Official references: +- https://cohere.com/blog/embed-5 +- https://docs.cohere.com/v2/docs/semantic-search-embed + +The standalone import API remains to be designed and implemented: current source +provisioning requires a verified managed-provider identity. Wiki/session imports +must instead bind a trusted publisher and scope without pretending to be Composio +accounts. Retained revisions, fenced writes, indexing work and read authorization +should reuse existing store machinery. A new public write endpoint must not expose +writer fences or let callers nominate arbitrary organizations/principals. From 4b29b18abeed0a77a6f8e3dc9dc9066d0f84e50d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 01:48:20 -0700 Subject: [PATCH 183/266] Add isolated retrieval relevance and latency evaluation --- evaluation/README.md | 114 +++++++++++ evaluation/__init__.py | 1 + evaluation/retrieval.py | 312 +++++++++++++++++++++++++++++ evaluation/tests/test_retrieval.py | 156 +++++++++++++++ 4 files changed, 583 insertions(+) create mode 100644 evaluation/README.md create mode 100644 evaluation/__init__.py create mode 100644 evaluation/retrieval.py create mode 100644 evaluation/tests/test_retrieval.py diff --git a/evaluation/README.md b/evaluation/README.md new file mode 100644 index 0000000000..03811c2522 --- /dev/null +++ b/evaluation/README.md @@ -0,0 +1,114 @@ +# Retrieval evaluation + +Offline evaluation of the **delivered record ranking**, with separately observed +latency samples. This does not run sync, query a provider, or change an index. +The production service does not import this package. `uv` installs the pinned +evaluation dependencies separately; no additional production infrastructure. + +```sh +uv run evaluation/retrieval.py /private/path/dataset.json +# Prints the dataset fingerprint. Record it in the corresponding run file. +uv run evaluation/retrieval.py /private/path/dataset.json \ + --run /private/path/run.json --output /private/path/report.json + +uv run --no-project --with ir-measures==0.4.3 --with pydantic==2.11.9 \ + --with pytest python -m pytest evaluation/tests -q --confcutdir=evaluation +``` + +Reports are created with mode 0600 and refuse to overwrite files. Keep real +queries, judgments, record IDs and response captures out of Git. Reports omit +query text but can still reveal sensitive labels through query IDs and tags. +Validation errors can echo input: do not publish raw error logs from private runs. + +## Dataset + +```json +{ + "corpus_id": "frozen-corpus-manifest-sha256", + "version": "judgments-v1", + "queries": [{ + "id": "meeting-hindi", + "text": "परियोजना बैठक कब है?", + "tags": ["language:hi", "intent:known-item", "split:holdout"], + "expectation": "relevant_records", + "expectation_assessor": "human", + "judgments": [ + {"record_id": "calendar:meeting-1", "relevance": 3, "assessor": "human"}, + {"record_id": "gmail:promotion-2", "relevance": 0, "assessor": "human"} + ] + }] +} +``` + +Use stable record identities, not chunk IDs. The corpus manifest should pin the +record revisions, permissions/scopes, extraction and index versions. Its content +is the operator's responsibility: this offline scorer verifies the supplied +corpus ID and exact dataset fingerprint, **not the live index's contents**. +Include applied account/date/type filters in the frozen experiment specification; +do not compare runs using different filters as ranking improvements. + +Grades: 0 irrelevant, 1 marginally useful, 2 useful, 3 directly satisfies intent. +Binary precision, recall and reciprocal rank count grades >= 1 as relevant. +`ir-measures` supplies nDCG with its default linear graded gain, RR, P and R; +we do not maintain custom implementations of that scoring math. + +Every known-answer query needs a positive judgment. `no_answer` queries explicitly +assert no relevant answer in this corpus and have no positive judgments. Their +empty-success outcome is reported separately, not averaged into nDCG/recall. +An error returning nothing is not a correct abstention. No-answer labels need +evidence; merely failing to find an item does not establish absence. + +## Run + +```json +{ + "system": "commit+embedding+retrieval+reranker+configuration", + "dataset_sha256": "", + "corpus_id": "frozen-corpus-manifest-sha256", + "results": [{ + "query_id": "meeting-hindi", + "status": "success", + "record_ids": ["gmail:promotion-2", "calendar:meeting-1"] + }], + "timings": [{ + "request_id": "trial-001", + "query_id": "meeting-hindi", + "phase": "end_to_end", + "condition": "unknown", + "status": "success", + "duration_ms": 812.4 + }] +} +``` + +The array order is authoritative. Unique descending synthetic scores preserve it +when calling the metrics library; scores from different retrieval engines are +not compared. Duplicate record IDs are rejected rather than silently repaired. +Include every query exactly once, recording errors/timeouts with empty results. +Partial responses keep their returned ranking and an explicit `partial` status. +For repeated ranking trials, save separate runs instead of selecting the best. + +Latency may have repeated samples per query, with unique request IDs. Supported +phases: end_to_end, retrieval, reranking, read. Measure with a monotonic clock; +record the actual condition as cold/warm/unknown. First request does not prove +cold caches. Timing groups preserve phase, cache condition and outcome, including +errors/timeouts. No timings are inferred from relevance results. End-to-end HTTP +latency does not establish browser paint latency. Percentiles use nearest rank; +sample count and maximum accompany p50/p95. Do not draw tail conclusions from a +handful of measurements or mix concurrency/load conditions in one run. + +## What constitutes improvement + +Compare the same frozen corpus, scopes, labels and held-out queries. Pool records +from candidate systems and judge previously unjudged records before concluding +one wins. Reports expose per-query top-10 judged/unjudged counts and assessor +provenance. Missing judgments receive no relevance credit; they are not relabeled +irrelevant. Recall means recall of **judged relevant records**, not account-wide +recall. Failed answerable queries contribute zero instead of disappearing from +the denominator. Inspect per-query and tagged regressions alongside aggregates. + +Real-corpus collection, request timing instrumentation, reviewed judgments, +snippet evaluation and corpus-freezing automation are still separate work. +The four focused synthetic tests verify evaluator behavior, not search quality. + +Reference: [ir-measures interfaces and formats](https://ir-measur.es/en/latest/getting-started.html). diff --git a/evaluation/__init__.py b/evaluation/__init__.py new file mode 100644 index 0000000000..a2cce858a4 --- /dev/null +++ b/evaluation/__init__.py @@ -0,0 +1 @@ +"""Offline evaluation tools; not imported by the service runtime.""" diff --git a/evaluation/retrieval.py b/evaluation/retrieval.py new file mode 100644 index 0000000000..22f4c4f9e3 --- /dev/null +++ b/evaluation/retrieval.py @@ -0,0 +1,312 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = ["ir-measures==0.4.3", "pydantic==2.11.9"] +# /// +"""Offline relevance and latency reports, independent of production credentials.""" + +import argparse +import hashlib +import json +import math +import os +from collections import defaultdict +from pathlib import Path +from statistics import mean +from typing import Annotated, Literal + +import ir_measures +from pydantic import BaseModel, ConfigDict, Field, model_validator + +Identifier = Annotated[str, Field(min_length=1, pattern=r"^\S+$")] +Milliseconds = Annotated[float, Field(ge=0, allow_inf_nan=False)] + + +class Value(BaseModel): + """Strict immutable file boundary.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + +class Judgment(Value): + """Record-level relevance; zero is judged irrelevant, absence is unjudged.""" + + record_id: Identifier + relevance: int = Field(ge=0, le=3, strict=True) + assessor: Literal["human", "agent", "synthetic"] + + +class Query(Value): + """A fixed intent with explicit known-answer or no-answer expectation.""" + + id: Identifier + text: str = Field(min_length=1) + tags: tuple[Identifier, ...] = () + expectation: Literal["relevant_records", "no_answer"] + expectation_assessor: Literal["human", "agent", "synthetic"] + judgments: tuple[Judgment, ...] + + @model_validator(mode="after") + def validate_judgments(self): + """Never silently treat a query with no positive judgments as an answerable query.""" + ids = [j.record_id for j in self.judgments] + if len(ids) != len(set(ids)): + raise ValueError("Duplicate record judgment") + positive = any(j.relevance > 0 for j in self.judgments) + if positive != (self.expectation == "relevant_records"): + raise ValueError("Query expectation conflicts with positive judgments") + if not self.text.strip() or len(self.tags) != len(set(self.tags)): + raise ValueError("Query text must be nonblank and tags unique") + return self + + +class Dataset(Value): + """Versioned corpus and labels; real content belongs in a private operator directory.""" + + corpus_id: Identifier + version: Identifier + queries: tuple[Query, ...] = Field(min_length=1) + + @model_validator(mode="after") + def unique_queries(self): + """Prevent duplicate intents from silently weighting the aggregate.""" + if len({q.id for q in self.queries}) != len(self.queries): + raise ValueError("Duplicate query ID") + return self + + def fingerprint(self) -> str: + """Bind a run to the exact validated corpus/query/judgment specification.""" + encoded = json.dumps( + self.model_dump(mode="json"), sort_keys=True, ensure_ascii=False + ) + return hashlib.sha256(encoded.encode()).hexdigest() + + +class Result(Value): + """One delivered ranking per query, before any evaluator-side filtering.""" + + query_id: Identifier + status: Literal["success", "partial", "error", "timeout"] + record_ids: tuple[Identifier, ...] = () + + @model_validator(mode="after") + def validate_ranking(self): + """Require record-level deduplication and unambiguous failed outcomes.""" + if len(self.record_ids) != len(set(self.record_ids)): + raise ValueError( + "Duplicate ranked record: evaluate the record-level response" + ) + if self.status in {"error", "timeout"} and self.record_ids: + raise ValueError("Failed requests cannot contain a delivered ranking") + return self + + +class Timing(Value): + """One observed phase, with explicit request identity and cache condition.""" + + request_id: Identifier + query_id: Identifier + phase: Literal["end_to_end", "retrieval", "reranking", "read"] + condition: Literal["cold", "warm", "unknown"] + status: Literal["success", "partial", "error", "timeout"] + duration_ms: Milliseconds + + +class Run(Value): + """A named system configuration and its observed outputs; no model-score assumptions.""" + + system: str = Field(min_length=1) + dataset_sha256: str = Field(pattern=r"^[a-f0-9]{64}$") + corpus_id: Identifier + results: tuple[Result, ...] + timings: tuple[Timing, ...] = () + + @model_validator(mode="after") + def unique_observations(self): + """Reject accidentally duplicated queries or phase measurements.""" + if len({r.query_id for r in self.results}) != len(self.results): + raise ValueError("Duplicate query result") + keys = [(t.request_id, t.phase) for t in self.timings] + if len(keys) != len(set(keys)): + raise ValueError("Duplicate request phase timing") + requests: dict[str, tuple[str, str]] = {} + for timing in self.timings: + identity = (timing.query_id, timing.condition) + if requests.setdefault(timing.request_id, identity) != identity: + raise ValueError("Request timing identity changed between phases") + return self + + +class QueryReport(Value): + """Scores retain failures and label provenance, without copying private query text.""" + + query_id: str + tags: tuple[str, ...] + status: str + expectation: str + assessors: tuple[str, ...] + scores: dict[str, float] + returned: int + judged_in_top_10: int + unjudged_in_top_10: int + correctly_empty: bool | None + + +class LatencyReport(Value): + """Percentiles by phase/condition/outcome; failures never disappear into fast successes.""" + + phase: str + condition: str + status: str + samples: int + p50_ms: float + p95_ms: float + max_ms: float + + +class Report(Value): + """Descriptive evidence, not a production-readiness verdict.""" + + system: str + corpus_id: str + dataset_sha256: str + answerable_queries: int + aggregate: dict[str, float] + queries: tuple[QueryReport, ...] + latency: tuple[LatencyReport, ...] + warnings: tuple[str, ...] + + +def summarize(dataset: Dataset, run: Run) -> Report: + """Use established IR measures, preserving API rank order with unique synthetic scores.""" + if ( + run.corpus_id != dataset.corpus_id + or run.dataset_sha256 != dataset.fingerprint() + ): + raise ValueError("Run does not match the exact dataset and corpus") + query_ids = {q.id for q in dataset.queries} + if {r.query_id for r in run.results} != query_ids: + raise ValueError( + "Every dataset query needs exactly one result, including failures" + ) + if any(t.query_id not in query_ids for t in run.timings): + raise ValueError("Timing references an unknown query") + measures = [ + ir_measures.nDCG @ 10, + ir_measures.RR @ 10, + ir_measures.P @ 5, + ir_measures.R @ 20, + ] + names = dict( + zip(measures, ("nDCG@10", "MRR@10", "Precision@5", "Recall@20"), strict=True) + ) + reports = [] + results = {r.query_id: r for r in run.results} + for query in dataset.queries: + result = results[query.id] + scores = {} + if query.expectation == "relevant_records": + qrels = [ + ir_measures.Qrel(query.id, j.record_id, j.relevance) + for j in query.judgments + ] + ranked = [ + ir_measures.ScoredDoc(query.id, record, float(-rank)) + for rank, record in enumerate(result.record_ids) + ] + scores = { + names[m]: float(v) + for m, v in ir_measures.calc_aggregate(measures, qrels, ranked).items() + } + judged = {j.record_id for j in query.judgments} + top = result.record_ids[:10] + reports.append( + QueryReport( + query_id=query.id, + tags=query.tags, + status=result.status, + expectation=query.expectation, + assessors=tuple( + sorted( + {query.expectation_assessor} + | {j.assessor for j in query.judgments} + ) + ), + scores=scores, + returned=len(result.record_ids), + judged_in_top_10=sum(record in judged for record in top), + unjudged_in_top_10=sum(record not in judged for record in top), + correctly_empty=(result.status == "success" and not result.record_ids) + if query.expectation == "no_answer" + else None, + ) + ) + answerable = [q for q in reports if q.expectation == "relevant_records"] + groups: dict[tuple[str, str, str], list[float]] = defaultdict(list) + for timing in run.timings: + groups[(timing.phase, timing.condition, timing.status)].append( + timing.duration_ms + ) + latencies = [] + for (phase, condition, status), samples in sorted(groups.items()): + samples.sort() + # Nearest-rank percentiles are defined even for one observation; counts remain visible. + latencies.append( + LatencyReport( + phase=phase, + condition=condition, + status=status, + samples=len(samples), + p50_ms=samples[math.ceil(0.5 * len(samples)) - 1], + p95_ms=samples[math.ceil(0.95 * len(samples)) - 1], + max_ms=samples[-1], + ) + ) + warnings = [ + "Recall is against judged relevant records, not all connected-account content.", + "Unjudged records receive no relevance credit; pool and judge them before conclusions.", + "Latency uses nearest-rank percentiles; small samples do not establish tail reliability.", + ] + if any(assessor != "human" for q in reports for assessor in q.assessors): + warnings.append( + "Dataset includes agent or synthetic judgments, not exclusively human labels." + ) + if not run.timings: + warnings.append("No latency measurements supplied.") + return Report( + system=run.system, + corpus_id=run.corpus_id, + dataset_sha256=run.dataset_sha256, + answerable_queries=len(answerable), + aggregate={ + name: mean(q.scores[name] for q in answerable) for name in names.values() + } + if answerable + else {}, + queries=tuple(reports), + latency=tuple(latencies), + warnings=tuple(warnings), + ) + + +def main() -> None: + """Print a dataset fingerprint or save a report without overwriting existing artifacts.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("dataset", type=Path) + parser.add_argument("--run", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + dataset = Dataset.model_validate_json(args.dataset.read_bytes()) + if args.run is None: + print(dataset.fingerprint()) + return + if args.output is None: + parser.error("--output is required with --run") + report = summarize(dataset, Run.model_validate_json(args.run.read_bytes())) + # Private reports can contain record IDs; do not print them or use world-readable defaults. + fd = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w") as output: + output.write(report.model_dump_json(indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/evaluation/tests/test_retrieval.py b/evaluation/tests/test_retrieval.py new file mode 100644 index 0000000000..671e177e83 --- /dev/null +++ b/evaluation/tests/test_retrieval.py @@ -0,0 +1,156 @@ +"""Small adversarial datasets exercise evaluation honesty, not production relevance.""" + +import pytest +from pydantic import ValidationError + +from evaluation.retrieval import ( + Dataset, + Judgment, + Query, + Result, + Run, + Timing, + summarize, +) + + +def corpus(): + return Dataset( + corpus_id="fixture-v1", + version="labels-v1", + queries=( + Query( + id="known", + text="परियोजना बैठक", + expectation="relevant_records", + expectation_assessor="synthetic", + judgments=( + Judgment(record_id="meeting", relevance=3, assessor="synthetic"), + Judgment(record_id="spam", relevance=0, assessor="synthetic"), + ), + ), + Query( + id="failed", + text="a known document", + expectation="relevant_records", + expectation_assessor="synthetic", + judgments=( + Judgment(record_id="document", relevance=1, assessor="synthetic"), + ), + ), + Query( + id="absent", + text="nonexistent meeting", + expectation="no_answer", + expectation_assessor="synthetic", + judgments=(), + ), + ), + ) + + +def run(dataset, **changes): + values = dict( + system="synthetic", + corpus_id=dataset.corpus_id, + dataset_sha256=dataset.fingerprint(), + results=( + Result( + query_id="known", + status="success", + record_ids=("spam", "meeting", "unknown"), + ), + Result(query_id="failed", status="timeout"), + Result(query_id="absent", status="error"), + ), + ) + return Run(**(values | changes)) + + +def test_failed_query_counts_and_unjudged_is_not_labeled_irrelevant(): + dataset = corpus() + report = summarize(dataset, run(dataset)) + assert report.answerable_queries == 2 + assert report.aggregate["MRR@10"] == 0.25 + assert report.aggregate["Precision@5"] == 0.1 + assert report.aggregate["Recall@20"] == 0.5 + assert 0 < report.aggregate["nDCG@10"] < 0.5 + assert report.queries[0].judged_in_top_10 == 2 + assert report.queries[0].unjudged_in_top_10 == 1 + assert report.queries[2].correctly_empty is False + assert any("synthetic" in warning for warning in report.warnings) + + +def test_latency_conditions_failures_and_tail_are_separate(): + dataset = corpus() + timings = tuple( + Timing( + request_id=f"warm-{i}", + query_id="known", + phase="end_to_end", + condition="warm", + status="success", + duration_ms=i, + ) + for i in range(1, 21) + ) + ( + Timing( + request_id="cold", + query_id="known", + phase="end_to_end", + condition="cold", + status="success", + duration_ms=500, + ), + Timing( + request_id="timeout", + query_id="failed", + phase="end_to_end", + condition="warm", + status="timeout", + duration_ms=10000, + ), + ) + report = summarize(dataset, run(dataset, timings=timings)) + warm = next( + row + for row in report.latency + if row.condition == "warm" and row.status == "success" + ) + assert (warm.samples, warm.p50_ms, warm.p95_ms, warm.max_ms) == (20, 10, 19, 20) + assert len(report.latency) == 3 + + +def test_missing_query_or_changed_labels_cannot_improve_score_silently(): + dataset = corpus() + original = run(dataset) + with pytest.raises(ValueError, match="Every dataset query"): + summarize( + dataset, original.model_copy(update={"results": original.results[:1]}) + ) + with pytest.raises(ValueError, match="exact dataset"): + summarize(dataset.model_copy(update={"version": "new-labels"}), original) + with pytest.raises(ValidationError, match="Duplicate ranked record"): + Result(query_id="known", status="success", record_ids=("meeting", "meeting")) + with pytest.raises(ValidationError): + Timing( + request_id="x", + query_id="known", + phase="end_to_end", + condition="warm", + status="success", + duration_ms=float("nan"), + ) + + +def test_no_answer_success_is_separate_from_answerable_metrics(): + dataset = Dataset(corpus_id="empty", version="v1", queries=(corpus().queries[2],)) + observation = Run( + system="empty", + corpus_id=dataset.corpus_id, + dataset_sha256=dataset.fingerprint(), + results=(Result(query_id="absent", status="success"),), + ) + report = summarize(dataset, observation) + assert report.aggregate == {} + assert report.queries[0].correctly_empty is True From 1e4421e4a950ee745fcaa5cb89c01d928b1b7098 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:03:46 -0700 Subject: [PATCH 184/266] Admit native snapshots with authoritative revision ordering --- .../domains/native_ingestion/README.md | 44 +++ .../domains/native_ingestion/__init__.py | 1 + .../domains/native_ingestion/errors.py | 9 + .../domains/native_ingestion/models.py | 94 ++++++ .../domains/native_ingestion/service.py | 21 ++ .../domains/native_ingestion/store.py | 156 +++++++++ .../native_ingestion/tests/conftest.py | 5 + .../native_ingestion/tests/test_ingestion.py | 302 ++++++++++++++++++ 8 files changed, 632 insertions(+) create mode 100644 backend/airweave/domains/native_ingestion/README.md create mode 100644 backend/airweave/domains/native_ingestion/__init__.py create mode 100644 backend/airweave/domains/native_ingestion/errors.py create mode 100644 backend/airweave/domains/native_ingestion/models.py create mode 100644 backend/airweave/domains/native_ingestion/service.py create mode 100644 backend/airweave/domains/native_ingestion/store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/conftest.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_ingestion.py diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md new file mode 100644 index 0000000000..943abeb6a7 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/README.md @@ -0,0 +1,44 @@ +# Native snapshot admission + +This internal service retains versioned Almanac snapshots in the existing +canonical store. It does not provision sources, authenticate callers, enumerate +Almanac data, publish a search projection, or provide a public import endpoint. +Almanac remains the authority for content and access. + +The future trusted publisher must bind a source connection with +`short_name=almanac` and `config_fields={owner_id, dataset}`. Datasets are +`knowledge` and `sessions`. The service checks that binding against the fenced +organization/sync and every snapshot. It does not accept a fabricated Composio +account. A future API must issue the fence after authenticating the publisher; +the internal command itself is not an authorization credential. + +`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. + +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; explicit renewal needs a separately +designed access-authority operation. 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. 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/errors.py b/backend/airweave/domains/native_ingestion/errors.py new file mode 100644 index 0000000000..d8e3bd1c0e --- /dev/null +++ b/backend/airweave/domains/native_ingestion/errors.py @@ -0,0 +1,9 @@ +"""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" diff --git a/backend/airweave/domains/native_ingestion/models.py b/backend/airweave/domains/native_ingestion/models.py new file mode 100644 index 0000000000..adb92b9bf7 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/models.py @@ -0,0 +1,94 @@ +"""Versioned native snapshots, distinct from provider observations and originals.""" + +from typing import Annotated, Literal + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, model_validator + +from airweave.domains.entities.canonical.requests import RecordIdentity, WriterFence + + +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 diff --git a/backend/airweave/domains/native_ingestion/service.py b/backend/airweave/domains/native_ingestion/service.py new file mode 100644 index 0000000000..3aa3046d55 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/service.py @@ -0,0 +1,21 @@ +"""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.native_ingestion.models import IngestNativeBatch +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) diff --git a/backend/airweave/domains/native_ingestion/store.py b/backend/airweave/domains/native_ingestion/store.py new file mode 100644 index 0000000000..354c0607ec --- /dev/null +++ b/backend/airweave/domains/native_ingestion/store.py @@ -0,0 +1,156 @@ +"""Native admission composed with the existing canonical writer transaction.""" + +from pydantic import ValidationError +from sqlalchemy import select +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.store import CanonicalRecordStore +from airweave.domains.native_ingestion.errors import NativeAdmissionError +from airweave.domains.native_ingestion.models import ( + IngestNativeBatch, + NativeSnapshot, + NativeSourceBinding, + NativeVersion, +) +from airweave.models.entity import Entity +from airweave.models.source_connection import SourceConnection + + +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) -> 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") + # A retry acknowledges retention, never restores local visibility. + return False + if entity.removal_reason in ("access_revoked", "scope_removed", "absent"): + raise NativeAdmissionError("Native source visibility requires explicit renewal") + return True + + +async def _validate_message_parent( + db: AsyncSession, request: IngestNativeBatch, snapshot: NativeSnapshot +) -> 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. + attested = next((item for item in request.snapshots if item.identity == parent), None) + if attested is None: + payload = await db.scalar( + select(Entity.source_payload).where( + Entity.organization_id == request.fence.organization_id, + Entity.sync_id == request.fence.sync_id, + Entity.entity_definition_short_name == parent.record_type, + Entity.entity_id == parent.entity_key, + ) + ) + try: + attested = NativeSnapshot.model_validate(payload) + 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") + + +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) + 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 + accepted = [] + unchanged = 0 + for snapshot in request.snapshots: + 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 = await db.scalar( + select(Entity) + .where( + Entity.organization_id == request.fence.organization_id, + Entity.sync_id == sync.id, + Entity.entity_definition_short_name == snapshot.identity.record_type, + Entity.entity_id == snapshot.identity.entity_key, + ) + .execution_options(populate_existing=True) + ) + if entity is not None and not _admit(entity, snapshot): + unchanged += 1 + continue + await _validate_message_parent(db, request, snapshot) + accepted.append( + 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, + ) + ) + result = await self.canonical._capture_locked( + db, sync, CaptureBatch(fence=request.fence, records=tuple(accepted)) + ) + return result.model_copy(update={"unchanged": result.unchanged + unchanged}) 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_ingestion.py b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py new file mode 100644 index 0000000000..34aed39522 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py @@ -0,0 +1,302 @@ +"""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.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"): + async with database() as db: + db.add( + SourceConnection( + id=uuid4(), + organization_id=fence.organization_id, + sync_id=fence.sync_id, + name="Native test", + short_name=provider, + 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"] From 7bf81fddfc48c8e8c48d175a1bd948f9ceb8f799 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:06:54 -0700 Subject: [PATCH 185/266] Batch native snapshot admission reads under the writer lock --- .../domains/native_ingestion/README.md | 4 ++ .../domains/native_ingestion/store.py | 57 +++++++++++-------- 2 files changed, 37 insertions(+), 24 deletions(-) diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 943abeb6a7..f505215ba4 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -18,6 +18,10 @@ the internal command itself is not an authorization credential. 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 diff --git a/backend/airweave/domains/native_ingestion/store.py b/backend/airweave/domains/native_ingestion/store.py index 354c0607ec..f2c507b957 100644 --- a/backend/airweave/domains/native_ingestion/store.py +++ b/backend/airweave/domains/native_ingestion/store.py @@ -1,7 +1,7 @@ """Native admission composed with the existing canonical writer transaction.""" from pydantic import ValidationError -from sqlalchemy import select +from sqlalchemy import select, tuple_ from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.models import CaptureResult @@ -55,26 +55,22 @@ def _admit(entity: Entity, snapshot: NativeSnapshot) -> bool: return True -async def _validate_message_parent( - db: AsyncSession, request: IngestNativeBatch, snapshot: NativeSnapshot +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. - attested = next((item for item in request.snapshots if item.identity == parent), None) + key = (parent.record_type, parent.entity_key) + attested = submitted.get(key) if attested is None: - payload = await db.scalar( - select(Entity.source_payload).where( - Entity.organization_id == request.fence.organization_id, - Entity.sync_id == request.fence.sync_id, - Entity.entity_definition_short_name == parent.record_type, - Entity.entity_id == parent.entity_key, - ) - ) + entity = retained.get(key) try: - attested = NativeSnapshot.model_validate(payload) + 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 ( @@ -118,26 +114,39 @@ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureR 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} accepted = [] unchanged = 0 for snapshot in request.snapshots: 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 = await db.scalar( - select(Entity) - .where( - Entity.organization_id == request.fence.organization_id, - Entity.sync_id == sync.id, - Entity.entity_definition_short_name == snapshot.identity.record_type, - Entity.entity_id == snapshot.identity.entity_key, - ) - .execution_options(populate_existing=True) - ) + entity = retained.get((snapshot.identity.record_type, snapshot.identity.entity_key)) if entity is not None and not _admit(entity, snapshot): unchanged += 1 continue - await _validate_message_parent(db, request, snapshot) + _validate_message_parent(snapshot, submitted, retained) accepted.append( CaptureRecord( identity=snapshot.identity, From 62409464f068be58188ddb7cdaa56b7bedc5c09c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:17:01 -0700 Subject: [PATCH 186/266] Project native knowledge and original conversation text into shared search --- .../entities/canonical/projection_inputs.py | 12 +- .../entities/canonical/projection_mappers.py | 9 + .../domains/entities/canonical/projector.py | 4 +- .../domains/entities/canonical/source.py | 13 ++ .../domains/native_ingestion/README.md | 27 ++- .../domains/native_ingestion/projection.py | 218 ++++++++++++++++++ .../native_ingestion/tests/test_projection.py | 213 +++++++++++++++++ .../native_ingestion/tests/test_search.py | 101 ++++++++ .../domains/search/canonical_visibility.py | 5 +- backend/airweave/domains/search/owned.py | 7 +- 10 files changed, 598 insertions(+), 11 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/projection.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_projection.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_search.py diff --git a/backend/airweave/domains/entities/canonical/projection_inputs.py b/backend/airweave/domains/entities/canonical/projection_inputs.py index 8747b22aa6..791ce7fd14 100644 --- a/backend/airweave/domains/entities/canonical/projection_inputs.py +++ b/backend/airweave/domains/entities/canonical/projection_inputs.py @@ -1,6 +1,8 @@ """Typed mapper outputs; source descriptors survive filtering and entity-ID stamping.""" -from pydantic import BaseModel, ConfigDict +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 @@ -14,6 +16,14 @@ class ProjectionInput(BaseModel): 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): diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 885f6aeccd..80d1167d96 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -250,6 +250,10 @@ def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: def excluded_from_search(record: SourceRecord, source_name: str) -> bool: """Retained calendar exclusions intentionally publish no searchable meeting.""" + if source_name == "almanac": + from airweave.domains.native_ingestion.projection import excluded_native + + return excluded_native(record) return source_name == "google_calendar" and ( record.identity.record_type == "event_occurrence" or (record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload)) @@ -268,6 +272,11 @@ async def map_record( if excluded_from_search(record, source_name): yield ProjectionInputs(parts=()) return + if source_name == "almanac": + from airweave.domains.native_ingestion.projection import map_native + + yield map_native(record) + return if ( source_name == "google_drive" and record.identity.record_type == "file" diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 6e8d0bc4fc..2a4bbc7d30 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -226,7 +226,9 @@ def _select_inputs( outcomes = [] for item in mapped.parts: part, entity = item.part, item.entity - if entity is None: + if item.omission == "unsupported_format": + outcome, reason = "unsupported", "unsupported_format" + elif entity is None: outcome, reason = "unavailable_original", "original_not_captured" elif ( part.kind == "file" diff --git a/backend/airweave/domains/entities/canonical/source.py b/backend/airweave/domains/entities/canonical/source.py index e3c646840d..c3557387b0 100644 --- a/backend/airweave/domains/entities/canonical/source.py +++ b/backend/airweave/domains/entities/canonical/source.py @@ -14,12 +14,25 @@ if TYPE_CHECKING: from airweave.domains.browse_tree.types import NodeSelectionData + from airweave.domains.sources.protocols import SourceRegistryProtocol from airweave.domains.storage.file_service import FileService from airweave.domains.syncs.cursors.cursor import SyncCursor SourceObservation = CaptureRecord | CompletedScope | RemovedScope | StartedScope +def indexed_record_types(short_name: str, registry: SourceRegistryProtocol) -> tuple[str, ...]: + """Index capability is independent of provider authentication or crawling. + + Native snapshots use canonical publications but have no provider lifecycle. + This declaration grants no access: callers must still authorize the stored + source connection and validate its current publication. + """ + if short_name == "almanac": + return ("knowledge", "session", "message") + return getattr(registry.get(short_name).source_class_ref, "canonical_record_types", ()) + + @runtime_checkable class CanonicalSource(Protocol): """Opt-in contract; declared types are audited operational records, not search chunks.""" diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index f505215ba4..4f8ec53a79 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -2,7 +2,7 @@ This internal service retains versioned Almanac snapshots in the existing canonical store. It does not provision sources, authenticate callers, enumerate -Almanac data, publish a search projection, or provide a public import endpoint. +Almanac data, or provide a public import endpoint. Almanac remains the authority for content and access. The future trusted publisher must bind a source connection with @@ -46,3 +46,28 @@ The publisher must not invent a version from capture time or an observation coun Tests use the existing disposable PostgreSQL fixture, including actual migration schemas and transaction rollback. They do not contact a provider or production. + +## 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. Other structured knowledge fields are retained but + do not yet have a curated searchable representation. +- 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 trusted publisher and +owner-binding API must still be implemented before exposing native ingestion. diff --git a/backend/airweave/domains/native_ingestion/projection.py b/backend/airweave/domains/native_ingestion/projection.py new file mode 100644 index 0000000000..a263162e8d --- /dev/null +++ b/backend/airweave/domains/native_ingestion/projection.py @@ -0,0 +1,218 @@ +"""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.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) + 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, + 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/tests/test_projection.py b/backend/airweave/domains/native_ingestion/tests/test_projection.py new file mode 100644 index 0000000000..01417617b3 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/tests/test_projection.py @@ -0,0 +1,213 @@ +"""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 excluded_from_search, map_record +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(), 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 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..12c45e41a7 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/tests/test_search.py @@ -0,0 +1,101 @@ +"""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 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, snapshot +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 +from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata + + +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, snapshot()) + async with database() as db: + deployment = VectorDbDeploymentMetadata( + dense_embedder="fake", embedding_dimensions=3, sparse_embedder="fake" + ) + db.add(deployment) + await db.flush() + collection = Collection( + name="Native synthetic", + readable_id="native-search", + organization_id=fence.organization_id, + vector_db_deployment_metadata_id=deployment.id, + ) + db.add(collection) + await db.flush() + await db.execute( + update(SourceConnection) + .where(SourceConnection.sync_id == fence.sync_id) + .values(readable_collection_id=collection.readable_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, 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/search/canonical_visibility.py b/backend/airweave/domains/search/canonical_visibility.py index 9670d79526..2a6aebb03a 100644 --- a/backend/airweave/domains/search/canonical_visibility.py +++ b/backend/airweave/domains/search/canonical_visibility.py @@ -8,6 +8,7 @@ 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 @@ -45,9 +46,7 @@ async def visible_results( scope = scopes.get(metadata.sync_id) if scope is None or metadata.source_name != scope.short_name: continue - canonical = bool( - getattr(registry.get(scope.short_name).source_class_ref, "canonical_record_types", ()) - ) + canonical = bool(indexed_record_types(scope.short_name, registry)) try: locator = ProjectionLocator.parse(metadata.original_entity_id) except ValueError: diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 479d37dcd4..a739fc561f 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -19,6 +19,7 @@ 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 ( OwnedSearchCoverage, @@ -256,11 +257,7 @@ async def _resolve_scopes( allowed_types = set() groups = defaultdict(list) for connection in snapshots: - types = getattr( - self._registry.get(connection.short_name).source_class_ref, - "canonical_record_types", - (), - ) + 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) From 5231b9e4e225aadc607736bdab64e526710ad84d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:20:12 -0700 Subject: [PATCH 187/266] Treat canonical search locators without web URLs as valid --- .../domains/search/adapters/vector_db/vespa_client.py | 4 +++- backend/airweave/domains/search/types/results.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) 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 0c7322beec..7bc6840f50 100644 --- a/backend/airweave/domains/search/adapters/vector_db/vespa_client.py +++ b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py @@ -461,7 +461,9 @@ def _convert_hits_to_results(self, hits: List[Dict[str, Any]]) -> SearchResults: ), 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/types/results.py b/backend/airweave/domains/search/types/results.py index 9c6442c608..8db44554cb 100644 --- a/backend/airweave/domains/search/types/results.py +++ b/backend/airweave/domains/search/types/results.py @@ -106,7 +106,7 @@ class SearchResult(BaseModel): 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( From 915489d7fe467fc1cb1d77faa2da75a4cb808d34 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:28:22 -0700 Subject: [PATCH 188/266] ci: exercise native ingestion and retrieval evaluation --- .github/workflows/owned-source-store.yml | 16 ++++++++++++++++ Makefile | 4 ++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 16b23be45d..db5c014ae6 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -5,6 +5,7 @@ on: branches: [dev, main] paths: - backend/** + - evaluation/** - vespa/app/** - Makefile - .github/workflows/owned-source-store.yml @@ -12,6 +13,7 @@ on: branches: [dev] paths: - backend/** + - evaluation/** - vespa/app/** - Makefile - .github/workflows/owned-source-store.yml @@ -24,6 +26,20 @@ concurrency: cancel-in-progress: true jobs: + evaluation: + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 + with: + python-version: '3.13' + - name: Verify offline retrieval metrics and file boundaries + run: | + python -m venv .venv-evaluation + .venv-evaluation/bin/python -m pip install ir-measures==0.4.3 pydantic==2.11.9 pytest==8.3.5 + .venv-evaluation/bin/python -m pytest evaluation/tests -q --confcutdir=evaluation + correctness: runs-on: ubuntu-latest timeout-minutes: 25 diff --git a/Makefile b/Makefile index 0613b8a598..eb343eea01 100644 --- a/Makefile +++ b/Makefile @@ -32,7 +32,7 @@ setup: test-store: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/syncs/tests + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/native_ingestion/tests airweave/domains/syncs/tests test-capture: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py @@ -42,7 +42,7 @@ test-provisioning: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/owned_provisioning/tests airweave/domains/source_connections/tests airweave/domains/temporal/tests airweave/domains/sources/tests/test_lifecycle.py airweave/domains/sync_pipeline/tests/test_factory.py test-search: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search airweave/domains/embedders test-auth: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/api/tests/test_service_auth.py airweave/api/tests/test_sync_authorization.py airweave/api/tests/test_context_resolver_auth.py From ff58a84a182d07ca1d1234ac2569c860a8b88402 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:31:38 -0700 Subject: [PATCH 189/266] fix: reserve native source mutations for import lifecycle --- .../airweave/domains/owned_provisioning/guard.py | 11 ++++++++++- .../airweave/domains/source_connections/create.py | 2 ++ .../airweave/domains/source_connections/delete.py | 4 +++- .../domains/source_connections/service.py | 2 ++ .../source_connections/tests/test_create.py | 10 ++++++++++ .../source_connections/tests/test_delete.py | 14 ++++++++++++++ .../source_connections/tests/test_service.py | 14 ++++++++++++++ .../source_connections/tests/test_update.py | 15 +++++++++++++++ .../airweave/domains/source_connections/update.py | 4 +++- 9 files changed, 73 insertions(+), 3 deletions(-) diff --git a/backend/airweave/domains/owned_provisioning/guard.py b/backend/airweave/domains/owned_provisioning/guard.py index 57a63f6f00..b9e8469baa 100644 --- a/backend/airweave/domains/owned_provisioning/guard.py +++ b/backend/airweave/domains/owned_provisioning/guard.py @@ -9,8 +9,17 @@ from airweave.models.owned_provisioning import OwnedProvisioning -async def require_unmanaged_source(db: AsyncSession, source_id: UUID, organization: UUID) -> None: +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, diff --git a/backend/airweave/domains/source_connections/create.py b/backend/airweave/domains/source_connections/create.py index 76a7217ba5..e6ca597b47 100644 --- a/backend/airweave/domains/source_connections/create.py +++ b/backend/airweave/domains/source_connections/create.py @@ -23,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, @@ -198,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, diff --git a/backend/airweave/domains/source_connections/delete.py b/backend/airweave/domains/source_connections/delete.py index aed95f15a5..5e36e61501 100644 --- a/backend/airweave/domains/source_connections/delete.py +++ b/backend/airweave/domains/source_connections/delete.py @@ -50,7 +50,9 @@ async def delete( if not source_conn: raise NotFoundException("Source connection not found") - await require_unmanaged_source(db, id, ctx.organization.id) + 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( 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 17d69140d0..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 @@ -212,3 +213,16 @@ async def test_delete_sync_service_failure_propagates(): with pytest.raises(RuntimeError, match="sync delete boom"): 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 0449431731..5af8384332 100644 --- a/backend/airweave/domains/source_connections/tests/test_update.py +++ b/backend/airweave/domains/source_connections/tests/test_update.py @@ -656,3 +656,18 @@ async def test_credential_update_unpause_failure_is_nonfatal(): obj_in = SourceConnectionUpdate(authentication={"credentials": {"token": "new_secret"}}) 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() diff --git a/backend/airweave/domains/source_connections/update.py b/backend/airweave/domains/source_connections/update.py index facced9c4c..f26ce795be 100644 --- a/backend/airweave/domains/source_connections/update.py +++ b/backend/airweave/domains/source_connections/update.py @@ -94,7 +94,9 @@ async def update( if not source_conn: raise NotFoundException("Source connection not found") - await require_unmanaged_source(db, id, ctx.organization.id) + 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) From 92da4c8f8a6dfb4b4adf483efa97ff8691c2f1d5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:34:33 -0700 Subject: [PATCH 190/266] Compose native snapshot admission with atomic scan pages --- .../domains/entities/canonical/scan_store.py | 27 +- .../domains/native_ingestion/README.md | 23 ++ .../domains/native_ingestion/models.py | 28 ++- .../domains/native_ingestion/service.py | 8 +- .../domains/native_ingestion/store.py | 147 +++++++++-- .../native_ingestion/tests/test_pages.py | 230 ++++++++++++++++++ 6 files changed, 433 insertions(+), 30 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_pages.py diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 2178ed2810..355d9f3b0c 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -2,6 +2,7 @@ import json from datetime import datetime, timezone +from typing import Protocol from uuid import UUID, uuid4 from pydantic import JsonValue @@ -55,6 +56,7 @@ ) from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity +from airweave.models.sync import Sync from airweave.models.sync_cursor import SyncCursor @@ -214,6 +216,20 @@ def attest_planned_owner(cycle: CaptureCycle, request: BeginScan, parent: Entity raise ScanConflict("Scope owner changed during plan selection") +class LockedPageCapture(Protocol): + """Internal composition: capture the entire page under the caller's writer lock. + + Implementations must not commit. This is code-owned behavior, never HTTP data. + The surrounding scan transaction owns scope checks and continuation advancement. + """ + + async def __call__( + self, db: AsyncSession, sync: Sync, batch: CaptureBatch, *, seen_id: UUID + ) -> CaptureResult: + """Capture changes and sightings atomically using the existing sweep identity.""" + ... + + class CanonicalScanStore: """Reuse record capture, lock ordering and journal semantics for recoverable scans.""" @@ -368,7 +384,13 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: await db.flush() return await self._state(db, row) - async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: + async def page( + self, + db: AsyncSession, + request: CommitScanPage, + *, + capture_page: LockedPageCapture | None = None, + ) -> ScanResult: """Capture and advance the page as one transaction, never a partial acknowledgement.""" sync = await self.records._fenced_sync(db, request.fence) _, cycle = await attest_cycle(db, request.fence, request.cycle_id) @@ -414,7 +436,8 @@ async def page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: for record in request.discovered_records ): raise ScanConflict("Discovered originals must be verified independent declared roots") - captured = await self.records._capture_locked( + capture = capture_page or self.records._capture_locked + captured = await capture( db, sync, CaptureBatch(fence=request.fence, records=request.records), diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 4f8ec53a79..35d7ca3c6f 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -47,6 +47,29 @@ The publisher must not invent a version from capture time or an observation coun 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, request receipts, import provisioning, explicit renewal and +publisher completeness attestation remain unfinished. Cancellation must resolve +the job's actual bound source; a caller-supplied source ID is insufficient evidence +for authorizing a job cancellation. + ## Native search projection `projection.py` supplies offline inputs to the existing canonical projector. diff --git a/backend/airweave/domains/native_ingestion/models.py b/backend/airweave/domains/native_ingestion/models.py index adb92b9bf7..3d7a0b4198 100644 --- a/backend/airweave/domains/native_ingestion/models.py +++ b/backend/airweave/domains/native_ingestion/models.py @@ -1,10 +1,18 @@ """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 RecordIdentity, WriterFence +from airweave.domains.entities.canonical.requests import ( + CaptureRecord, + CompletedScope, + RecordIdentity, + ScanVersion, + WriterFence, +) +from airweave.domains.entities.canonical.scan_models import ScanContinuation class NativeModel(BaseModel): @@ -92,3 +100,21 @@ def unique_records(self) -> "IngestNativeBatch": 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/service.py b/backend/airweave/domains/native_ingestion/service.py index 3aa3046d55..ae618839c8 100644 --- a/backend/airweave/domains/native_ingestion/service.py +++ b/backend/airweave/domains/native_ingestion/service.py @@ -4,7 +4,8 @@ from airweave.db.unit_of_work import UnitOfWork from airweave.domains.entities.canonical.models import CaptureResult -from airweave.domains.native_ingestion.models import IngestNativeBatch +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 @@ -19,3 +20,8 @@ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureR """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/store.py b/backend/airweave/domains/native_ingestion/store.py index f2c507b957..575d0a3d26 100644 --- a/backend/airweave/domains/native_ingestion/store.py +++ b/backend/airweave/domains/native_ingestion/store.py @@ -1,21 +1,28 @@ """Native admission composed with the existing canonical writer transaction.""" +from uuid import UUID + from pydantic import ValidationError -from sqlalchemy import select, tuple_ +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.store import CanonicalRecordStore +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: @@ -82,6 +89,50 @@ def _validate_message_parent( 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, +) -> 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): + 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. @@ -97,6 +148,23 @@ def __init__(self, canonical: CanonicalRecordStore): 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 + ) -> 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) @@ -136,30 +204,57 @@ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureR .execution_options(populate_existing=True) ) retained = {(row.entity_definition_short_name, row.entity_id): row for row in rows} - accepted = [] - unchanged = 0 - for snapshot in request.snapshots: - 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): - unchanged += 1 - continue - _validate_message_parent(snapshot, submitted, retained) - accepted.append( - 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, + admitted = _admit_snapshots( + request.snapshots, binding, submitted, retained, request, sightings=sightings + ) + unchanged = admitted.unchanged_ids + if sightings and unchanged: + available = set( + await db.scalars( + select(Entity.id).where(Entity.id.in_(unchanged), content_is_available()) ) ) - result = await self.canonical._capture_locked( - db, sync, CaptureBatch(fence=request.fence, records=tuple(accepted)) + 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, ) - return result.model_copy(update={"unchanged": result.unchanged + unchanged}) 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 From 7c3655959fcb429fc6957fdc69ed056250fc8748 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:35:52 -0700 Subject: [PATCH 191/266] test: qualify native test modules in combined store suite --- backend/airweave/domains/native_ingestion/tests/__init__.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 backend/airweave/domains/native_ingestion/tests/__init__.py 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.""" From 36c45fc50eeffb90ce8733c1aa67732142f3de9d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:43:40 -0700 Subject: [PATCH 192/266] Add backend-only immutable native source binding API --- backend/airweave/api/backend_actor.py | 13 ++ backend/airweave/api/v1/api.py | 2 + .../api/v1/endpoints/native_sources.py | 38 +++++ .../api/v1/endpoints/owned_provisioning.py | 10 +- backend/airweave/core/container/container.py | 2 + backend/airweave/core/container/factory.py | 3 + .../domains/native_ingestion/source_models.py | 70 +++++++++ .../native_ingestion/source_service.py | 29 ++++ .../domains/native_ingestion/source_store.py | 140 +++++++++++++++++ .../native_ingestion/tests/test_sources.py | 141 ++++++++++++++++++ 10 files changed, 440 insertions(+), 8 deletions(-) create mode 100644 backend/airweave/api/backend_actor.py create mode 100644 backend/airweave/api/v1/endpoints/native_sources.py create mode 100644 backend/airweave/domains/native_ingestion/source_models.py create mode 100644 backend/airweave/domains/native_ingestion/source_service.py create mode 100644 backend/airweave/domains/native_ingestion/source_store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_sources.py diff --git a/backend/airweave/api/backend_actor.py b/backend/airweave/api/backend_actor.py new file mode 100644 index 0000000000..aa17b5fe7b --- /dev/null +++ b/backend/airweave/api/backend_actor.py @@ -0,0 +1,13 @@ +"""Shared backend-only authority gate; organization keys retain their existing scope.""" + +from fastapi import Depends, HTTPException + +from airweave.api import deps +from airweave.api.context import ApiContext + + +def backend_actor(ctx: ApiContext = Depends(deps.get_context)) -> ApiContext: + """A user session cannot act as Almanac's credential/identity authority.""" + if not ctx.is_api_key_auth: + raise HTTPException(status_code=403, detail="Owned provisioning requires a backend API key") + return ctx diff --git a/backend/airweave/api/v1/api.py b/backend/airweave/api/v1/api.py index 0d6f7557b1..526285b70b 100644 --- a/backend/airweave/api/v1/api.py +++ b/backend/airweave/api/v1/api.py @@ -14,6 +14,7 @@ entity_counts, file_retrieval, health, + native_sources, organizations, owned_provisioning, records, @@ -63,6 +64,7 @@ ) api_router.include_router(webhooks.router, prefix="/webhooks", tags=["webhooks"]) api_router.include_router(connect.router, prefix="/connect", tags=["connect"]) +api_router.include_router(native_sources.router, prefix="/native/sources", tags=["native-sources"]) api_router.include_router(calendar_records.router, prefix="/sync", tags=["records"]) diff --git a/backend/airweave/api/v1/endpoints/native_sources.py b/backend/airweave/api/v1/endpoints/native_sources.py new file mode 100644 index 0000000000..0af6ffc50a --- /dev/null +++ b/backend/airweave/api/v1/endpoints/native_sources.py @@ -0,0 +1,38 @@ +"""Backend-key native source binding; no OAuth, imports or provider execution.""" + +from uuid import UUID + +from fastapi import Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api import deps +from airweave.api.backend_actor import backend_actor +from airweave.api.context import ApiContext +from airweave.api.router import TrailingSlashRouter +from airweave.core.container import Container +from airweave.db.session import get_db +from airweave.domains.native_ingestion.source_models import EnsureNativeSource, NativeSource + +router = TrailingSlashRouter() + + +@router.put("", response_model=NativeSource) +async def ensure_native_source( + request: EnsureNativeSource, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeSource: + """Ensure one immutable owner/dataset binding in the authenticated organization.""" + return await container.native_sources.ensure(db, ctx.organization.id, request) + + +@router.get("/{source_id}", response_model=NativeSource) +async def read_native_source( + source_id: UUID, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeSource: + """Read the existing binding without starting synchronization.""" + return await container.native_sources.get(db, ctx.organization.id, source_id) diff --git a/backend/airweave/api/v1/endpoints/owned_provisioning.py b/backend/airweave/api/v1/endpoints/owned_provisioning.py index 6b52682371..aee9117771 100644 --- a/backend/airweave/api/v1/endpoints/owned_provisioning.py +++ b/backend/airweave/api/v1/endpoints/owned_provisioning.py @@ -2,10 +2,11 @@ from uuid import UUID -from fastapi import Depends, HTTPException +from fastapi import Depends from sqlalchemy.ext.asyncio import AsyncSession from airweave.api import deps +from airweave.api.backend_actor import backend_actor from airweave.api.context import ApiContext from airweave.api.router import TrailingSlashRouter from airweave.core.container import Container @@ -15,13 +16,6 @@ router = TrailingSlashRouter() -def backend_actor(ctx: ApiContext = Depends(deps.get_context)) -> ApiContext: - """A user session cannot act as Almanac's credential/identity authority.""" - if not ctx.is_api_key_auth: - raise HTTPException(status_code=403, detail="Owned provisioning requires a backend API key") - return ctx - - @router.put("/{account_id}", response_model=ProvisionedSource) async def ensure_owned_source( account_id: UUID, diff --git a/backend/airweave/core/container/container.py b/backend/airweave/core/container/container.py index 00ab5fda18..121f7cc809 100644 --- a/backend/airweave/core/container/container.py +++ b/backend/airweave/core/container/container.py @@ -62,6 +62,7 @@ EntityDefinitionRegistryProtocol, EntityRepositoryProtocol, ) +from airweave.domains.native_ingestion.source_service import NativeSources from airweave.domains.oauth.protocols import ( OAuth1ServiceProtocol, OAuth2ServiceProtocol, @@ -255,6 +256,7 @@ async def my_endpoint(event_bus: EventBus = Inject(EventBus)): # Search domain (v2 tiers) owned_search: OwnedSearchService owned_provisioning: OwnedProvisioningService + native_sources: NativeSources instant_search: InstantSearchServiceProtocol classic_search: ClassicSearchServiceProtocol agentic_search: AgenticSearchServiceProtocol diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index fedda80239..6f3ba2e41a 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -88,6 +88,8 @@ from airweave.domains.entities.entity_count_repository import EntityCountRepository from airweave.domains.entities.entity_repository import EntityRepository from airweave.domains.entities.registry import EntityDefinitionRegistry +from airweave.domains.native_ingestion.source_service import NativeSources +from airweave.domains.native_ingestion.source_store import NativeSourceStore from airweave.domains.oauth.callback_service import OAuthCallbackService from airweave.domains.oauth.flow_service import OAuthFlowService from airweave.domains.oauth.oauth1_service import OAuth1Service @@ -617,6 +619,7 @@ def create_container(settings: Settings) -> Container: user_service=user_service, email_service=email_service, owned_search=search_deps["owned_search"], + native_sources=NativeSources(NativeSourceStore()), owned_provisioning=OwnedProvisioningService( store=ProvisioningStore(create_service, source_validation), lifecycle=source_deps["source_lifecycle_service"], 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..31541197d9 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/source_models.py @@ -0,0 +1,70 @@ +"""Immutable native source identity and backend-only ensure/read contracts.""" + +import json +from uuid import UUID, uuid5 + +from pydantic import BaseModel, ConfigDict, Field + +from airweave.domains.native_ingestion.models import NativeModel, NativeSourceBinding +from airweave.models.source_connection import SourceConnection +from airweave.models.sync import Sync + +# 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 + + +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, + ) 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..20e5e24b5d --- /dev/null +++ b/backend/airweave/domains/native_ingestion/source_store.py @@ -0,0 +1,140 @@ +"""Primary-key serialized native binding; no JSON discovery or external side effects.""" + +from uuid import UUID + +from fastapi import HTTPException +from pydantic import 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, + LockedNativeSource, + 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 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/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() From 136faf08a52fbaf9e578448f15caf82955f14966 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:45:06 -0700 Subject: [PATCH 193/266] feat: persist native import identity and atomic writer start --- .../domains/native_ingestion/README.md | 23 +++ .../domains/native_ingestion/import_models.py | 39 ++++ .../domains/native_ingestion/import_store.py | 192 ++++++++++++++++++ .../native_ingestion/tests/test_imports.py | 141 +++++++++++++ 4 files changed, 395 insertions(+) create mode 100644 backend/airweave/domains/native_ingestion/import_models.py create mode 100644 backend/airweave/domains/native_ingestion/import_store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_imports.py diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 35d7ca3c6f..36e4a4cc8b 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -94,3 +94,26 @@ Native indexed sources use the same publication and access validation as provide sources, without a provider registry entry or fabricated OAuth connection. This is an index capability, not an authentication grant. The trusted publisher and owner-binding API must still be implemented before exposing native 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` is the internal start/read transaction primitive, not yet an +HTTP publisher 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: final scopes, stable export boundaries, receipts, +renewal and import completion/cancellation still need their HTTP lifecycle. 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..20dce9255f --- /dev/null +++ b/backend/airweave/domains/native_ingestion/import_models.py @@ -0,0 +1,39 @@ +"""Native import intent and durable identity, without client-controlled writer fences.""" + +from typing import Literal +from uuid import UUID + +from pydantic import Field + +from airweave.core.shared_models import SyncJobStatus +from airweave.domains.entities.canonical.requests import WriterFence +from airweave.domains.native_ingestion.models import NativeModel + + +class StartNativeImport(NativeModel): + """Publisher attestation; bounded imports never authorize absence deletion.""" + + snapshot_id: str = Field(min_length=1, max_length=256) + coverage: Literal["bounded", "complete"] + + +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 + + +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 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..c7a1eafc36 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -0,0 +1,192 @@ +"""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 select +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.core.shared_models import SyncStatus +from airweave.domains.entities.canonical.cycle_models import ( + BeginCycle, + CycleConfiguration, + RestartCycle, +) +from airweave.domains.entities.canonical.cycle_store import ( + begin_cycle, + cursor_row, + cycle_state, + restart_cycle, +) +from airweave.domains.entities.canonical.store import CanonicalRecordStore +from airweave.domains.native_ingestion.errors import NativeAdmissionError +from airweave.domains.native_ingestion.import_models import ( + NativeImportReceipt, + NativeImportState, + StartNativeImport, +) +from airweave.domains.native_ingestion.source_store import NativeSourceStore +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.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, + ) + + +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) + 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",)} + ) + policy = "exhaustive" if request.coverage == "complete" else "discovery_only" + configuration = CycleConfiguration( + fingerprint=fingerprint, + parents=parents, + completion_policies={kind: policy for kind in parents}, + ) + 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.""" + 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, + ) + ) + if job is None: + raise NativeAdmissionError("Native import does not exist") + return import_state(job, receipt(job, source_id, request_key)) 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..d3b13ea6a4 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/tests/test_imports.py @@ -0,0 +1,141 @@ +"""Real PostgreSQL import request recovery, concurrent creation and rollback.""" + +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.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_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 From 4dc2b65b10455335ff1efff2da3745ec732fe7d9 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:51:00 -0700 Subject: [PATCH 194/266] Expose authenticated native import start and recovery --- backend/airweave/api/middleware.py | 43 +++++- backend/airweave/api/v1/api.py | 2 + .../api/v1/endpoints/native_imports.py | 55 +++++++ backend/airweave/core/container/container.py | 2 + backend/airweave/core/container/factory.py | 6 + .../domains/native_ingestion/errors.py | 6 + .../native_ingestion/import_service.py | 36 +++++ .../native_ingestion/tests/test_import_api.py | 143 ++++++++++++++++++ 8 files changed, 286 insertions(+), 7 deletions(-) create mode 100644 backend/airweave/api/v1/endpoints/native_imports.py create mode 100644 backend/airweave/domains/native_ingestion/import_service.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_import_api.py diff --git a/backend/airweave/api/middleware.py b/backend/airweave/api/middleware.py index 6237aefa0a..791349c00d 100644 --- a/backend/airweave/api/middleware.py +++ b/backend/airweave/api/middleware.py @@ -33,6 +33,19 @@ from airweave.core.protocols import HttpMetrics +def _native_request(request: Request) -> bool: + """Native publisher paths carry identities and retained source data.""" + path = request.url.path + return "/native/sources/" in path or path.endswith("/native/sources") + + +def _observable_url(request: Request) -> str: + """Keep route shape for native requests without publisher keys or query values.""" + if not _native_request(request): + return str(request.url) + return _build_endpoint_name(request, fallback="/native/sources/{native_path}") + + async def add_request_id(request: Request, call_next: callable) -> Response: """Middleware to generate and add a request ID to the request for tracing. @@ -68,7 +81,8 @@ async def log_requests(request: Request, call_next: callable) -> Response: duration = time.time() - start_time logger.info( ( - f"Handled request {request.method} {request.url} in {duration:.2f} seconds." + f"Handled request {request.method} {_observable_url(request)} " + f"in {duration:.2f} seconds." f"Response code: {response.status_code}" ) ) @@ -92,6 +106,9 @@ async def exception_logging_middleware(request: Request, call_next: callable) -> response = await call_next(request) return response except Exception as exc: + if _native_request(request): + logger.error(f"Native request failed: {type(exc).__name__}") + return JSONResponse(status_code=500, content={"detail": "Internal Server Error"}) # Always log the full exception details logger.error(f"Unhandled exception: {exc}\n{traceback.format_exc()}") @@ -144,7 +161,7 @@ async def request_timeout_middleware(request: Request, call_next: callable) -> R except asyncio.TimeoutError: logger.warning( f"Request timeout after {settings.API_REQUEST_TIMEOUT_SECONDS}s: " - f"{request.method} {request.url}" + f"{request.method} {_observable_url(request)}" ) return JSONResponse( status_code=504, @@ -180,7 +197,7 @@ async def request_body_size_middleware(request: Request, call_next: callable) -> actual_size_mb = content_length_bytes / (1024 * 1024) logger.warning( f"Request body too large: {actual_size_mb:.2f}MB exceeds limit " - f"of {max_size_mb:.2f}MB for {request.method} {request.url}" + f"of {max_size_mb:.2f}MB for {request.method} {_observable_url(request)}" ) return JSONResponse( status_code=413, @@ -288,12 +305,18 @@ async def dispatch(self, request: Request, call_next): ) response.headers["Access-Control-Allow-Headers"] = "*" response.headers["Access-Control-Allow-Credentials"] = "true" - logger.debug(f"Handled OPTIONS preflight for {path} from origin {origin}") + logger.debug( + f"Handled OPTIONS preflight for " + f"{_observable_url(request) if _native_request(request) else path} " + f"from origin {origin}" + ) return response else: # Not allowed, return 403 logger.debug( - f"Rejected OPTIONS preflight for {path} from disallowed origin {origin}" + f"Rejected OPTIONS preflight for " + f"{_observable_url(request) if _native_request(request) else path} " + f"from disallowed origin {origin}" ) return Response(status_code=403) @@ -358,6 +381,9 @@ async def validation_exception_handler( """ # Extract basic error messages error_messages = unpack_validation_error(exc) + if _native_request(request): + logger.error(f"Native request validation failed: {len(exc.errors())} errors") + return JSONResponse(status_code=422, content=error_messages) logger.error(f"Validation error: {error_messages}") if settings.LOCAL_CURSOR_DEVELOPMENT: @@ -705,6 +731,9 @@ def _build_endpoint_name(request: Request, *, fallback: str | None = None) -> st # e.g., "/collections/{readable_id}/sources/{source_short_name}" return route.path + if _native_request(request): + return fallback if fallback is not None else "/native/sources/{native_path}" + # Fallback to raw path if route not available (shouldn't happen in normal operation) return fallback if fallback is not None else request.url.path.rstrip("/") @@ -732,14 +761,14 @@ async def _track_api_call_async( properties = { "endpoint": endpoint, "request_method": request.method, - "request_path": request.url.path, + "request_path": _observable_url(request) if _native_request(request) else request.url.path, "status_code": response.status_code, "duration_ms": duration_ms, } # Flatten all path parameters directly into properties # Convert all values to strings to ensure JSON serialization works - if request.path_params: + if request.path_params and not _native_request(request): properties.update({k: str(v) for k, v in request.path_params.items()}) # Determine event name diff --git a/backend/airweave/api/v1/api.py b/backend/airweave/api/v1/api.py index 526285b70b..0b03c530ba 100644 --- a/backend/airweave/api/v1/api.py +++ b/backend/airweave/api/v1/api.py @@ -14,6 +14,7 @@ entity_counts, file_retrieval, health, + native_imports, native_sources, organizations, owned_provisioning, @@ -64,6 +65,7 @@ ) api_router.include_router(webhooks.router, prefix="/webhooks", tags=["webhooks"]) api_router.include_router(connect.router, prefix="/connect", tags=["connect"]) +api_router.include_router(native_imports.router, prefix="/native/sources", tags=["native-imports"]) api_router.include_router(native_sources.router, prefix="/native/sources", tags=["native-sources"]) diff --git a/backend/airweave/api/v1/endpoints/native_imports.py b/backend/airweave/api/v1/endpoints/native_imports.py new file mode 100644 index 0000000000..424a01d8df --- /dev/null +++ b/backend/airweave/api/v1/endpoints/native_imports.py @@ -0,0 +1,55 @@ +"""Backend-only native import intent/recovery; no publisher-controlled fences.""" + +from typing import Annotated +from uuid import UUID + +from fastapi import Depends, HTTPException, Path +from sqlalchemy.ext.asyncio import AsyncSession + +from airweave.api import deps +from airweave.api.backend_actor import backend_actor +from airweave.api.context import ApiContext +from airweave.api.router import TrailingSlashRouter +from airweave.core.container import Container +from airweave.db.session import get_db +from airweave.domains.entities.canonical.store import WriterBusy +from airweave.domains.native_ingestion.errors import NativeAdmissionError, NativeImportNotFound +from airweave.domains.native_ingestion.import_models import NativeImportState, StartNativeImport + +router = TrailingSlashRouter() +RequestKey = Annotated[str, Path(min_length=1, max_length=128)] + + +@router.put("/{source_id}/imports/{request_key}", response_model=NativeImportState) +async def start_native_import( + source_id: UUID, + request_key: RequestKey, + request: StartNativeImport, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeImportState: + """Start or recover identical intent in the authenticated organization.""" + try: + return await container.native_imports.start( + db, ctx.organization.id, source_id, request_key, request + ) + except (NativeAdmissionError, WriterBusy) as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.get("/{source_id}/imports/{request_key}", response_model=NativeImportState) +async def read_native_import( + source_id: UUID, + request_key: RequestKey, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeImportState: + """Read import state without activating any writer.""" + try: + return await container.native_imports.read(db, ctx.organization.id, source_id, request_key) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except NativeAdmissionError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error diff --git a/backend/airweave/core/container/container.py b/backend/airweave/core/container/container.py index 121f7cc809..ad6f2af012 100644 --- a/backend/airweave/core/container/container.py +++ b/backend/airweave/core/container/container.py @@ -62,6 +62,7 @@ EntityDefinitionRegistryProtocol, EntityRepositoryProtocol, ) +from airweave.domains.native_ingestion.import_service import NativeImports from airweave.domains.native_ingestion.source_service import NativeSources from airweave.domains.oauth.protocols import ( OAuth1ServiceProtocol, @@ -257,6 +258,7 @@ async def my_endpoint(event_bus: EventBus = Inject(EventBus)): owned_search: OwnedSearchService owned_provisioning: OwnedProvisioningService native_sources: NativeSources + native_imports: NativeImports instant_search: InstantSearchServiceProtocol classic_search: ClassicSearchServiceProtocol agentic_search: AgenticSearchServiceProtocol diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 6f3ba2e41a..f7f6124a7a 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -85,9 +85,12 @@ from airweave.domains.embedders.sparse.fastembed import ( FastEmbedSparseEmbedder as DomainFastEmbedSparseEmbedder, ) +from airweave.domains.entities.canonical.store import CanonicalRecordStore from airweave.domains.entities.entity_count_repository import EntityCountRepository from airweave.domains.entities.entity_repository import EntityRepository from airweave.domains.entities.registry import EntityDefinitionRegistry +from airweave.domains.native_ingestion.import_service import NativeImports +from airweave.domains.native_ingestion.import_store import NativeImportStore from airweave.domains.native_ingestion.source_service import NativeSources from airweave.domains.native_ingestion.source_store import NativeSourceStore from airweave.domains.oauth.callback_service import OAuthCallbackService @@ -620,6 +623,9 @@ def create_container(settings: Settings) -> Container: email_service=email_service, owned_search=search_deps["owned_search"], native_sources=NativeSources(NativeSourceStore()), + native_imports=NativeImports( + NativeImportStore(NativeSourceStore(), CanonicalRecordStore()) + ), owned_provisioning=OwnedProvisioningService( store=ProvisioningStore(create_service, source_validation), lifecycle=source_deps["source_lifecycle_service"], diff --git a/backend/airweave/domains/native_ingestion/errors.py b/backend/airweave/domains/native_ingestion/errors.py index d8e3bd1c0e..a3b0f91e09 100644 --- a/backend/airweave/domains/native_ingestion/errors.py +++ b/backend/airweave/domains/native_ingestion/errors.py @@ -7,3 +7,9 @@ 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_service.py b/backend/airweave/domains/native_ingestion/import_service.py new file mode 100644 index 0000000000..7ce7201b35 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/import_service.py @@ -0,0 +1,36 @@ +"""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.import_models import NativeImportState, StartNativeImport +from airweave.domains.native_ingestion.import_store import NativeImportStore + + +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 + + 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) 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..2b7ab01307 --- /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 +async def native_api(database, monkeypatch): + 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="knowledge", 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() From 378a675a51dab88dbf6b025279cd4af91c67963a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:52:09 -0700 Subject: [PATCH 195/266] feat: retain atomic native page retry receipts --- .../domains/native_ingestion/README.md | 18 +++ .../domains/native_ingestion/import_store.py | 31 +++- .../domains/native_ingestion/page_models.py | 56 +++++++ .../domains/native_ingestion/page_store.py | 100 ++++++++++++ .../tests/test_page_receipts.py | 147 ++++++++++++++++++ 5 files changed, 350 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/page_models.py create mode 100644 backend/airweave/domains/native_ingestion/page_store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_page_receipts.py diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 36e4a4cc8b..f89ce799e1 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -117,3 +117,21 @@ Publisher-declared bounded snapshots use discovery-only completion policies; complete dataset snapshots use exhaustive policies. This declaration does not itself certify completeness: final scopes, stable export boundaries, receipts, renewal and import completion/cancellation still need their HTTP lifecycle. + +## 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` remains an internal transaction primitive until scope +endpoints and completion are wired. 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. diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py index c7a1eafc36..c986476f69 100644 --- a/backend/airweave/domains/native_ingestion/import_store.py +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -22,7 +22,7 @@ restart_cycle, ) from airweave.domains.entities.canonical.store import CanonicalRecordStore -from airweave.domains.native_ingestion.errors import NativeAdmissionError +from airweave.domains.native_ingestion.errors import NativeAdmissionError, NativeImportNotFound from airweave.domains.native_ingestion.import_models import ( NativeImportReceipt, NativeImportState, @@ -188,5 +188,32 @@ async def read( ) ) if job is None: - raise NativeAdmissionError("Native import does not exist") + raise NativeImportNotFound("Native import does not exist") return import_state(job, receipt(job, source_id, request_key)) + + 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 = await self.sources.require(db, organization_id, source_id) + if not bound.source.is_authenticated or bound.sync.status != SyncStatus.ACTIVE: + raise NativeAdmissionError("Native source is unavailable for import") + 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") + saved = receipt(job, source_id, request_key) + await self.canonical._fenced_sync(db, saved.fence) + return saved 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/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 + ) From d281362c244d1d775946f5f83bd726d33af3a273 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:57:02 -0700 Subject: [PATCH 196/266] Bound native import child work to observed membership --- .../entities/canonical/cycle_models.py | 3 + .../domains/entities/canonical/cycle_store.py | 18 +- .../domains/native_ingestion/import_store.py | 1 + .../tests/test_observed_membership.py | 177 ++++++++++++++++++ 4 files changed, 197 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_observed_membership.py diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 1a38e0b2df..38b7507ef9 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -47,10 +47,13 @@ class CycleConfiguration(BaseModel): completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) known_object_validation: tuple[RecordKind, ...] = () scope_changes: tuple[RecordKind, ...] = () + membership: Literal["retained", "observed"] = "retained" def digest(self) -> str: """Bind history to scope, topology and guarantees, including default policies.""" value = self.model_dump(mode="json") + if self.membership == "retained": + value.pop("membership") # Preserve existing provider checkpoint digests. if not self.known_object_validation: value.pop("known_object_validation") # Preserve existing default configuration digests. if not self.scope_changes: diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index 260444c504..f123514887 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -175,8 +175,14 @@ def inventory_complete(subject, fence: WriterFence, state: CaptureCycle): .correlate(subject) .scalar_subquery() ) + sightings = ( + (subject.last_seen_run_id == CaptureScan.sweep_id,) + if state.configuration.membership == "observed" + else () + ) return exists( select(CaptureScan.id).where( + *sightings, CaptureScan.organization_id == fence.organization_id, CaptureScan.sync_id == fence.sync_id, CaptureScan.cycle_id == state.version.cycle_id, @@ -384,7 +390,7 @@ async def terminal_checkpoint( async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) -> CaptureCycle: - """Publish the checkpoint only when SQL proves all currently visible scopes complete.""" + """Publish only when every parent eligible under the declared membership is complete.""" cursor, state = await attest_cycle(db, request.fence, request.expected.cycle_id) if state.version != request.expected: raise CycleConflict("Cycle changed before completion") @@ -403,6 +409,13 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - Entity.record_revision > 0, Entity.deleted_at.is_(None), content_is_available(), + # The same current-sweep parent/ancestor predicate governs admission + # and frontier selection. Retained mode preserves provider semantics. + *( + (membership_ready(request.fence, state),) + if state.configuration.membership == "observed" + else () + ), ~child_scope_complete(request.fence, state, child_type), ) .limit(1) @@ -418,7 +431,8 @@ async def complete_cycle(db: AsyncSession, sync: Sync, request: CompleteCycle) - completed_at=now, configuration_digest=state.configuration.digest(), discovery="scope_enumeration_complete" - if all( + if state.configuration.membership == "retained" + and all( state.configuration.policy(kind) == "exhaustive" for kind in state.configuration.parents ) diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py index c986476f69..b95b66fe42 100644 --- a/backend/airweave/domains/native_ingestion/import_store.py +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -132,6 +132,7 @@ async def start( fingerprint=fingerprint, parents=parents, completion_policies={kind: policy 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( 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() + ) From 278965a2ba4693bfdbb86ec8755708b164b29c52 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 02:59:54 -0700 Subject: [PATCH 197/266] feat: expose native scope capture and recovery workflow --- .../api/v1/endpoints/native_imports.py | 89 +++++++++- .../domains/native_ingestion/README.md | 19 ++ .../native_ingestion/import_service.py | 60 +++++++ .../domains/native_ingestion/scope_models.py | 42 +++++ .../domains/native_ingestion/scope_store.py | 116 ++++++++++++ .../native_ingestion/tests/test_import_api.py | 6 +- .../native_ingestion/tests/test_scope_api.py | 168 ++++++++++++++++++ 7 files changed, 496 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/scope_models.py create mode 100644 backend/airweave/domains/native_ingestion/scope_store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_scope_api.py diff --git a/backend/airweave/api/v1/endpoints/native_imports.py b/backend/airweave/api/v1/endpoints/native_imports.py index 424a01d8df..11cd382729 100644 --- a/backend/airweave/api/v1/endpoints/native_imports.py +++ b/backend/airweave/api/v1/endpoints/native_imports.py @@ -12,9 +12,16 @@ from airweave.api.router import TrailingSlashRouter from airweave.core.container import Container from airweave.db.session import get_db -from airweave.domains.entities.canonical.store import WriterBusy +from airweave.domains.entities.canonical.store import CanonicalStoreError, WriterBusy from airweave.domains.native_ingestion.errors import NativeAdmissionError, NativeImportNotFound from airweave.domains.native_ingestion.import_models import NativeImportState, StartNativeImport +from airweave.domains.native_ingestion.page_models import CommitNativePage, NativePageAck +from airweave.domains.native_ingestion.scope_models import ( + BeginNativeScope, + NativeScopeRef, + NativeScopeState, + ReconcileNativeScope, +) router = TrailingSlashRouter() RequestKey = Annotated[str, Path(min_length=1, max_length=128)] @@ -53,3 +60,83 @@ async def read_native_import( raise HTTPException(404, {"code": error.code, "message": str(error)}) from error except NativeAdmissionError as error: raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.put("/{source_id}/imports/{request_key}/scopes", response_model=NativeScopeState) +async def begin_scope( + source_id: UUID, + request_key: RequestKey, + request: BeginNativeScope, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeScopeState: + """Begin scope using server-held import authority.""" + try: + return await container.native_imports.begin_scope( + db, ctx.organization.id, source_id, request_key, request + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.post("/{source_id}/imports/{request_key}/scopes/read", response_model=NativeScopeState) +async def read_scope( + source_id: UUID, + request_key: RequestKey, + request: NativeScopeRef, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeScopeState: + """Read scope using server-held import authority.""" + try: + return await container.native_imports.read_scope( + db, ctx.organization.id, source_id, request_key, request + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.post("/{source_id}/imports/{request_key}/scopes/reconcile", response_model=NativeScopeState) +async def reconcile_scope( + source_id: UUID, + request_key: RequestKey, + request: ReconcileNativeScope, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeScopeState: + """Reconcile scope using server-held import authority.""" + try: + return await container.native_imports.reconcile_scope( + db, ctx.organization.id, source_id, request_key, request + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.put("/{source_id}/imports/{request_key}/pages", response_model=NativePageAck) +async def page( + source_id: UUID, + request_key: RequestKey, + request: CommitNativePage, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativePageAck: + """Page using server-held import authority.""" + try: + return await container.native_imports.page( + db, ctx.organization.id, source_id, request_key, request + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index f89ce799e1..37b2f0e3d4 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -135,3 +135,22 @@ 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. diff --git a/backend/airweave/domains/native_ingestion/import_service.py b/backend/airweave/domains/native_ingestion/import_service.py index 7ce7201b35..022ba77f85 100644 --- a/backend/airweave/domains/native_ingestion/import_service.py +++ b/backend/airweave/domains/native_ingestion/import_service.py @@ -7,6 +7,16 @@ from airweave.db.unit_of_work import UnitOfWork 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: @@ -15,6 +25,8 @@ class NativeImports: def __init__(self, store: NativeImportStore): """Inject the native import persistence boundary.""" self.store = store + self.scopes = NativeScopeStore(store) + self.pages = NativePageStore(store, NativeIngestionStore(store.canonical)) async def start( self, @@ -34,3 +46,51 @@ async def read( """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) 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..1e19ffd415 --- /dev/null +++ b/backend/airweave/domains/native_ingestion/scope_store.py @@ -0,0 +1,116 @@ +"""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 current active import progress without starting a missing scope.""" + imported = await self.imports.active(db, organization_id, source_id, request_key) + state = await self.scans.read(db, imported.fence, request.scope) + if state is None or state.cycle_id != imported.cycle_id: + raise NativeAdmissionError("Scope has not been started for this import") + 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/tests/test_import_api.py b/backend/airweave/domains/native_ingestion/tests/test_import_api.py index 2b7ab01307..0ebc3cae24 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_import_api.py +++ b/backend/airweave/domains/native_ingestion/tests/test_import_api.py @@ -22,8 +22,8 @@ from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata -@pytest.fixture -async def native_api(database, monkeypatch): +@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")) @@ -46,7 +46,7 @@ async def native_api(database, monkeypatch): source = await NativeSources(sources).ensure( db, org, - EnsureNativeSource(owner_id="owner", dataset="knowledge", collection="native-api"), + 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)) 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 From ff99a07818765b2324633e69e8ee6c5c706ff992 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:04:50 -0700 Subject: [PATCH 198/266] Finalize native imports with durable capture outcomes --- .../api/v1/endpoints/native_imports.py | 38 ++++++ .../domains/native_ingestion/README.md | 18 +++ .../domains/native_ingestion/import_models.py | 31 ++++- .../native_ingestion/import_service.py | 15 +++ .../domains/native_ingestion/import_store.py | 104 ++++++++++++--- .../domains/native_ingestion/scope_store.py | 17 ++- .../tests/test_terminal_api.py | 124 ++++++++++++++++++ 7 files changed, 324 insertions(+), 23 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_terminal_api.py diff --git a/backend/airweave/api/v1/endpoints/native_imports.py b/backend/airweave/api/v1/endpoints/native_imports.py index 11cd382729..aab69ff2b9 100644 --- a/backend/airweave/api/v1/endpoints/native_imports.py +++ b/backend/airweave/api/v1/endpoints/native_imports.py @@ -140,3 +140,41 @@ async def page( raise HTTPException(404, {"code": error.code, "message": str(error)}) from error except CanonicalStoreError as error: raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.post("/{source_id}/imports/{request_key}/complete", response_model=NativeImportState) +async def complete_native_import( + source_id: UUID, + request_key: RequestKey, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeImportState: + """Complete this authorized import; never change another writer.""" + try: + return await container.native_imports.finish( + db, ctx.organization.id, source_id, request_key, cancel=False + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.post("/{source_id}/imports/{request_key}/cancel", response_model=NativeImportState) +async def cancel_native_import( + source_id: UUID, + request_key: RequestKey, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeImportState: + """Cancel this authorized import; never change another writer.""" + try: + return await container.native_imports.finish( + db, ctx.organization.id, source_id, request_key, cancel=True + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 37b2f0e3d4..b2a37d534d 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -154,3 +154,21 @@ 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. diff --git a/backend/airweave/domains/native_ingestion/import_models.py b/backend/airweave/domains/native_ingestion/import_models.py index 20dce9255f..13216a13ac 100644 --- a/backend/airweave/domains/native_ingestion/import_models.py +++ b/backend/airweave/domains/native_ingestion/import_models.py @@ -3,7 +3,7 @@ from typing import Literal from uuid import UUID -from pydantic import Field +from pydantic import AwareDatetime, Field, model_validator from airweave.core.shared_models import SyncJobStatus from airweave.domains.entities.canonical.requests import WriterFence @@ -17,6 +17,26 @@ class StartNativeImport(NativeModel): coverage: Literal["bounded", "complete"] +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.""" @@ -26,6 +46,14 @@ class NativeImportReceipt(NativeModel): 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): @@ -37,3 +65,4 @@ class NativeImportState(NativeModel): 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 index 022ba77f85..97b0094b6b 100644 --- a/backend/airweave/domains/native_ingestion/import_service.py +++ b/backend/airweave/domains/native_ingestion/import_service.py @@ -94,3 +94,18 @@ async def page( """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 + ) diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py index b95b66fe42..f9bc9caa29 100644 --- a/backend/airweave/domains/native_ingestion/import_store.py +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -6,17 +6,19 @@ from uuid import UUID, uuid4, uuid5 from pydantic import ValidationError -from sqlalchemy import select +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, @@ -26,9 +28,12 @@ from airweave.domains.native_ingestion.import_models import ( NativeImportReceipt, NativeImportState, + NativeImportSummary, StartNativeImport, ) +from airweave.domains.native_ingestion.source_models import LockedNativeSource from airweave.domains.native_ingestion.source_store import 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. @@ -49,7 +54,15 @@ def receipt(job: SyncJob, source_id: UUID, request_key: str) -> NativeImportRece except ValidationError as error: raise NativeAdmissionError("Native import receipt is malformed") from error if ( - value.source_id != source_id + ( + 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 @@ -68,6 +81,7 @@ def import_state(job: SyncJob, value: NativeImportReceipt) -> NativeImportState: request=value.request, status=job.status, cycle_id=value.cycle_id, + summary=value.summary, ) @@ -180,17 +194,8 @@ async def read( request_key: str, ) -> NativeImportState: """Authorize actual source/job association, including terminal import retries.""" - 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, - ) - ) - if job is None: - raise NativeImportNotFound("Native import does not exist") - return import_state(job, receipt(job, source_id, request_key)) + _, job, saved = await self.load(db, organization_id, source_id, request_key) + return import_state(job, saved) async def active( self, @@ -200,9 +205,17 @@ async def active( request_key: str, ) -> NativeImportReceipt: """Resolve stored authority for a mutation, never a publisher-supplied fence.""" - bound = await self.sources.require(db, organization_id, source_id) + 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( @@ -215,6 +228,63 @@ async def active( ) if job is None: raise NativeImportNotFound("Native import does not exist") - saved = receipt(job, source_id, request_key) - await self.canonical._fenced_sync(db, saved.fence) - return saved + 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/scope_store.py b/backend/airweave/domains/native_ingestion/scope_store.py index 1e19ffd415..401ed0b12f 100644 --- a/backend/airweave/domains/native_ingestion/scope_store.py +++ b/backend/airweave/domains/native_ingestion/scope_store.py @@ -85,11 +85,18 @@ async def read( request_key: str, request: NativeScopeRef, ) -> NativeScopeState: - """Read current active import progress without starting a missing scope.""" - imported = await self.imports.active(db, organization_id, source_id, request_key) - state = await self.scans.read(db, imported.fence, request.scope) - if state is None or state.cycle_id != imported.cycle_id: - raise NativeAdmissionError("Scope has not been started for this import") + """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": "exhaustive" + if imported.request.coverage == "complete" + else "discovery_only" + } + ) return scope_state(state) async def reconcile( 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 From fd8fb649f502ccf39a3c703d4c5c0d1185ba8bb7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:07:16 -0700 Subject: [PATCH 199/266] fix: route native job cancellation through import ownership --- .../native_ingestion/tests/test_imports.py | 13 +++++++ .../domains/owned_provisioning/guard.py | 16 +++++++++ backend/airweave/domains/syncs/service.py | 3 ++ .../domains/syncs/tests/test_service.py | 34 +++++++++++++++---- 4 files changed, 60 insertions(+), 6 deletions(-) diff --git a/backend/airweave/domains/native_ingestion/tests/test_imports.py b/backend/airweave/domains/native_ingestion/tests/test_imports.py index d3b13ea6a4..14d09f7cba 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_imports.py +++ b/backend/airweave/domains/native_ingestion/tests/test_imports.py @@ -139,3 +139,16 @@ async def fail(*args, **kwargs): db, uuid4(), native.source_connection_id, "request-one" ) assert error.value.status_code == 404 + + +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/owned_provisioning/guard.py b/backend/airweave/domains/owned_provisioning/guard.py index b9e8469baa..4193a9961c 100644 --- a/backend/airweave/domains/owned_provisioning/guard.py +++ b/backend/airweave/domains/owned_provisioning/guard.py @@ -7,6 +7,7 @@ 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: @@ -30,3 +31,18 @@ async def require_unmanaged_source( 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/syncs/service.py b/backend/airweave/domains/syncs/service.py index 009aac357e..2e3a3e577c 100644 --- a/backend/airweave/domains/syncs/service.py +++ b/backend/airweave/domains/syncs/service.py @@ -32,6 +32,7 @@ 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 @@ -315,6 +316,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( diff --git a/backend/airweave/domains/syncs/tests/test_service.py b/backend/airweave/domains/syncs/tests/test_service.py index c4bb68b24e..9a27deed49 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 @@ -1215,3 +1221,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() From 53d819da1b50278745b7f0169c3e278d92d99fc0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:10:24 -0700 Subject: [PATCH 200/266] test: surface early search failures and clean up interrupted probes --- .../canonical/tests/test_search_sessions.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py b/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py index 8894006e02..c322b93416 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_search_sessions.py @@ -161,14 +161,26 @@ async def interrupt_after_sql(db, *args): json={"query": "budget", "sync_ids": [str(fence.sync_id)], "mode": "keyword"}, ) ) - await asyncio.wait_for(entered.wait(), timeout=5) - if outcome == "cancellation": + 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() - with pytest.raises(asyncio.CancelledError): - await task - else: - with pytest.raises(RuntimeError, match="synthetic enrichment failure"): - await task + await asyncio.gather(waiting, task, return_exceptions=True) await connection_is_free(single_connection) From 7c9bf4c1dc854f108531c94284b1c21e7fdac28b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:13:28 -0700 Subject: [PATCH 201/266] docs: describe current native import lifecycle without stale gates --- .../domains/native_ingestion/README.md | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index b2a37d534d..7a8d6c16eb 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -1,16 +1,16 @@ # Native snapshot admission -This internal service retains versioned Almanac snapshots in the existing -canonical store. It does not provision sources, authenticate callers, enumerate -Almanac data, or provide a public import endpoint. -Almanac remains the authority for content and access. - -The future trusted publisher must bind a source connection with -`short_name=almanac` and `config_fields={owner_id, dataset}`. Datasets are -`knowledge` and `sessions`. The service checks that binding against the fenced -organization/sync and every snapshot. It does not accept a fabricated Composio -account. A future API must issue the fence after authenticating the publisher; -the internal command itself is not an authorization credential. +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 @@ -65,10 +65,10 @@ changes-feed contract. Empty pages are valid; only `final=true` begins the exist reconciliation phase, and source completion remains a separate cycle barrier. Admission, capture, sightings and continuation either commit together or roll back. -HTTP authentication, request receipts, import provisioning, explicit renewal and -publisher completeness attestation remain unfinished. Cancellation must resolve -the job's actual bound source; a caller-supplied source ID is insufficient evidence -for authorizing a job cancellation. +HTTP authentication, receipts, provisioning and completion/cancellation are +implemented below. Explicit access renewal and authoritative publisher +completeness remain pending. Native jobs cannot be cancelled through the generic +provider route: the guard checks the actual job's sync identity. ## Native search projection @@ -92,8 +92,8 @@ 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 trusted publisher and -owner-binding API must still be implemented before exposing native ingestion. +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 @@ -103,8 +103,7 @@ 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` is the internal start/read transaction primitive, not yet an -HTTP publisher API. A bounded request key derives an existing SyncJob primary +`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. @@ -115,8 +114,9 @@ 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: final scopes, stable export boundaries, receipts, -renewal and import completion/cancellation still need their HTTP lifecycle. +itself certify completeness. The HTTP lifecycle verifies final scopes and +receipts; the publisher must establish stable export boundaries. Explicit access +renewal remains a separate unfinished capability. ## Import HTTP and page receipts @@ -126,8 +126,8 @@ 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` remains an internal transaction primitive until scope -endpoints and completion are wired. It resolves stored writer authority, admits +`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 From 0e5604f34d2383004e1fa675051786e7d29333ba Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:14:42 -0700 Subject: [PATCH 202/266] test: align Vespa fixture with embedding purpose contract --- .../entities/canonical/tests/test_real_vespa_projection.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 index 0e4b1d29c9..4cb42c317f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py @@ -19,7 +19,7 @@ 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, SparseEmbedding +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 @@ -46,7 +46,9 @@ class FixedDense(FakeDenseEmbedder): """Only model inference is substituted; all processing stays production code.""" - async def embed_many(self, texts): + async def embed_many( + self, texts: list[str], *, purpose: EmbeddingPurpose = "document" + ) -> list[DenseEmbedding]: return [DenseEmbedding(vector=[1.0] + [0.0] * 383) for _ in texts] From b4678a188004e0d3acf3ff88c281aabc750bc9f2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:25:46 -0700 Subject: [PATCH 203/266] Publish staged native snapshots with resumable bounded imports --- .../native_ingestion/tests/test_publisher.py | 148 +++++++++ backend/pytest.ini | 2 +- evaluation/README.md | 13 + evaluation/native_import.py | 311 ++++++++++++++++++ 4 files changed, 473 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_publisher.py create mode 100644 evaluation/native_import.py 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/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/evaluation/README.md b/evaluation/README.md index 03811c2522..b4e6b29728 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -112,3 +112,16 @@ snippet evaluation and corpus-freezing automation are still separate work. The four focused synthetic tests verify evaluator behavior, not search quality. Reference: [ir-measures interfaces and formats](https://ir-measur.es/en/latest/getting-started.html). + +## Bounded native publishing helper + +`evaluation.native_import.publish_native` publishes already-staged `NativeSnapshot` +values through the authenticated native HTTP API; it neither discovers nor fetches +source content. Run with the backend environment and `PYTHONPATH=backend:.` from the +repository root. The injected `httpx.AsyncClient` base URL must target the API root +with a trailing slash. Retain the exact input and request key after a transport error: +reinvoking resumes the durable cursor; changed input with that key conflicts. The +returned terminal summary certifies bounded capture, not search publication. + +Tests live in the backend native-ingestion suite and require disposable PostgreSQL. +The standalone retrieval-metrics test environment does not import this helper. diff --git a/evaluation/native_import.py b/evaluation/native_import.py new file mode 100644 index 0000000000..1ed985288d --- /dev/null +++ b/evaluation/native_import.py @@ -0,0 +1,311 @@ +"""Bounded staged-native publisher; retry by retaining the same input and request key.""" + +import hashlib +import json +from typing import Literal +from urllib.parse import quote +from uuid import UUID, uuid5 + +import httpx +from pydantic import Field, ValidationError + +from airweave.domains.entities.canonical.requests import CompletedScope +from airweave.domains.native_ingestion.import_models import ( + NativeImportState, + StartNativeImport, +) +from airweave.domains.native_ingestion.models import NativeModel, NativeSnapshot +from airweave.domains.native_ingestion.page_models import ( + CommitNativePage, + NativePageAck, +) +from airweave.domains.native_ingestion.scope_models import ( + BeginNativeScope, + NativeScopeState, + ReconcileNativeScope, +) +from airweave.domains.native_ingestion.source_models import ( + EnsureNativeSource, + NativeSource, +) + +_PAGE_NAMESPACE = UUID("3c2ba962-a612-4ca3-a29b-342521fd35bc") +_PAGE_SIZE = 500 + + +class NativePublishError(ValueError): + """Publication cannot safely resume; no original content appears in diagnostics.""" + + +class _Cursor(NativeModel): + snapshot_id: str + next_offset: int = Field(ge=0) + + +class _Scope(NativeModel): + scope: CompletedScope + snapshots: tuple[NativeSnapshot, ...] + + +def _prepare( + owner: str, + dataset: Literal["knowledge", "sessions"], + snapshots: tuple[NativeSnapshot, ...], +) -> tuple[_Scope, ...]: + # Validate copied/constructed models again before any server mutation. + checked = tuple( + NativeSnapshot.model_validate(item.model_dump(mode="json")) + for item in snapshots + ) + unique = set() + for item in checked: + key = (item.identity.record_type, item.identity.entity_key) + if key in unique or item.owner_id != owner or item.operation != "upsert": + raise NativePublishError( + "Staged native identity, owner or operation is invalid" + ) + unique.add(key) + allowed = ("knowledge",) if dataset == "knowledge" else ("session", "message") + if item.identity.record_type not in allowed: + raise NativePublishError("Staged native dataset is mixed") + kind = "knowledge" if dataset == "knowledge" else "session" + roots = tuple( + sorted( + (item for item in checked if item.identity.record_type == kind), + key=lambda item: item.identity.native_id, + ) + ) + scopes = [_Scope(scope=CompletedScope(record_type=kind), snapshots=roots)] + if dataset == "sessions": + parents = {item.identity.native_id: item for item in roots} + children: dict[str, list[NativeSnapshot]] = {key: [] for key in parents} + for item in checked: + if item.identity.record_type != "message": + continue + parent = parents.get(item.identity.container_id) + if ( + parent is None + or item.parent != parent.identity + or item.version != parent.version + ): + raise NativePublishError( + "Staged message has no matching session/version" + ) + children[parent.identity.native_id].append(item) + for parent in roots: + scopes.append( + _Scope( + scope=CompletedScope( + record_type="message", + container_id=parent.identity.native_id, + parent=parent.identity, + ), + snapshots=tuple( + sorted( + children[parent.identity.native_id], + key=lambda item: item.identity.native_id, + ) + ), + ) + ) + return tuple(scopes) + + +async def _call[T: NativeModel]( + client: httpx.AsyncClient, + method: str, + path: str, + model: type[T], + body: NativeModel | None = None, +) -> T: + response = await client.request( + method, path, json=body.model_dump(mode="json") if body else None + ) + if response.status_code != 200: + raise NativePublishError( + f"Native destination rejected {method} with HTTP {response.status_code}" + ) + try: + return model.model_validate_json(response.content) + except ValidationError: + raise NativePublishError( + "Native destination returned malformed progress" + ) from None + + +def _attest_import( + state: NativeImportState, source: NativeSource, key: str, digest: str +) -> None: + if ( + state.source_id != source.source_connection_id + or state.request_key != key + or state.request.snapshot_id != digest + or state.request.coverage != "bounded" + ): + raise NativePublishError("Native import response differs from staged intent") + if state.status == "completed" and ( + state.summary is None + or not state.summary.capture_complete + or state.summary.outcome != "completed" + or state.summary.coverage != "bounded" + ): + raise NativePublishError( + "Completed import has no valid bounded capture summary" + ) + + +def _attest_scope( + state: NativeScopeState, scope: _Scope, imported: NativeImportState +) -> None: + if ( + state.cycle_id != imported.cycle_id + or state.scope != scope.scope + or state.coverage != "bounded" + ): + raise NativePublishError("Native scope does not belong to this bounded import") + + +async def publish_native( + client: httpx.AsyncClient, + *, + owner_id: str, + dataset: Literal["knowledge", "sessions"], + collection: str, + request_key: str, + snapshots: tuple[NativeSnapshot, ...], +) -> NativeImportState: + """Publish already staged originals; transport failures require identical caller retry. + + The authenticated client base URL targets the destination API root. No source + reads, credentials, discovery, automatic retries or indexing claims are involved. + """ + ensure = EnsureNativeSource( + owner_id=owner_id, dataset=dataset, collection=collection + ) + if not 1 <= len(request_key) <= 128: + raise NativePublishError("Native request key must contain 1 to 128 characters") + scopes = _prepare(owner_id, dataset, snapshots) + digest = hashlib.sha256( + json.dumps( + { + "publisher_version": 1, + "page_size": _PAGE_SIZE, + "source": ensure.model_dump(mode="json"), + "scopes": [scope.model_dump(mode="json") for scope in scopes], + }, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode() + ).hexdigest() + source = await _call(client, "PUT", "native/sources", NativeSource, ensure) + if source.binding != ensure.binding() or source.collection != collection: + raise NativePublishError( + "Native destination source binding differs from staged input" + ) + base = f"native/sources/{source.source_connection_id}/imports/{quote(request_key, safe='')}" + started = await _call( + client, + "PUT", + base, + NativeImportState, + StartNativeImport(snapshot_id=digest, coverage="bounded"), + ) + _attest_import(started, source, request_key, digest) + if started.status == "completed": + return started + if started.status != "running": + raise NativePublishError("Native import is terminal or unavailable") + for scope in scopes: + state = await _call( + client, + "PUT", + base + "/scopes", + NativeScopeState, + BeginNativeScope(scope=scope.scope), + ) + _attest_scope(state, scope, started) + if not state.cursor: + if state.phase != "collecting" or state.last_page is not None: + raise NativePublishError( + "Native scope is missing its staged input cursor" + ) + offset = 0 + else: + try: + cursor = _Cursor.model_validate(state.cursor) + except ValidationError: + raise NativePublishError("Native scope cursor is malformed") from None + offset = cursor.next_offset + if cursor.snapshot_id != digest or offset > len(scope.snapshots): + raise NativePublishError( + "Native scope cursor differs from staged input" + ) + if state.phase != "collecting" and offset != len(scope.snapshots): + raise NativePublishError("Native scope ended before the staged input") + if state.phase == "collecting" and ( + offset == len(scope.snapshots) or offset % _PAGE_SIZE + ): + raise NativePublishError( + "Native scope cursor is not a deterministic page boundary" + ) + while state.phase == "collecting": + end = min(offset + _PAGE_SIZE, len(scope.snapshots)) + page_id = uuid5( + _PAGE_NAMESPACE, + json.dumps( + [digest, scope.scope.model_dump(mode="json"), offset], + sort_keys=True, + ), + ) + ack = await _call( + client, + "PUT", + base + "/pages", + NativePageAck, + CommitNativePage( + page_id=page_id, + scope=scope.scope, + expected=state.version, + snapshots=scope.snapshots[offset:end], + cursor=_Cursor(snapshot_id=digest, next_offset=end).model_dump( + mode="json" + ), + final=end == len(scope.snapshots), + ), + ) + if ( + ack.page_id != page_id + or ack.version.sweep_id != state.version.sweep_id + or ack.version.revision != state.version.revision + 1 + or ack.phase + != ("reconciling" if end == len(scope.snapshots) else "collecting") + ): + raise NativePublishError( + "Native destination acknowledged inconsistent page progress" + ) + state = state.model_copy( + update={"phase": ack.phase, "version": ack.version} + ) + offset = end + if state.phase == "reconciling": + state = await _call( + client, + "POST", + base + "/scopes/reconcile", + NativeScopeState, + ReconcileNativeScope(scope=scope.scope, expected=state.version), + ) + _attest_scope(state, scope, started) + if state.phase != "complete": + raise NativePublishError("Bounded native scope did not complete") + completed = await _call(client, "POST", base + "/complete", NativeImportState) + _attest_import(completed, source, request_key, digest) + if ( + completed.import_id != started.import_id + or completed.cycle_id != started.cycle_id + or completed.status != "completed" + ): + raise NativePublishError("Native destination did not complete this import") + return completed From 955c65b1d83674a61ce0d787bb391b3fde9083db Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:55:56 -0700 Subject: [PATCH 204/266] Add bounded native import operator CLI without server credentials --- .../domains/native_ingestion/import_store.py | 3 +- .../domains/native_ingestion/source_models.py | 24 +-- .../domains/native_ingestion/source_store.py | 24 ++- .../tests/test_publish_cli.py | 171 ++++++++++++++++ evaluation/README.md | 35 ++++ evaluation/native_import_cli.py | 183 ++++++++++++++++++ 6 files changed, 413 insertions(+), 27 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_publish_cli.py create mode 100644 evaluation/native_import_cli.py diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py index f9bc9caa29..fb50cdfcd0 100644 --- a/backend/airweave/domains/native_ingestion/import_store.py +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -31,8 +31,7 @@ NativeImportSummary, StartNativeImport, ) -from airweave.domains.native_ingestion.source_models import LockedNativeSource -from airweave.domains.native_ingestion.source_store import NativeSourceStore +from airweave.domains.native_ingestion.source_store import LockedNativeSource, NativeSourceStore from airweave.models.capture_scan import CaptureScan from airweave.models.sync_job import SyncJob diff --git a/backend/airweave/domains/native_ingestion/source_models.py b/backend/airweave/domains/native_ingestion/source_models.py index 31541197d9..0216e14eec 100644 --- a/backend/airweave/domains/native_ingestion/source_models.py +++ b/backend/airweave/domains/native_ingestion/source_models.py @@ -3,11 +3,9 @@ import json from uuid import UUID, uuid5 -from pydantic import BaseModel, ConfigDict, Field +from pydantic import Field from airweave.domains.native_ingestion.models import NativeModel, NativeSourceBinding -from airweave.models.source_connection import SourceConnection -from airweave.models.sync import Sync # Versioned protocol constants: changing these would create duplicate source identities. _SOURCE_NAMESPACE_V1 = UUID("b13e2720-64ab-4ac5-9a0f-6264e2fb783f") @@ -48,23 +46,3 @@ class NativeSource(NativeModel): binding: NativeSourceBinding collection: str available: bool - - -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, - ) diff --git a/backend/airweave/domains/native_ingestion/source_store.py b/backend/airweave/domains/native_ingestion/source_store.py index 20e5e24b5d..baf18b3de4 100644 --- a/backend/airweave/domains/native_ingestion/source_store.py +++ b/backend/airweave/domains/native_ingestion/source_store.py @@ -3,7 +3,7 @@ from uuid import UUID from fastapi import HTTPException -from pydantic import ValidationError +from pydantic import BaseModel, ConfigDict, ValidationError from sqlalchemy import select from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession @@ -11,7 +11,7 @@ from airweave.domains.native_ingestion.models import NativeSourceBinding from airweave.domains.native_ingestion.source_models import ( EnsureNativeSource, - LockedNativeSource, + NativeSource, native_source_id, native_sync_id, ) @@ -20,6 +20,26 @@ 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.""" 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/evaluation/README.md b/evaluation/README.md index b4e6b29728..4f33ac6235 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -125,3 +125,38 @@ returned terminal summary certifies bounded capture, not search publication. Tests live in the backend native-ingestion suite and require disposable PostgreSQL. The standalone retrieval-metrics test environment does not import this helper. + +### Operator command + +From the repository root, using the backend Python environment: + +```sh +PYTHONPATH=backend:. backend/.venv/bin/python -m evaluation.native_import_cli \ + --url http://127.0.0.1:18086/api/v1/ \ + --input staged-knowledge.json --owner OWNER_ID --dataset knowledge \ + --collection COLLECTION_ID --request-key stable-import-key +``` + +Supply `AIRWEAVE_API_KEY` through the process environment (for example a secret +manager); there is no key argument or dotenv loading. No server database/settings +credentials are needed. Use the backend API key authorized to attest this owner. +The file is a UTF-8 JSON array of `NativeSnapshot` values, at most **32 MiB**. +Use `--dataset sessions` for staged session roots and their original messages. +This command publishes the supplied bounded set, not an account-wide backfill; +it does not infer deletions from records omitted from the file. + +External URLs must use HTTPS; HTTP is allowed only for loopback IPs or localhost. +URL credentials, queries and fragments are rejected. Redirects and environment +proxies are disabled. Each HTTP operation has a 10-second connect timeout and +60-second read/write/pool timeout, not a whole-import deadline. No automatic +retry occurs. Keep the exact staged file and request key on failure; repeat the +same command to recover durable progress. Do not change the file under that key. + +Success writes one JSON capture summary to stdout, including +`capture_complete: true`, `coverage: "bounded"`, and `indexing: "not_verified"`. +Search publication must be verified separately. Errors use stderr without raw +payloads, keys, file paths or destination URLs. Exit codes: 0 completed capture; +2 invalid configuration/input; 3 unknown network outcome; 4 destination rejection +or inconsistent import progress (including cross-record preflight rejection); +130 interrupted; 1 unexpected failure. Failed commands may have committed earlier +pages: a nonzero exit is not evidence of rollback. diff --git a/evaluation/native_import_cli.py b/evaluation/native_import_cli.py new file mode 100644 index 0000000000..fb6fc31216 --- /dev/null +++ b/evaluation/native_import_cli.py @@ -0,0 +1,183 @@ +"""Publish one bounded staged snapshot using the existing durable native importer.""" + +import argparse +import asyncio +import ipaddress +import sys +from pathlib import Path +from typing import Literal +from urllib.parse import urlsplit + +import httpx +from pydantic import ( + Field, + HttpUrl, + SecretStr, + TypeAdapter, + ValidationError, + field_validator, +) +from pydantic_settings import BaseSettings, SettingsConfigDict + +from airweave.domains.native_ingestion.models import NativeModel, NativeSnapshot +from evaluation.native_import import NativePublishError, publish_native + +MAX_INPUT_BYTES = 32 * 1024 * 1024 +SNAPSHOTS = TypeAdapter(tuple[NativeSnapshot, ...]) + + +class Settings(BaseSettings): + """Read authentication only from the process environment, never dotenv files.""" + + model_config = SettingsConfigDict(env_prefix="AIRWEAVE_", extra="ignore") + api_key: SecretStr + + @field_validator("api_key") + @classmethod + def valid_key(cls, value: SecretStr) -> SecretStr: + """Require a nonempty printable header value without exposing it.""" + key = value.get_secret_value() + if not key or any(ord(char) < 33 or ord(char) > 126 for char in key): + raise ValueError("Invalid API key") + return value + + +class PublishArguments(NativeModel): + """CLI input boundary; all state and retry decisions belong to the publisher.""" + + url: str + input: Path + owner: str = Field(min_length=1) + dataset: Literal["knowledge", "sessions"] + collection: str = Field(min_length=1, max_length=255) + request_key: str = Field(min_length=1, max_length=128) + + @field_validator("url") + @classmethod + def destination_url(cls, value: str) -> str: + """Allow explicit TLS destinations and local HTTP evaluation only.""" + if any(char.isspace() or ord(char) < 32 for char in value): + raise ValueError("Invalid destination") + parsed = urlsplit(value) + # Reject even empty query/fragment delimiters and URL-carried credentials. + if any(mark in value for mark in ("?", "#")) or parsed.username is not None: + raise ValueError( + "Destination must not include credentials, query or fragment" + ) + url = TypeAdapter(HttpUrl).validate_python(value) + if url.username is not None or url.password is not None: + raise ValueError("Destination must not include credentials") + if url.scheme == "http": + hostname = parsed.hostname or "" + try: + loopback = ipaddress.ip_address(hostname).is_loopback + except ValueError: + loopback = hostname == "localhost" + if not loopback: + raise ValueError("HTTP requires a loopback destination") + return str(url).rstrip("/") + "/" + + +class Parser(argparse.ArgumentParser): + """Argument errors do not echo potentially sensitive caller input.""" + + def error(self, message: str) -> None: + """Print static usage guidance instead of the original parser diagnostic.""" + self.exit(2, "Invalid arguments. Use --help for the accepted options.\n") + + +def parse_args(argv: list[str] | None) -> PublishArguments: + """Parse explicit publisher options into their validated boundary.""" + parser = Parser(description=__doc__) + parser.add_argument( + "--url", required=True, help="API root, e.g. https://host/api/v1/" + ) + parser.add_argument( + "--input", + required=True, + help="JSON array of NativeSnapshot values (max 32 MiB)", + ) + parser.add_argument("--owner", required=True) + parser.add_argument("--dataset", required=True, choices=("knowledge", "sessions")) + parser.add_argument("--collection", required=True) + parser.add_argument( + "--request-key", + required=True, + help="Retain this key and exact input for retries", + ) + return PublishArguments.model_validate(vars(parser.parse_args(argv))) + + +async def publish(arguments: PublishArguments, settings: Settings) -> str: + """Read bounded originals then delegate the entire workflow to the publisher.""" + with arguments.input.open("rb") as staged: + data = staged.read(MAX_INPUT_BYTES + 1) + if len(data) > MAX_INPUT_BYTES: + raise ValueError("Staged input exceeds the size limit") + snapshots = SNAPSHOTS.validate_json(data) + async with httpx.AsyncClient( + base_url=arguments.url, + headers={"X-API-Key": settings.api_key.get_secret_value()}, + follow_redirects=False, + trust_env=False, + timeout=httpx.Timeout(60, connect=10), + ) as client: + state = await publish_native( + client, + owner_id=arguments.owner, + dataset=arguments.dataset, + collection=arguments.collection, + request_key=arguments.request_key, + snapshots=snapshots, + ) + # The publisher attests this summary; never turn absent progress into success. + if state.summary is None: + raise NativePublishError("Destination returned no capture summary") + return state.summary.model_dump_json() + + +def main(argv: list[str] | None = None) -> int: + """Emit only the capture summary or a sanitized diagnostic with a stable exit code.""" + try: + arguments = parse_args(argv) + settings = Settings() + result = asyncio.run(publish(arguments, settings)) + except NativePublishError: + print( + "Publication rejected or progress invalid. Retain input and request key; " + "inspect destination state.", + file=sys.stderr, + ) + return 4 + except httpx.HTTPError: + print( + "Network outcome unknown. Retry only with the identical staged input and request key.", + file=sys.stderr, + ) + return 3 + except (ValidationError, ValueError, OSError): + print( + "Invalid configuration or staged input. " + "Check credentials, destination and JSON (maximum 32 MiB).", + file=sys.stderr, + ) + return 2 + except KeyboardInterrupt: + print( + "Interrupted; outcome may be unknown. Retain the exact staged input and request key.", + file=sys.stderr, + ) + return 130 + except Exception: + print( + "Unexpected failure; outcome may be unknown. " + "Retain the exact staged input and request key.", + file=sys.stderr, + ) + return 1 + print(result) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From cd5d380b46410016e9e92978fbb765ed75c97ae3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 03:59:22 -0700 Subject: [PATCH 205/266] Route native retained records through canonical projection activity --- .../activities/project_canonical_records.py | 6 +- .../tests/test_project_canonical_records.py | 131 ++++++++++++++++++ 2 files changed, 133 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py diff --git a/backend/airweave/domains/temporal/activities/project_canonical_records.py b/backend/airweave/domains/temporal/activities/project_canonical_records.py index 422f17ac86..87dc2d9fdf 100644 --- a/backend/airweave/domains/temporal/activities/project_canonical_records.py +++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py @@ -8,11 +8,10 @@ from airweave.core.logging import logger from airweave.db.session import get_db_context -from airweave.domains.entities.canonical.page_source import CanonicalPageSource 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 CanonicalSource +from airweave.domains.entities.canonical.source import indexed_record_types from airweave.domains.sources.protocols import SourceRegistryProtocol from airweave.models.collection import Collection from airweave.models.source_connection import SourceConnection @@ -49,8 +48,7 @@ async def run(self, organization_id: str, sync_id: str, after_id: str | None = N if row is None: return ProjectionBatchResult().model_dump(mode="json") source_name, collection_id = row - entry = self.source_registry.get(source_name) - if not isinstance(entry.source_class_ref, (CanonicalSource, CanonicalPageSource)): + 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( 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..c2d67a2ac5 --- /dev/null +++ b/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py @@ -0,0 +1,131 @@ +"""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", + 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() From 4fab9621144c0cca249899b0a36bb211ef4d91a5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:01:27 -0700 Subject: [PATCH 206/266] Keep native imports outside provider timeout maintenance --- backend/airweave/crud/crud_sync_job.py | 14 +++++++++- .../tests/test_maintenance_ownership.py | 27 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_maintenance_ownership.py diff --git a/backend/airweave/crud/crud_sync_job.py b/backend/airweave/crud/crud_sync_job.py index 89590122bf..6df91e7d0c 100644 --- a/backend/airweave/crud/crud_sync_job.py +++ b/backend/airweave/crud/crud_sync_job.py @@ -11,6 +11,7 @@ from airweave.core.context import BaseContext from airweave.crud._base_organization import CRUDBaseOrganization from airweave.db.unit_of_work import UnitOfWork +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, SyncJobUpdate @@ -165,7 +166,18 @@ async def get_stuck_jobs_by_status( Returns: List of stuck sync jobs """ - stmt = select(SyncJob).where(SyncJob.status.in_(status)) + # Native imports share durable jobs, but own their explicit resume/cancel + # lifecycle. Provider heartbeat timeouts must never terminate them. + native_source = ( + select(SourceConnection.id) + .where( + SourceConnection.sync_id == SyncJob.sync_id, + SourceConnection.organization_id == SyncJob.organization_id, + SourceConnection.short_name == "almanac", + ) + .exists() + ) + stmt = select(SyncJob).where(SyncJob.status.in_(status), ~native_source) # Apply timestamp filters based on what's provided if modified_before is not None: 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 From 6ad087bba01a783e4831cc1d78c043406df668f9 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:07:53 -0700 Subject: [PATCH 207/266] Select bounded fresh projection work for native recovery --- .../entities/canonical/projection_models.py | 16 +++ .../entities/canonical/projection_store.py | 70 +++++++++- .../domains/entities/canonical/projector.py | 2 + .../tests/test_projection_recovery.py | 129 ++++++++++++++++++ 4 files changed, 210 insertions(+), 7 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py diff --git a/backend/airweave/domains/entities/canonical/projection_models.py b/backend/airweave/domains/entities/canonical/projection_models.py index c23c74a77e..11c8e51804 100644 --- a/backend/airweave/domains/entities/canonical/projection_models.py +++ b/backend/airweave/domains/entities/canonical/projection_models.py @@ -52,6 +52,22 @@ class ProjectionWork(BaseModel): previous_generation: UUID | None +class ProjectionSourceRef(BaseModel): + """Stored tenant/source identity eligible for recovery, without original payloads.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + organization_id: UUID + sync_id: UUID + + +class ProjectionSourcePage(BaseModel): + """Bounded source keyset page; a cursor exists only when another row was observed.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + sources: tuple[ProjectionSourceRef, ...] + next_cursor: UUID | None = None + + class ProjectionBatchResult(BaseModel): """One bounded scan page; failures remain pending for a later retry.""" diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index 85d20fa9b7..a0d56d68b4 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -14,6 +14,8 @@ from airweave.domains.entities.canonical.projection_models import ( ProjectionDocument, ProjectionLocator, + ProjectionSourcePage, + ProjectionSourceRef, ProjectionWork, projection_document_locator, ) @@ -22,6 +24,7 @@ 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.source_connection import SourceConnection from airweave.models.sync import Sync @@ -77,6 +80,19 @@ def publications_match(locators: Iterable[ProjectionLocator]): ) +def _pending_record(): + """Shared exact publication/visibility eligibility for discovery and execution.""" + return and_( + Entity.record_revision > 0, + 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), + ), + or_(Entity.deleted_at.is_not(None), content_is_available()), + ) + + class CanonicalProjectionStore: """Optimistic projection computation, serialized CAS against source capture.""" @@ -88,6 +104,7 @@ async def pending( *, after_id: UUID | None = None, limit: int = 25, + skip_failed: bool = False, ) -> tuple[ProjectionWork, ...]: """One bounded UUID page; deleted records publish an empty generation.""" if not 1 <= limit <= 100: @@ -99,15 +116,11 @@ async def pending( Entity.organization_id == organization_id, Entity.sync_id == sync_id, Sync.organization_id == organization_id, - Entity.record_revision > 0, - 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), - ), - or_(Entity.deleted_at.is_not(None), content_is_available()), + _pending_record(), ) ) + if skip_failed: + statement = statement.where(Entity.projection_error.is_(None)) if after_id is not None: statement = statement.where(Entity.id > after_id) rows = await db.execute(statement.order_by(Entity.id).limit(limit + 1)) @@ -121,6 +134,49 @@ async def pending( for row, version in rows ) + async def pending_sources( + self, + db: AsyncSession, + source_names: tuple[str, ...], + *, + after_id: UUID | None = None, + limit: int = 20, + ) -> ProjectionSourcePage: + """Discover fresh pending work fairly; failures require an explicit retry. + + This system-worker query crosses tenants but returns their exact stored + scope. Execution must revalidate that scope before reading any originals. + """ + if not 1 <= limit <= 100: + raise ValueError("Projection source page size must be between 1 and 100") + owned_source = exists( + select(SourceConnection.id).where( + SourceConnection.sync_id == Sync.id, + SourceConnection.organization_id == Sync.organization_id, + SourceConnection.short_name.in_(source_names), + ) + ) + fresh_work = exists( + select(Entity.id).where( + Entity.sync_id == Sync.id, + Entity.organization_id == Sync.organization_id, + _pending_record(), + Entity.projection_error.is_(None), + ) + ) + statement = select(Sync.organization_id, Sync.id).where(owned_source, fresh_work) + if after_id is not None: + statement = statement.where(Sync.id > after_id) + rows = (await db.execute(statement.order_by(Sync.id).limit(limit + 1))).all() + sources = tuple( + ProjectionSourceRef(organization_id=organization, sync_id=sync) + for organization, sync in rows[:limit] + ) + return ProjectionSourcePage( + sources=sources, + next_cursor=sources[-1].sync_id if len(rows) > limit else None, + ) + async def _current(self, db: AsyncSession, work: ProjectionWork) -> Entity | None: # Capture locks Sync first too: parent visibility and revision cannot change # between validation and publication. This transaction contains no provider I/O. diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 2a4bbc7d30..7c8b0a3cdc 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -179,6 +179,7 @@ async def batch( *, after_id: UUID | None = None, limit: int = 25, + skip_failed: bool = False, ) -> ProjectionBatchResult: """Failed rows stay pending; other rows in the page continue to make progress.""" async with self._sessions() as db: @@ -188,6 +189,7 @@ async def batch( sync_id, after_id=after_id, limit=limit, + **({"skip_failed": True} if skip_failed else {}), ) published = superseded = failed = 0 for work in pending[:limit]: 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..bf656a88d5 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py @@ -0,0 +1,129 @@ +"""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 capture, observation +from airweave.models import Entity, Organization, SourceConnection, Sync, SyncJob + + +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")) + await db.commit() + expected = set() + for case in ( + "fresh", + "fresh-two", + "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 "almanac", + is_authenticated=False, + ) + ) + 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", "fresh-two", "tombstone"): + expected.add(sync_id) + store = CanonicalProjectionStore() + async with database() as db: + first = await store.pending_sources(db, ("almanac",), limit=2) + assert len(first.sources) == 2 + assert first.next_cursor == first.sources[-1].sync_id + second = await store.pending_sources(db, ("almanac",), 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 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() From 0513c3784a8bc140d9ff5336d2460c84eba9ffe0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:13:14 -0700 Subject: [PATCH 208/266] Recover pending native publications through existing Temporal maintenance --- .../activities/discover_native_projection.py | 24 +++ .../activities/project_canonical_records.py | 11 +- .../temporal/worker/tests/test_wiring.py | 6 +- .../domains/temporal/worker/wiring.py | 8 + .../workflows/cleanup_stuck_sync_jobs.py | 35 ++-- .../workflows/project_canonical_records.py | 14 +- .../workflows/recover_native_projection.py | 51 ++++++ .../workflows/tests/test_cleanup_workflows.py | 14 +- .../tests/test_native_projection_recovery.py | 165 ++++++++++++++++++ 9 files changed, 310 insertions(+), 18 deletions(-) create mode 100644 backend/airweave/domains/temporal/activities/discover_native_projection.py create mode 100644 backend/airweave/domains/temporal/workflows/recover_native_projection.py create mode 100644 backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py 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..96d20a3570 --- /dev/null +++ b/backend/airweave/domains/temporal/activities/discover_native_projection.py @@ -0,0 +1,24 @@ +"""Discover committed native work without a second dispatch journal.""" + +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 + + +class DiscoverNativeProjectionActivity: + """Read the existing pending state; no credentials or content cross Temporal.""" + + @activity.defn(name="discover_native_projection_activity") + async def run(self, after_id: str | None = None) -> dict: + """Native HTTP imports lack the provider workflow's durable projection child.""" + async with get_db_context() as db: + page = await CanonicalProjectionStore().pending_sources( + db, + source_names=("almanac",), + 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 index 87dc2d9fdf..1b3dd5ad7f 100644 --- a/backend/airweave/domains/temporal/activities/project_canonical_records.py +++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py @@ -26,7 +26,13 @@ class ProjectCanonicalRecordsActivity: 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) -> dict: + 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) @@ -57,6 +63,7 @@ async def run(self, organization_id: str, sync_id: str, after_id: str | None = N 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") @@ -67,6 +74,7 @@ async def run(self, organization_id: str, sync_id: str, after_id: str | None = N soft_fail=False, ) try: + options = {"skip_failed": True} if skip_failed else {} result = await self.projector.batch( organization, sync, @@ -74,6 +82,7 @@ async def run(self, organization_id: str, sync_id: str, after_id: str | None = N destination, logger, after_id=UUID(after_id) if after_id else None, + **options, ) return result.model_dump(mode="json") finally: diff --git a/backend/airweave/domains/temporal/worker/tests/test_wiring.py b/backend/airweave/domains/temporal/worker/tests/test_wiring.py index 9d490271cc..cb5f87774e 100644 --- a/backend/airweave/domains/temporal/worker/tests/test_wiring.py +++ b/backend/airweave/domains/temporal/worker/tests/test_wiring.py @@ -32,10 +32,11 @@ def test_create_activities_returns_list(): result = create_activities() assert isinstance(result, list) - assert len(result) == 9 + 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 @pytest.mark.unit @@ -54,9 +55,10 @@ def test_get_workflows_returns_classes(): result = get_workflows() assert isinstance(result, list) - assert len(result) == 5 + 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/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py index a51c34d850..6e76f83c54 100644 --- a/backend/airweave/domains/temporal/worker/wiring.py +++ b/backend/airweave/domains/temporal/worker/wiring.py @@ -40,6 +40,9 @@ def create_activities() -> list: 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, ) @@ -61,6 +64,7 @@ def create_activities() -> list: logger.debug("Wiring activities with container dependencies") return [ + DiscoverNativeProjectionActivity().run, CleanupProjectionGenerationsActivity(container.storage_backend).run, ProjectCanonicalRecordsActivity( projector=CanonicalProjector( @@ -129,8 +133,12 @@ def get_workflows() -> list: 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, 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 28b391ec0c..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,13 +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( @@ -32,16 +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, - ) - - if workflow.patched("canonical-generation-gc-v1"): + """Preserve cleanup failures while independently dispatching native recovery.""" + try: await workflow.execute_activity( - CleanupProjectionGenerationsActivity.run, + 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 index db6a689f7d..84371d9d09 100644 --- a/backend/airweave/domains/temporal/workflows/project_canonical_records.py +++ b/backend/airweave/domains/temporal/workflows/project_canonical_records.py @@ -24,12 +24,17 @@ async def run( 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], + 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), ) @@ -38,7 +43,7 @@ async def run( if not result["has_more"]: if not failures: return - if sweep >= 2: + if skip_failed or sweep >= 2: raise ApplicationError( "Canonical projection still has failed records; retained as pending", non_retryable=True, @@ -46,4 +51,7 @@ async def run( 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. - workflow.continue_as_new(args=[organization_id, sync_id, after_id, failures, sweep]) + 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..bc8f00bfa7 --- /dev/null +++ b/backend/airweave/domains/temporal/workflows/recover_native_projection.py @@ -0,0 +1,51 @@ +"""Recover native publications independently of unrelated 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 + 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=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/tests/test_cleanup_workflows.py b/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py index 5c980701d7..66f3447945 100644 --- a/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py +++ b/backend/airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py @@ -13,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, @@ -32,19 +35,26 @@ async def test_cleanup_stuck_sync_jobs_calls_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), cleanup_generations], + 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") 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..2ceeb29753 --- /dev/null +++ b/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py @@ -0,0 +1,165 @@ +"""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"native-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"native-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] From e7fa08f946f92b9c815b511641e671f7e307d1fa Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:17:53 -0700 Subject: [PATCH 209/266] ci: verify owned projection recovery workflows --- .github/workflows/owned-source-store.yml | 2 ++ Makefile | 9 +++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index db5c014ae6..8b91e8f747 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -98,6 +98,8 @@ jobs: run: make test-auth - name: Verify dependency readiness run: make test-health + - name: Verify projection recovery and maintenance workflows + run: make test-worker - name: Record tested Vespa image run: docker image inspect vespaengine/vespa@sha256:5c30f5c41e7563498c4f925db6a837a3848f04726a3ed26aed4a7c8ab69f18fd --format '{{json .RepoDigests}}' - name: Verify real Vespa schema and retrieval diff --git a/Makefile b/Makefile index eb343eea01..6eb0cc620c 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest IMAGE ?= almanac-source-store:local -.PHONY: help setup check test-store test-provisioning test-capture test-search test-auth test-health test-index validate-deploy build +.PHONY: help setup check test-store test-provisioning test-capture test-search test-auth test-health test-worker test-index validate-deploy build help: @echo 'setup Install the locked backend and development dependencies' @@ -14,6 +14,7 @@ help: @echo 'test-search Verify search services and visibility behavior' @echo 'test-auth Verify service-key and source authorization boundaries' @echo 'test-health Verify dependency probes and readiness behavior' + @echo 'test-worker Verify owned projection, recovery and maintenance workflows' @echo 'test-index Deploy/test schemas on explicitly disposable local Vespa' @echo 'check Run store, capture, search, authentication and health checks' @echo 'validate-deploy Validate the Porter application manifest locally (no deployment)' @@ -50,8 +51,12 @@ test-auth: test-health: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/core/health/tests airweave/adapters/health/tests airweave/core/container/tests/test_health_wiring.py +test-worker: + @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py + test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py tests/integration/test_canonical_search_prefilters.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py -check: test-store test-provisioning test-capture test-search test-auth test-health +check: test-store test-provisioning test-capture test-search test-auth test-health test-worker From ff9f22764b7663e13c5ad0259fcdcb4bcae13bd7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:36:17 -0700 Subject: [PATCH 210/266] Bind Outlook managed Graph reads to an attested native principal --- .../auth_provider/providers/composio.py | 2 + backend/airweave/platform/configs/config.py | 14 +- .../platform/sources/outlook_calendar.py | 86 +++--- .../platform/sources/outlook_graph.py | 147 ++++++++++ .../airweave/platform/sources/outlook_mail.py | 98 +++---- .../sources/tests/test_outlook_identity.py | 263 ++++++++++++++++++ 6 files changed, 502 insertions(+), 108 deletions(-) create mode 100644 backend/airweave/platform/sources/outlook_graph.py create mode 100644 backend/airweave/platform/sources/tests/test_outlook_identity.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index 7d8c4434de..d5cb6cf586 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -113,6 +113,8 @@ async def get_auth_result( hosts = { "gmail": frozenset({"gmail.googleapis.com"}), + "outlook_mail": frozenset({"graph.microsoft.com"}), + "outlook_calendar": frozenset({"graph.microsoft.com"}), "google_calendar": frozenset({"www.googleapis.com"}), "google_drive": frozenset( {"www.googleapis.com", "docs.googleapis.com", "sheets.googleapis.com"} diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 34f1114763..27e61d94ff 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -567,12 +567,24 @@ class OracleConfig(SourceConfig): class OutlookCalendarConfig(SourceConfig): """Outlook Calendar configuration schema.""" - pass + 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.""" + 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", diff --git a/backend/airweave/platform/sources/outlook_calendar.py b/backend/airweave/platform/sources/outlook_calendar.py index 9d656b6abd..c7d1703176 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.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.http_helpers import raise_for_status +from airweave.platform.sources.outlook_graph import OutlookBoundaryError, OutlookGraphClient from airweave.schemas.source_connection import AuthenticationMethod, OAuthType @@ -69,34 +71,23 @@ class OutlookCalendarSource(BaseSource): 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 + ) + 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 +95,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 +115,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 +158,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 +189,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 +210,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 +230,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 +278,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 +297,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 +319,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 +392,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 +406,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 +424,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 +447,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 +459,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_graph.py b/backend/airweave/platform/sources/outlook_graph.py new file mode 100644 index 0000000000..9e8189a091 --- /dev/null +++ b/backend/airweave/platform/sources/outlook_graph.py @@ -0,0 +1,147 @@ +"""Outlook's native mailbox boundary over the existing managed/direct transport.""" + +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.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 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) -> httpx.Response: + self._validate_url(url) + headers = await authorization_headers(self.auth) + 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) -> 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) + 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) + if response.status_code == 401: + self.verified_principal_id = None + 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 diff --git a/backend/airweave/platform/sources/outlook_mail.py b/backend/airweave/platform/sources/outlook_mail.py index cf92280568..b820745330 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.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.http_helpers import raise_for_status +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 @@ -75,13 +77,17 @@ class OutlookMailSource(BaseSource): 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 + ) + 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 +97,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 +104,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 +211,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 +240,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 +255,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 +287,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 +364,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 +377,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 +454,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 +475,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 +543,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 +649,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 +670,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 +694,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 +718,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 +751,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 +796,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 +816,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 +846,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 +855,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 +988,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 +1007,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 +1061,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 +1117,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 +1161,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 +1173,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/tests/test_outlook_identity.py b/backend/airweave/platform/sources/tests/test_outlook_identity.py new file mode 100644 index 0000000000..7b3316e001 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_outlook_identity.py @@ -0,0 +1,263 @@ +"""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) From 0d4470566480898e474a13b12ebe15797529ed93 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 04:40:47 -0700 Subject: [PATCH 211/266] Keep raw HTTP transport traces out of service debug logs --- Makefile | 2 +- backend/airweave/core/logging.py | 6 +++ .../http_client/test_logging_privacy.py | 42 +++++++++++++++++++ 3 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 backend/tests/unit/platform/http_client/test_logging_privacy.py diff --git a/Makefile b/Makefile index 6eb0cc620c..5df29be23a 100644 --- a/Makefile +++ b/Makefile @@ -36,7 +36,7 @@ test-store: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/entities/canonical/tests airweave/domains/native_ingestion/tests airweave/domains/syncs/tests test-capture: - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities tests/unit/platform/http_client/test_composio_transport.py airweave/domains/storage/tests/test_file_service.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/platform/sources/tests tests/unit/platform/sources/records tests/unit/platform/sources/test_*_capture.py tests/unit/domains/entities tests/unit/platform/http_client/test_composio_transport.py tests/unit/platform/http_client/test_logging_privacy.py airweave/domains/storage/tests/test_file_service.py test-provisioning: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) diff --git a/backend/airweave/core/logging.py b/backend/airweave/core/logging.py index 4387ed0826..03dbe9c94f 100644 --- a/backend/airweave/core/logging.py +++ b/backend/airweave/core/logging.py @@ -331,5 +331,11 @@ def configure_logger( return ContextualLogger(logger, prefix, dimensions) +# HTTPX INFO includes full URLs; HTTPCore DEBUG includes response headers/cookies. +# Application DEBUG must not implicitly enable those raw credential-bearing traces. +# Sanitized source diagnostics remain owned by AirweaveHttpClient. +logging.getLogger("httpx").setLevel(logging.WARNING) +logging.getLogger("httpcore").setLevel(logging.WARNING) + # Default logger instance logger = LoggerConfigurator.configure_logger(__name__) diff --git a/backend/tests/unit/platform/http_client/test_logging_privacy.py b/backend/tests/unit/platform/http_client/test_logging_privacy.py new file mode 100644 index 0000000000..7d0a6867b5 --- /dev/null +++ b/backend/tests/unit/platform/http_client/test_logging_privacy.py @@ -0,0 +1,42 @@ +"""Verbose application logging must not enable raw HTTP transport traces.""" + +import asyncio +import logging + +import httpx +import pytest + +from airweave.core.logging import logger # noqa: F401 - service logging initialization + + +@pytest.mark.asyncio +async def test_verbose_application_logging_excludes_url_and_response_cookie(caplog): + """Exercise actual HTTPX/HTTPCore logging over a disposable loopback connection.""" + caplog.set_level(logging.DEBUG) + + async def respond(reader, writer): + try: + await reader.readuntil(b"\r\n\r\n") + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n" + b"Set-Cookie: session=private-cookie-sentinel\r\n" + b"Connection: close\r\n\r\nOK" + ) + await writer.drain() + finally: + writer.close() + await writer.wait_closed() + + async with await asyncio.start_server(respond, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + async with httpx.AsyncClient(trust_env=False) as client: + response = await client.get( + f"http://127.0.0.1:{port}/private-path-sentinel?token=private-query-sentinel" + ) + assert response.text == "OK" + + logging.getLogger("httpx").warning("Transport warning remains observable") + assert "Transport warning remains observable" in caplog.text + assert "private-path-sentinel" not in caplog.text + assert "private-query-sentinel" not in caplog.text + assert "private-cookie-sentinel" not in caplog.text From 1822eeee80bc716e1dd0439d5919ff65b0fb1dfe Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 05:02:05 -0700 Subject: [PATCH 212/266] Retain resumable mailbox-owned Outlook originals and project MIME offline --- .../entities/canonical/outlook_projection.py | 310 +++++++++++++++++ .../entities/canonical/projection_mappers.py | 5 + .../tests/test_outlook_projection.py | 205 +++++++++++ .../canonical/tests/test_outlook_recovery.py | 149 ++++++++ .../platform/entities/outlook_mail.py | 19 +- .../platform/sources/outlook_graph.py | 60 +++- .../airweave/platform/sources/outlook_mail.py | 4 + .../platform/sources/outlook_mail_capture.py | 318 ++++++++++++++++++ .../platform/sources/outlook_mail_models.py | 92 +++++ .../tests/test_outlook_mail_capture.py | 239 +++++++++++++ 10 files changed, 1394 insertions(+), 7 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/outlook_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py create mode 100644 backend/airweave/platform/sources/outlook_mail_capture.py create mode 100644 backend/airweave/platform/sources/outlook_mail_models.py create mode 100644 backend/airweave/platform/sources/tests/test_outlook_mail_capture.py 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..03f418909e --- /dev/null +++ b/backend/airweave/domains/entities/canonical/outlook_projection.py @@ -0,0 +1,310 @@ +"""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 ExtractionPart +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] = [] + + 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: + text = _decoded(part).decode(part.get_content_charset() or "us-ascii", errors="strict") + 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", + ), + 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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 80d1167d96..18649a42d7 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -307,6 +307,11 @@ async def map_record( yield await map_gmail(record, storage, directory) return + elif source_name == "outlook_mail": + from airweave.domains.entities.canonical.outlook_projection import map_outlook + + yield await map_outlook(record, storage, directory) + return elif source_name == "slack" and record.identity.record_type == "file": from airweave.domains.entities.canonical.slack_projection import map_slack_file diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py new file mode 100644 index 0000000000..6a21e53cb2 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py @@ -0,0 +1,205 @@ +"""Retained MIME replay: no snippet substitution, bytes loss, or invented attachment identity.""" + +import hashlib +from datetime import datetime, timezone +from email.message import EmailMessage +from pathlib import Path +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.outlook_projection import map_outlook +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity + + +def source(content: bytes): + sync_id = uuid4() + digest = hashlib.sha256(content).hexdigest() + record = SourceRecord( + id=uuid4(), + sync_id=sync_id, + identity=RecordIdentity(record_type="message", native_id="immutable-1"), + revision=1, + payload={ + "id": "immutable-1", + "changeKey": "change-1", + "parentFolderId": "folder-1", + "subject": "Meeting", + "conversationId": "conversation-1", + "bodyPreview": "lossy-preview", + "sender": {"emailAddress": {"name": "रोहन", "address": "assistant@example.com"}}, + "from": {"emailAddress": {"address": "author@example.com"}}, + "toRecipients": [ + {"emailAddress": {"name": "Doe, Jane", "address": "jane@example.com"}} + ], + "bccRecipients": [{"emailAddress": {"address": "bcc@example.com"}}], + "replyTo": [{"emailAddress": {"address": "reply@example.com"}}], + "receivedDateTime": "2026-10-01T01:02:03Z", + "nativeUnknown": {"retained": True}, + }, + 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=( + BlobReference( + key=f"canonical/{sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type="message/rfc822", + ), + ), + indexed_revision=None, + indexed_pipeline_version=None, + ) + storage = AsyncMock() + storage.read_file.return_value = content + return record, storage + + +@pytest.mark.asyncio +async def test_mixed_related_alternative_preserves_unicode_metadata_and_attachments(tmp_path): + message = EmailMessage() + message.set_content("duplicate plain body") + message.add_alternative("

नमस्ते café, पूरा पाठ

", subtype="html", charset="utf-8") + message.get_payload()[1].add_related( + b"image-original", maintype="image", subtype="png", cid="", filename="../../same.png" + ) + message.add_attachment( + b"pdf-original", maintype="application", subtype="pdf", filename="../../same.png" + ) + record, storage = source(message.as_bytes()) + before = record.model_dump() + result = await map_outlook(record, storage, tmp_path) + body, image, attachment = result.entities + assert "नमस्ते café, पूरा पाठ" in Path(body.local_path).read_text() + assert "duplicate plain body" not in Path(body.local_path).read_text() + assert "lossy-preview" not in Path(body.local_path).read_text() + assert body.sender == '"रोहन" ' + assert body.from_address == "author@example.com" + assert body.to_recipients == ['"Doe, Jane" '] + assert body.bcc_recipients == ["bcc@example.com"] and body.reply_to == ["reply@example.com"] + assert body.conversation_id == "conversation-1" and body.folder_id == "folder-1" + assert body.folder_name is None + assert image.is_inline and image.content_id == "" + assert attachment.file_type == "pdf" and image.file_type == "png" + assert image.composite_id != attachment.composite_id + assert [part.part.key for part in result.parts[1:-1]] == [ + "/mime/parts/0/parts/1/parts/1", + "/mime/parts/1", + ] + assert {Path(entity.local_path).read_bytes() for entity in result.entities[1:]} == { + b"image-original", + b"pdf-original", + } + assert all(Path(entity.local_path).parent == tmp_path for entity in result.entities) + assert result.parts[-1].part.key == "/attachment_inventory" and result.parts[-1].entity is None + assert record.model_dump() == before + storage.read_file.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_latin_charset_and_explicit_related_root_are_respected(tmp_path): + message = EmailMessage() + message.set_content("café ", charset="iso-8859-1") + record, storage = source(message.as_bytes()) + record.payload["subject"] = " " + result = await map_outlook(record, storage, tmp_path) + assert result.entities[0].subject == "Email message" + assert "café <full text>" in Path(result.entities[0].local_path).read_text() + message = EmailMessage() + message.make_related() + message.set_param("start", "") + image = EmailMessage() + image.set_content(b"image", maintype="image", subtype="png") + image["Content-ID"] = "" + body = EmailMessage() + body.set_content("selected body") + body["Content-ID"] = "" + message.attach(image) + message.attach(body) + record, storage = source(message.as_bytes()) + result = await map_outlook(record, storage, tmp_path) + assert "selected body" in Path(result.entities[0].local_path).read_text() + assert result.entities[1].content_id == "" + + +@pytest.mark.asyncio +async def test_missing_and_unsupported_originals_remain_explicit(tmp_path): + record, storage = source(b"Subject: empty\r\n\r\n") + omitted = record.model_copy(update={"blobs": (), "completeness": "metadata_only"}) + result = await map_outlook(omitted, storage, tmp_path) + assert not result.entities and result.parts[0].entity is None + storage.read_file.assert_not_called() + with pytest.raises(ValueError, match="metadata-only"): + await map_outlook(record.model_copy(update={"blobs": ()}), storage, tmp_path) + message = EmailMessage() + message.set_content("retained body") + nested = EmailMessage() + nested.set_content("nested mail is not silently indexed") + message.add_attachment(nested) + record, storage = source(message.as_bytes()) + result = await map_outlook(record, storage, tmp_path) + assert len(result.entities) == 1 + assert result.parts[1].omission == "unsupported_format" + assert result.parts[1].part.media_type == "message/rfc822" + encrypted = b"MIME-Version: 1.0\r\nContent-Type: application/pkcs7-mime\r\n\r\nencrypted" + record, storage = source(encrypted) + result = await map_outlook(record, storage, tmp_path) + assert not result.entities and result.parts[0].omission == "unsupported_format" + + +@pytest.mark.asyncio +async def test_identity_and_blob_integrity_fail_before_materialization(tmp_path): + record, storage = source(b"Subject: retained\r\n\r\nbody") + bad = record.model_copy( + update={"identity": RecordIdentity(record_type="message", native_id="different")} + ) + with pytest.raises(ValueError, match="identity"): + await map_outlook(bad, storage, tmp_path) + storage.read_file.assert_not_called() + bad_ref = record.blobs[0].model_copy(update={"key": "canonical/other/blobs/sha256/hash"}) + with pytest.raises(ValueError, match="source scope"): + await map_outlook(record.model_copy(update={"blobs": (bad_ref,)}), storage, tmp_path) + storage.read_file.assert_not_called() + storage.read_file.return_value = b"bad bytes" + with pytest.raises(ValueError, match="digest"): + await map_outlook(record, storage, tmp_path) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_malformed_transfers_charset_and_structure_are_never_repaired_silently(tmp_path): + invalid_messages = [ + b"Content-Transfer-Encoding: base64\r\n\r\naGVsbG8=@@", + b"Content-Transfer-Encoding: quoted-printable\r\n\r\nhello=ZZ", + b"Content-Type: text/plain; charset=utf-8\r\n\r\n\xff", + b"Content-Type: multipart/mixed; boundary=missing\r\n\r\nno boundary", + ] + for content in invalid_messages: + record, storage = source(content) + with pytest.raises(ValueError): + await map_outlook(record, storage, tmp_path) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.asyncio +async def test_pathological_nesting_is_explicit_failure_not_partial_text(tmp_path): + message = EmailMessage() + message.set_content("body") + for _ in range(34): + parent = EmailMessage() + parent.make_mixed() + parent.attach(message) + message = parent + record, storage = source(message.as_bytes()) + with pytest.raises(ValueError, match="depth or part limit"): + await map_outlook(record, storage, tmp_path) + assert not list(tmp_path.iterdir()) diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py new file mode 100644 index 0000000000..760584bf34 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py @@ -0,0 +1,149 @@ +"""Outlook native HTTP fixtures through real canonical SQL and retained blob storage.""" + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.adapters.storage.filesystem import FilesystemBackend +from airweave.domains.entities.canonical.blob_materializer import read_blob +from airweave.domains.entities.canonical.cycle_models import CompleteCycle +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.sources.token_providers.static import StaticTokenProvider +from airweave.domains.storage.file_service import FileService +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 OutlookMailConfig +from airweave.platform.sources.outlook_graph import OutlookGraphClient +from airweave.platform.sources.outlook_mail_capture import OutlookMailCapture + +BASE = "https://graph.microsoft.com/v1.0/me" +MIME = ( + b"Subject: Original\r\nContent-Type: text/plain; charset=utf-8\r\n\r\n" + b"Original retained bytes\r\n" +) + + +class Mailbox: + """One interrupted list page, followed by a move of the same immutable message.""" + + def __init__(self): + self.fail_next_page = True + self.calls = [] + self.moved = False + + async def handle(self, request): + path = request.url.path + if path == "/v1.0/me": + return httpx.Response(200, json={"id": "mailbox"}) + assert request.headers["Prefer"] == 'IdType="ImmutableId"' + self.calls.append((path, str(request.url.query))) + if path == "/v1.0/me/messages": + if request.url.params.get("$skiptoken") == "second": + if self.fail_next_page: + raise httpx.ConnectError("Synthetic page interruption", request=request) + self.moved = True + return httpx.Response(200, json={"value": [{"id": "a"}, {"id": "b"}]}) + return httpx.Response( + 200, + json={ + "value": [{"id": "a"}], + "@odata.nextLink": BASE + "/messages?$skiptoken=second", + }, + ) + if path.endswith("/$value"): + return httpx.Response(200, content=MIME, headers={"Content-Type": "message/rfc822"}) + identity = path.rsplit("/", 1)[-1] + assert identity in {"a", "b"} + return httpx.Response( + 200, + json={ + "id": identity, + "changeKey": "moved" if self.moved and identity == "a" else "original", + "parentFolderId": "sent" if self.moved and identity == "a" else "inbox", + "subject": "Original", + "body": {"contentType": "text", "content": "Original retained bytes"}, + "createdDateTime": "2026-09-01T00:00:00Z", + "lastModifiedDateTime": "2026-09-02T00:00:00Z", + "unknownNativeField": {"retained": True}, + }, + ) + + +async def run_capture(database, service, fence, mailbox, storage): + async with httpx.AsyncClient(transport=httpx.MockTransport(mailbox.handle)) as client: + graph = OutlookGraphClient( + StaticTokenProvider("synthetic"), client, "outlook_mail", "mailbox" + ) + connector = await OutlookMailCapture.create( + graph=graph, + config=OutlookMailConfig( + expected_principal_id="mailbox", included_folders=[], excluded_folders=[] + ), + ) + files = FileService(fence.job_id, storage, sync_id=fence.sync_id) + try: + return await CanonicalScanDriver( + service, database, fence, connector, AsyncMock(), AsyncMock(), files + ).run() + finally: + await files.cleanup_sync_directory(MagicMock()) + + +async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place( + database, source, tmp_path +): + service, fence = source + mailbox = Mailbox() + storage = FilesystemBackend(tmp_path) + with pytest.raises(httpx.RequestError): + await run_capture(database, service, fence, mailbox, storage) + async with database() as db: + rows = list((await db.scalars(select(Entity))).all()) + assert len(rows) == 1 and rows[0].native_id == "a" and rows[0].deleted_at is None + initial_id = rows[0].id + initial_revision = rows[0].record_revision + scan = await db.scalar(select(CaptureScan)) + assert scan.phase == "collecting" + before = await service.read_cycle(db, fence) + renewed = await service.activate_writer( + db, + fence.organization_id, + fence.sync_id, + fence.job_id, + attempt_id=uuid4(), + attempt_number=2, + ) + mailbox.fail_next_page = False + split = len(mailbox.calls) + after = await run_capture(database, service, renewed, mailbox, storage) + async with database() as db: + after = await service.complete_cycle( + db, CompleteCycle(fence=renewed, expected=after.version) + ) + assert after.phase == "complete" + assert after.last_full_capture.discovery == "incomplete" + assert after.promoted_checkpoint is None + assert after.version.cycle_id == before.version.cycle_id + listing = [query for path, query in mailbox.calls[split:] if path == "/v1.0/me/messages"] + assert len(listing) == 1 and "second" in listing[0] + async with database() as db: + rows = list((await db.scalars(select(Entity))).all()) + assert len(rows) == 2 and all(row.deleted_at is None for row in rows) + moved = next(row for row in rows if row.native_id == "a") + assert moved.id == initial_id and moved.record_revision == initial_revision + 1 + assert moved.source_payload["parentFolderId"] == "sent" + assert moved.source_payload["unknownNativeField"] == {"retained": True} + retained = await service.store.read(db, fence.organization_id, fence.sync_id, moved.id) + assert retained.parent is None and retained.identity.container_id is None + assert retained.completeness == "partial" and len(retained.blobs) == 1 + assert await read_blob(retained, retained.blobs[0], storage) == MIME + # Actual mapper dispatch consumes retained bytes after the provider client has closed. + async with map_record(retained, "outlook_mail", storage) as projected: + assert len(projected.entities) == 1 + assert projected.entities[0].id == "a" + assert projected.parts[-1].part.key == "/attachment_inventory" + assert projected.parts[-1].entity is None 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/sources/outlook_graph.py b/backend/airweave/platform/sources/outlook_graph.py index 9e8189a091..6f9931ba40 100644 --- a/backend/airweave/platform/sources/outlook_graph.py +++ b/backend/airweave/platform/sources/outlook_graph.py @@ -1,5 +1,6 @@ """Outlook's native mailbox boundary over the existing managed/direct transport.""" +from contextlib import aclosing from urllib.parse import unquote, urlsplit import httpx @@ -11,6 +12,7 @@ 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 @@ -71,9 +73,13 @@ def _validate_url(self, url: str) -> None: 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) -> httpx.Response: + 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) + if immutable_ids: + headers["Prefer"] = 'IdType="ImmutableId"' try: return await self.http_client.get( url, headers=headers, params=params, follow_redirects=False @@ -120,19 +126,21 @@ async def verify_principal(self) -> None: ) self.verified_principal_id = principal.id - async def get(self, url: str, params: dict | None = None) -> dict: + 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) + 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) + response = await self._request(url, params, immutable_ids=immutable_ids) if response.status_code == 401: self.verified_principal_id = None self._check_response(response) @@ -145,3 +153,47 @@ async def get(self, url: str, params: dict | None = None) -> dict: 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 b820745330..09d1d10e0f 100644 --- a/backend/airweave/platform/sources/outlook_mail.py +++ b/backend/airweave/platform/sources/outlook_mail.py @@ -73,6 +73,10 @@ 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, 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..dc800985bc --- /dev/null +++ b/backend/airweave/platform/sources/outlook_mail_capture.py @@ -0,0 +1,318 @@ +"""Opt-in mailbox-owned full Outlook capture; no legacy cursor or activation changes.""" + +import hashlib +import json +from datetime import datetime, timezone +from urllib.parse import quote, 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 +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, OutlookGraphClient +from airweave.platform.sources.outlook_mail_models import ( + OutlookMailContinuation, + OutlookMessage, + OutlookMessageID, + OutlookMessagePage, +) + +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. Incremental folder deltas remain a required, unimplemented follow-up. + """ + + 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": 1, + "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: + """Every cycle is explicitly full; no delta completion is implied.""" + await self.graph.verify_principal() + self._require_principal() + return CapturePlan() + + def initial_continuation(self, cycle: CaptureCycle) -> ScanContinuation: + """Bind a fresh scan to its persisted principal, selection and ID format.""" + self._require_principal() + if cycle.configuration != self.capture_cycle_configuration: + raise OutlookBoundaryError( + "Outlook capture configuration changed", source_short_name="outlook_mail" + ) + return self._continuation(OutlookMailContinuation(fingerprint=self.fingerprint)) + + @staticmethod + def _continuation(state: OutlookMailContinuation) -> ScanContinuation: + return ScanContinuation(value=state.model_dump(mode="json")) + + def _state(self, continuation: ScanContinuation) -> OutlookMailContinuation: + try: + state = OutlookMailContinuation.model_validate(continuation.value) + except ValidationError: + raise OutlookBoundaryError( + "Invalid Outlook capture continuation", source_short_name="outlook_mail" + ) from None + if state.fingerprint != self.fingerprint: + raise OutlookBoundaryError( + "Outlook continuation belongs to another capture scope", + source_short_name="outlook_mail", + ) + return state + + def _list_url(self, url: str) -> None: + self.graph._validate_url(url) + if urlsplit(url).path != "/v1.0/me/messages": + raise OutlookBoundaryError( + "Outlook continuation changed the message collection", + source_short_name="outlook_mail", + ) + + async def capture_page( + self, + scope: CompletedScope, + continuation: ScanContinuation, + *, + files: FileService, + parent: SourceRecord | None = None, + ) -> CapturePage: + """Capture one complete message; durable pending IDs avoid refetching inventory.""" + self._require_principal() + if scope != CompletedScope(record_type="message") or parent is not None: + raise OutlookBoundaryError( + "Outlook requires the mailbox message root", source_short_name="outlook_mail" + ) + state = self._state(continuation) + if not state.pending_ids: + if state.started and state.next_link is None: + return CapturePage(records=(), continuation=continuation, final=True) + url = state.next_link or f"{BASE}/messages" + self._list_url(url) + raw = await self.graph.get( + url, + params=None if state.started else {"$select": "id", "$top": 25}, + immutable_ids=True, + ) + try: + page = OutlookMessagePage.model_validate(raw) + except ValidationError: + raise OutlookBoundaryError( + "Invalid Outlook message inventory", source_short_name="outlook_mail" + ) from None + if page.next_link: + self._list_url(page.next_link) + if page.next_link == url: + raise OutlookBoundaryError( + "Outlook continuation did not advance", source_short_name="outlook_mail" + ) + state = OutlookMailContinuation( + fingerprint=self.fingerprint, + started=True, + pending_ids=tuple(dict.fromkeys(item.id for item in page.value)), + next_link=page.next_link, + ) + self._continuation(state) # Validate the byte bound before hydrating any originals. + records = () + if state.pending_ids: + observation = await self.message(state.pending_ids[0], files=files) + # New out-of-selection objects are not stored as synthetic tombstones. + # Existing objects are checked by the shared exact omission phase. + records = (observation,) if observation.kind == "upsert" else () + state = state.model_copy(update={"pending_ids": state.pending_ids[1:]}) + return CapturePage( + records=records, + continuation=self._continuation(state), + final=not state.pending_ids and state.next_link is None, + ) + + @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..b5f34fa475 --- /dev/null +++ b/backend/airweave/platform/sources/outlook_mail_models.py @@ -0,0 +1,92 @@ +"""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 OutlookMailContinuation(BaseModel): + """Only pending native IDs and one native next link; no mailbox-sized folder map.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + version: Literal[1] = 1 + fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") + started: bool = False + pending_ids: tuple[str, ...] = Field(default=(), max_length=500) + next_link: str | None = None + + @model_validator(mode="after") + def coherent_progress(self): + """An initial cursor cannot masquerade as an acknowledged native page.""" + if len(set(self.pending_ids)) != len(self.pending_ids) or any( + not x for x in self.pending_ids + ): + raise ValueError("Pending Outlook identities must be nonempty and unique") + if not self.started and (self.pending_ids or self.next_link is not None): + raise ValueError("Initial Outlook continuation cannot contain progress") + return self 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..bc4b8796e8 --- /dev/null +++ b/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py @@ -0,0 +1,239 @@ +"""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") + ) + + +@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"' + if request.url.path == "/v1.0/me/messages": + return httpx.Response( + 200, json={"value": [{"id": "immutable-message"}, {"id": "second"}]} + ) + 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 source.capture_page(ROOT, initial(source), files=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/messages") == 1 + + +@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/messages?$skip=next" + + def graph(request): + requested.append(request.url.path) + if request.url.path == "/v1.0/me": + return httpx.Response(200, json={"id": "principal"}) + 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 message collection"): + await source.capture_page(ROOT, initial(source), files=files()) + next_link = "https://graph.microsoft.com/v1.0/me/messages?token=" + "x" * 66000 + with pytest.raises(ValueError, match="64 KiB"): + await source.capture_page(ROOT, initial(source), files=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()) From a37f5ddaf319d38a12dc9a52a99c5272e908b065 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 05:20:24 -0700 Subject: [PATCH 213/266] Wire owned Outlook originals into source lifecycle --- .../canonical/tests/test_outlook_recovery.py | 51 +++++++++++++ .../domains/oauth/callback_service.py | 1 + .../oauth/tests/test_callback_service.py | 14 +++- .../domains/owned_provisioning/models.py | 14 +++- .../tests/test_provisioning.py | 62 ++++++++++++++- .../domains/source_connections/create.py | 2 + .../source_connections/tests/test_update.py | 75 ++++++++++++++++++- .../domains/source_connections/update.py | 24 ++++++ .../airweave/domains/sync_pipeline/factory.py | 63 ++++++++++++++-- .../sync_pipeline/tests/test_factory.py | 28 +++++-- .../airweave/domains/syncs/fakes/service.py | 3 + backend/airweave/domains/syncs/protocols.py | 2 + backend/airweave/domains/syncs/service.py | 16 ++-- .../domains/syncs/tests/test_service.py | 26 +++++-- backend/airweave/platform/configs/config.py | 15 ++++ backend/airweave/platform/sources/_base.py | 12 +++ .../airweave/platform/sources/outlook_mail.py | 9 ++- .../sources/tests/test_outlook_identity.py | 29 +++++++ 18 files changed, 415 insertions(+), 31 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py index 760584bf34..a77988a1c3 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py @@ -147,3 +147,54 @@ async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place assert projected.entities[0].id == "a" assert projected.parts[-1].part.key == "/attachment_inventory" assert projected.parts[-1].entity is None + + +async def test_factory_rejects_mode_conflict_before_auth_or_cursor_reset(database, source): + """A full sync or skip-load cannot reinterpret a persisted Outlook cursor.""" + from airweave.domains.sync_pipeline.config import SyncConfig + from airweave.domains.sync_pipeline.factory import SyncFactory + from airweave.models.sync_cursor import SyncCursor + from airweave.platform.sources.outlook_mail import OutlookMailSource + + _, fence = source + factory = object.__new__(SyncFactory) + factory._source_lifecycle_service = MagicMock(create=AsyncMock()) + factory._source_registry = MagicMock() + factory._source_registry.get.return_value.source_class_ref = OutlookMailSource + factory._build_arf_replay_source = AsyncMock(return_value="offline legacy replay") + ctx = MagicMock() + ctx.organization.id = fence.organization_id + sc = MagicMock(short_name="outlook_mail", config_fields={}) + sync = MagicMock(id=fence.sync_id) + job = MagicMock(id=fence.job_id) + replay = SyncConfig.model_validate({"behavior": {"replay_from_arf": True}}) + async with database() as db: + cursor = await db.scalar(select(SyncCursor).where(SyncCursor.sync_id == fence.sync_id)) + if cursor is None: + cursor = SyncCursor( + sync_id=fence.sync_id, organization_id=fence.organization_id, cursor_data={} + ) + db.add(cursor) + cursor.cursor_data = {} + await db.flush() + assert await factory._build_source( + db, sync, job, ctx, MagicMock(), sc, False, replay + ) == "offline legacy replay" + factory._build_arf_replay_source.reset_mock() + sc.config_fields = {"capture_originals": True, "expected_principal_id": "mailbox"} + with pytest.raises(ValueError, match="not mutable ARF"): + await factory._build_source(db, sync, job, ctx, MagicMock(), sc, False, replay) + cursor.cursor_data = {"delta_link": "legacy continuation"} + await db.flush() + for force_full_sync in (False, True): + with pytest.raises(ValueError, match="retained state"): + await factory._build_source( + db, sync, job, ctx, MagicMock(), sc, force_full_sync, replay + ) + sc.config_fields = {} + cursor.cursor_data = {"canonical_cycle": {"phase": "active"}} + await db.flush() + with pytest.raises(ValueError, match="retained state"): + await factory._build_source(db, sync, job, ctx, MagicMock(), sc, True, replay) + factory._source_lifecycle_service.create.assert_not_called() + factory._build_arf_replay_source.assert_not_called() 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/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index b8509ae064..e79430178d 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -11,7 +11,7 @@ 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"] + provider: Literal["gmail", "google_calendar", "google_drive", "slack", "outlook_mail"] 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) @@ -39,6 +39,12 @@ def valid_schedule(cls, value: str) -> str: def source_config(self) -> dict[str, JsonValue]: """Expected identity cannot be overridden inside unstructured provider config.""" + if self.provider == "outlook_mail": + return { + **self.config, + "expected_principal_id": self.expected_identity, + "capture_originals": True, + } if self.provider == "slack": return { **self.config, @@ -100,6 +106,7 @@ def native_principal(provider: str, config: dict) -> tuple[str, str | None]: GmailConfig, GoogleCalendarConfig, GoogleDriveConfig, + OutlookMailConfig, SlackConfig, ) @@ -111,6 +118,11 @@ def native_principal(provider: str, config: dict) -> tuple[str, str | None]: identity = GoogleCalendarConfig.model_validate(config).expected_primary_calendar_id case "google_drive": identity = GoogleDriveConfig.model_validate(config).expected_permission_id + case "outlook_mail": + outlook = OutlookMailConfig.model_validate(config) + if not outlook.capture_originals: + raise ValueError("Owned Outlook source must capture originals") + identity = outlook.expected_principal_id case "slack": slack = SlackConfig.model_validate(config) identity, user = slack.expected_team_id, slack.expected_user_id diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index 2d3b233aac..7290828b17 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -13,7 +13,7 @@ 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 +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 @@ -453,3 +453,63 @@ async def test_slack_workspace_user_survives_pause_and_rejects_different_member( ) assert resumed.source_connection_id == first.source_connection_id assert (resumed.expected_identity, resumed.expected_user_identity) == ("T1", "U1") + + +async def test_outlook_fresh_source_retry_and_principal_cannot_downgrade_owned_capture( + database, setup +): + ctx, service, request, account, lifecycle, schedules, workflows = setup + creator = service.store.create + creator._source_registry.get.return_value.short_name = "outlook_mail" + expected_config = {"expected_principal_id": "native-owner", "capture_originals": True} + creator._source_validation.seed_config_result("outlook_mail", expected_config) + spec = ManagedSource.model_validate( + { + **request.source.model_dump(), + "provider": "outlook_mail", + "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("outlook_mail", spec.source_config()) == ("native-owner", None) + with pytest.raises(ValueError, match="must capture originals"): + native_principal("outlook_mail", {"expected_principal_id": "native-owner"}) + with pytest.raises(ValueError): + native_principal("outlook_mail", {"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 == "outlook_mail" 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() diff --git a/backend/airweave/domains/source_connections/create.py b/backend/airweave/domains/source_connections/create.py index e6ca597b47..cacfbc9c90 100644 --- a/backend/airweave/domains/source_connections/create.py +++ b/backend/airweave/domains/source_connections/create.py @@ -488,6 +488,7 @@ 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, @@ -702,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/tests/test_update.py b/backend/airweave/domains/source_connections/tests/test_update.py index 5af8384332..75fa6515cd 100644 --- a/backend/airweave/domains/source_connections/tests/test_update.py +++ b/backend/airweave/domains/source_connections/tests/test_update.py @@ -201,6 +201,77 @@ async def test_update_config_valid(): assert any(c[0] == "validate_config" for c in validation._calls) +@pytest.mark.parametrize("originals", [False, True]) +async def test_outlook_capture_mode_change_requires_new_source(originals: bool): + sc = _make_sc(short_name="outlook_mail") + 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("outlook_mail", 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("originals", [False, True]) +async def test_outlook_principal_change_preserves_capture_identity(originals: bool): + sc = _make_sc(short_name="outlook_mail") + 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("outlook_mail", 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 # --------------------------------------------------------------------------- @@ -665,7 +736,9 @@ async def test_native_update_cannot_rebind_identity(): 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"}), + db, + id=sc.id, + obj_in=SourceConnectionUpdate(config={"owner_id": "other"}), ctx=_make_ctx(), ) assert error.value.status_code == 409 diff --git a/backend/airweave/domains/source_connections/update.py b/backend/airweave/domains/source_connections/update.py index f26ce795be..c9db0bba1f 100644 --- a/backend/airweave/domains/source_connections/update.py +++ b/backend/airweave/domains/source_connections/update.py @@ -31,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 OutlookMailConfig from airweave.schemas.source_connection import ( AuthenticationMethod, ScheduleConfig, @@ -41,6 +42,27 @@ ) +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 != "outlook_mail": + return + previous = OutlookMailConfig.model_validate(source.config_fields or {}) + proposed = OutlookMailConfig.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", + ) + + class SourceConnectionUpdateService(SourceConnectionUpdateServiceProtocol): """Updates a source connection. @@ -115,6 +137,7 @@ 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) update_data["config_fields"] = validated_config del update_data["config"] @@ -238,6 +261,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/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index 2ce2a3307a..437a86b5d9 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 @@ -67,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 from airweave.platform.sources._base import BaseSource from .entity.pipeline import EntityPipeline @@ -199,11 +204,9 @@ async def create_orchestrator( 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) - canonical_source = ( - source_result.source - if isinstance(source_result.source, (CanonicalSource, CanonicalPageSource)) - else None - ) + 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 @@ -358,7 +361,7 @@ def _build_stream( sync_context: SyncContext, ) -> AsyncSourceStream | None: """Page sources are driven sequentially; legacy generators retain their stream.""" - if isinstance(runtime.source, CanonicalPageSource): + 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 @@ -397,9 +400,15 @@ async def _build_source( access_token: Optional[str] = None, ) -> SourceBuildResult: """Build source instance, cursor, file service, and node selections.""" + outlook_capture = False + if source_connection.short_name == "outlook_mail": + outlook_capture = OutlookMailConfig.model_validate( + source_connection.config_fields or {} + ).capture_originals + await self._validate_outlook_capture_mode(db, sync.id, ctx, outlook_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 isinstance(source_class, (CanonicalSource, CanonicalPageSource)): + if outlook_capture or isinstance(source_class, (CanonicalSource, CanonicalPageSource)): raise ValueError( "Canonical records must be reindexed from Postgres, not mutable ARF" ) @@ -438,6 +447,46 @@ async def _build_source( source=source, cursor=cursor, files=files, node_selections=node_selections ) + async def _validate_outlook_capture_mode( + self, db: AsyncSession, sync_id: UUID, ctx: BaseContext, capture_originals: bool + ) -> None: + """Never reinterpret persisted IDs or cursors when selecting an Outlook 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( + "Outlook 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/tests/test_factory.py b/backend/airweave/domains/sync_pipeline/tests/test_factory.py index 2daf27e40a..8be28460c4 100644 --- a/backend/airweave/domains/sync_pipeline/tests/test_factory.py +++ b/backend/airweave/domains/sync_pipeline/tests/test_factory.py @@ -141,7 +141,7 @@ async def test_create_orchestrator_raises_when_source_connection_missing(): @pytest.mark.asyncio -@pytest.mark.parametrize("canonical", [False, True, "page"]) +@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() @@ -191,8 +191,9 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): ) as mock_build_tracker, ): mock_source = MagicMock() + mock_source.capture_page_source = None mock_source.generate_entities = MagicMock(return_value=AsyncMock()) - if canonical == "page": + 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 @@ -205,6 +206,17 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): 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()) @@ -235,9 +247,12 @@ async def test_create_orchestrator_passes_entity_repo_to_pipeline(canonical): assert isinstance(orchestrator.entity_pipeline, CanonicalCapturePipeline) mock_build_destinations.assert_not_called() mock_disp_builder.assert_not_called() - if canonical == "page": + if canonical in {"page", "composed"}: assert orchestrator.stream is None - assert orchestrator.entity_pipeline.page_source is mock_source + 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() @@ -403,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() 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 416c2e8e5e..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 @@ -124,6 +125,7 @@ async def create( 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/service.py b/backend/airweave/domains/syncs/service.py index 2e3a3e577c..c8041f3e71 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 @@ -58,6 +59,7 @@ TemporalScheduleServiceProtocol, TemporalWorkflowServiceProtocol, ) +from airweave.platform.configs.config import OutlookMailConfig from airweave.schemas.source_connection import ScheduleConfig from airweave.schemas.sync import SyncCreate from airweave.schemas.sync_job import SyncJobCreate @@ -111,6 +113,7 @@ async def create( ctx: ApiContext, uow: UnitOfWork, defer_execution: bool = False, + source_config: dict[str, JsonValue] | None = None, ) -> SyncProvisionResult: """Create sync + optional job + Temporal schedule atomically. @@ -129,6 +132,13 @@ 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 == "outlook_mail": + # Outlook composes its canonical adapter only for explicit original capture. + canonical = OutlookMailConfig.model_validate(source_config or {}).capture_originals + sync_schema, sync_job_schema = await self._create_sync_records( uow.session, name=f"Sync for {name}", @@ -136,11 +146,7 @@ async def create( destination_connection_ids=destination_connection_ids, cron_schedule=cron, run_immediately=run_immediately, - initial_pipeline_version=( - SEARCH_METADATA_PIPELINE_VERSION - if isinstance(source_entry.source_class_ref, (CanonicalSource, CanonicalPageSource)) - else 1 - ), + initial_pipeline_version=(SEARCH_METADATA_PIPELINE_VERSION if canonical else 1), ctx=ctx, uow=uow, ) diff --git a/backend/airweave/domains/syncs/tests/test_service.py b/backend/airweave/domains/syncs/tests/test_service.py index 9a27deed49..010af70462 100644 --- a/backend/airweave/domains/syncs/tests/test_service.py +++ b/backend/airweave/domains/syncs/tests/test_service.py @@ -1073,9 +1073,18 @@ async def test_create_no_cron_no_run_immediately(): @pytest.mark.asyncio @pytest.mark.parametrize( - "source_kind, expected_version", [("legacy", 1), ("gmail", 2), ("slack", 2)] + "source_kind, source_config, expected_version", + [ + ("legacy", None, 1), + ("gmail", None, 2), + ("slack", None, 2), + ("outlook_mail", None, 1), + ("outlook_mail", {"capture_originals": True, "expected_principal_id": "native-owner"}, 2), + ], ) -async def test_create_with_cron_calls_temporal_schedule(source_kind, expected_version): +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() @@ -1097,12 +1106,16 @@ async def test_create_with_cron_calls_temporal_schedule(source_kind, expected_ve ) from airweave.platform.sources.gmail import GmailSource + from airweave.platform.sources.outlook_mail import OutlookMailSource from airweave.platform.sources.slack import SlackSource - source_entry = _mock_source_entry() - source_entry.source_class_ref = {"legacy": object, "gmail": GmailSource, "slack": SlackSource}[ - source_kind - ] + source_entry = _mock_source_entry(short_name=source_kind) + source_entry.source_class_ref = { + "legacy": object, + "gmail": GmailSource, + "slack": SlackSource, + "outlook_mail": OutlookMailSource, + }[source_kind] result = await svc.create( AsyncMock(), name="test", @@ -1111,6 +1124,7 @@ async def test_create_with_cron_calls_temporal_schedule(source_kind, expected_ve collection_id=uuid4(), collection_readable_id="col-x", source_entry=source_entry, + source_config=source_config, schedule_config=ScheduleConfig(cron="0 6 * * *"), run_immediately=False, ctx=_mock_ctx(), diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 27e61d94ff..c1997c7093 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -578,6 +578,21 @@ class OutlookCalendarConfig(SourceConfig): 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, 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/outlook_mail.py b/backend/airweave/platform/sources/outlook_mail.py index 09d1d10e0f..e50b565d73 100644 --- a/backend/airweave/platform/sources/outlook_mail.py +++ b/backend/airweave/platform/sources/outlook_mail.py @@ -91,7 +91,14 @@ async def create( instance.graph = OutlookGraphClient( auth, http_client, cls.short_name, config.expected_principal_id ) - await instance.graph.verify_principal() + 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 diff --git a/backend/airweave/platform/sources/tests/test_outlook_identity.py b/backend/airweave/platform/sources/tests/test_outlook_identity.py index 7b3316e001..134bb200fb 100644 --- a/backend/airweave/platform/sources/tests/test_outlook_identity.py +++ b/backend/airweave/platform/sources/tests/test_outlook_identity.py @@ -261,3 +261,32 @@ def graph(request): 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 From f28bc5b2c5c252316df909c9aca6d8bc3ab61935 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 05:25:32 -0700 Subject: [PATCH 214/266] Batch root coverage evidence without changing sync authority --- .../domains/entities/canonical/coverage.py | 33 ++++++++----- .../canonical/tests/test_mixed_scopes.py | 48 ++++++++++++++++++- 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/coverage.py b/backend/airweave/domains/entities/canonical/coverage.py index fe237eacb0..ce9967c6a0 100644 --- a/backend/airweave/domains/entities/canonical/coverage.py +++ b/backend/airweave/domains/entities/canonical/coverage.py @@ -4,7 +4,7 @@ from uuid import UUID from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError, model_validator -from sqlalchemy import and_, case, func, select +from sqlalchemy import and_, case, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession from airweave.domains.entities.canonical.cycle_models import ( @@ -12,7 +12,6 @@ CaptureCycle, CompletionPolicy, ) -from airweave.domains.entities.canonical.cycle_store import root_ready from airweave.domains.entities.canonical.requests import WriterFence from airweave.domains.entities.canonical.store import content_is_available from airweave.models.capture_scan import CaptureScan @@ -77,27 +76,37 @@ async def mixed_scope_summary( attempt_id=sync.writer_attempt_id, attempt_number=sync.writer_attempt_number, ) - if not await root_ready(db, fence, cycle): - return None - eligible = full = changes = 0 common = ( CaptureScan.organization_id == organization_id, CaptureScan.sync_id == sync_id, CaptureScan.cycle_id == cycle.version.cycle_id, CaptureScan.phase == "complete", ) - for kind in cycle.configuration.root_record_types: - eligible += 1 - mode = await db.scalar( - select(CaptureScan.execution_state["plan"]["mode"].astext).where( + roots = cycle.configuration.root_record_types + inventory_roots = tuple(kind for kind in roots if cycle.configuration.children_of(kind)) + modes: dict[str, str | None] = {} + if roots: + rows = await db.execute( + select( + CaptureScan.record_type, + CaptureScan.execution_state["plan"]["mode"].astext, + ).where( *common, - CaptureScan.record_type == kind, + CaptureScan.record_type.in_(roots), CaptureScan.parent_record_id.is_(None), CaptureScan.container_id.is_(None), + or_( + CaptureScan.record_type.not_in(inventory_roots), + CaptureScan.membership_attempt_id == fence.attempt_id, + ), ) ) - full += int(mode == "full") - changes += int(mode == "changes") + modes = dict(rows.tuples().all()) + if not set(roots).issubset(modes): + return None + eligible = len(roots) + full = sum(mode == "full" for mode in modes.values()) + changes = sum(mode == "changes" for mode in modes.values()) for kind, parents in cycle.configuration.parents.items(): parent_types = tuple(parent for parent in parents if parent is not None) if not parent_types: diff --git a/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py index 10e8681536..f048701256 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py +++ b/backend/airweave/domains/entities/canonical/tests/test_mixed_scopes.py @@ -7,7 +7,7 @@ import pytest from sqlalchemy import select -from airweave.domains.entities.canonical.coverage import capture_coverage +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, @@ -541,3 +541,49 @@ async def test_completed_coverage_survives_job_completion_but_not_new_writer(dat 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 From fd08ffb4df299567828d7b83d54161dc57ef9a85 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 05:37:23 -0700 Subject: [PATCH 215/266] Allow isolated worker control and metrics bindings --- Makefile | 2 +- backend/airweave/core/config/settings.py | 1 + .../domains/temporal/worker/__init__.py | 2 +- .../domains/temporal/worker/config.py | 4 +++ .../domains/temporal/worker/control_server.py | 4 +-- .../temporal/worker/tests/test_config.py | 4 +++ .../temporal/worker/tests/test_worker.py | 27 +++++++++++++++++-- deploy/README.md | 9 +++++++ 8 files changed, 47 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 5df29be23a..8a46e7c7b7 100644 --- a/Makefile +++ b/Makefile @@ -53,7 +53,7 @@ test-health: test-worker: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) diff --git a/backend/airweave/core/config/settings.py b/backend/airweave/core/config/settings.py index 080cbe5a2b..9b5a520465 100644 --- a/backend/airweave/core/config/settings.py +++ b/backend/airweave/core/config/settings.py @@ -233,6 +233,7 @@ class Settings(BaseSettings): # Temporal worker graceful shutdown configuration TEMPORAL_GRACEFUL_SHUTDOWN_TIMEOUT: int = 7200 # 2 hours in seconds + WORKER_BIND_HOST: str = "0.0.0.0" # Control server and Temporal SDK metrics WORKER_METRICS_PORT: int = 8888 # Port for /drain and /health endpoints METRICS_PORT: int = 9090 # Port for Prometheus metrics endpoint METRICS_HOST: str = "0.0.0.0" # Bind address for metrics server diff --git a/backend/airweave/domains/temporal/worker/__init__.py b/backend/airweave/domains/temporal/worker/__init__.py index d1c8e83e55..077aafdef2 100644 --- a/backend/airweave/domains/temporal/worker/__init__.py +++ b/backend/airweave/domains/temporal/worker/__init__.py @@ -63,7 +63,7 @@ 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}", ), ), ) 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_worker.py b/backend/airweave/domains/temporal/worker/tests/test_worker.py index 339edf7b87..4bbf121279 100644 --- a/backend/airweave/domains/temporal/worker/tests/test_worker.py +++ b/backend/airweave/domains/temporal/worker/tests/test_worker.py @@ -8,6 +8,7 @@ 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 @@ -33,7 +34,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 +47,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 +184,25 @@ 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() diff --git a/deploy/README.md b/deploy/README.md index 84c782e6f1..d70d2f673b 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -120,3 +120,12 @@ coverage from one healthy request. Manifest syntax reference: https://docs.porter.run/applications/configuration-as-code/reference Private service semantics: https://docs.porter.run/applications/configuration-as-code/services/web-service + +## Private local worker listeners + +Set `WORKER_BIND_HOST=127.0.0.1` for a host-local worker. This binds both its +health/metrics/drain control server and Temporal SDK metrics exporter to loopback. +The default remains `0.0.0.0` for container networking. Binding limits network +reachability; it does not authenticate these endpoints, including `POST /drain`. +Keep container worker ports private. API listener and API metrics bindings are +configured separately. From a35cd82d6aec57eb0627f21a1970b53d4ac24abe Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 05:57:21 -0700 Subject: [PATCH 216/266] fix: flush verified source readiness before job admission --- .../domains/owned_provisioning/service.py | 3 + .../owned_provisioning/tests/conftest.py | 15 ++++- .../tests/test_provisioning.py | 56 +++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/backend/airweave/domains/owned_provisioning/service.py b/backend/airweave/domains/owned_provisioning/service.py index 5bf7da6950..c2a7496cdc 100644 --- a/backend/airweave/domains/owned_provisioning/service.py +++ b/backend/airweave/domains/owned_provisioning/service.py @@ -132,6 +132,9 @@ async def reconcile( 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() diff --git a/backend/airweave/domains/owned_provisioning/tests/conftest.py b/backend/airweave/domains/owned_provisioning/tests/conftest.py index 530354a78b..2916906a10 100644 --- a/backend/airweave/domains/owned_provisioning/tests/conftest.py +++ b/backend/airweave/domains/owned_provisioning/tests/conftest.py @@ -1,3 +1,14 @@ -"""Reuse the isolated canonical PostgreSQL migration fixture.""" +"""Use isolated PostgreSQL with the production session autoflush policy.""" -from airweave.domains.entities.canonical.tests.conftest import database # noqa: F401 +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 index 7290828b17..1e201e5900 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -513,3 +513,59 @@ async def test_outlook_fresh_source_retry_and_principal_cannot_downgrade_owned_c 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" From 654f2b8fd083b393ba3e4bf451647f73b1635b1b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 06:13:20 -0700 Subject: [PATCH 217/266] fix: recover mislabeled UTF-8 MIME text with extraction evidence --- .../entities/canonical/extraction_models.py | 10 ++++++ .../entities/canonical/gmail_projection.py | 14 +++++--- .../domains/entities/canonical/mime_text.py | 29 ++++++++++++++++ .../entities/canonical/outlook_projection.py | 12 +++++-- .../tests/test_extraction_coverage.py | 7 ++++ .../canonical/tests/test_gmail_projection.py | 18 ++++++++++ .../canonical/tests/test_mime_text.py | 33 +++++++++++++++++++ .../tests/test_outlook_projection.py | 17 ++++++++++ 8 files changed, 134 insertions(+), 6 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/mime_text.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_mime_text.py diff --git a/backend/airweave/domains/entities/canonical/extraction_models.py b/backend/airweave/domains/entities/canonical/extraction_models.py index 5c070129e1..f51778c655 100644 --- a/backend/airweave/domains/entities/canonical/extraction_models.py +++ b/backend/airweave/domains/entities/canonical/extraction_models.py @@ -5,6 +5,15 @@ from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator +class CharsetRecovery(BaseModel): + """A disclosed strict UTF-8 recovery after the attempted charset rejected bytes.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + source_path: str = Field(min_length=1, max_length=2048) + from_charset: str = Field(min_length=1, max_length=256) + to_charset: Literal["utf-8"] = "utf-8" + + class ExtractionPart(BaseModel): """Stable source-local identity and format; no private content or download URLs.""" @@ -14,6 +23,7 @@ class ExtractionPart(BaseModel): kind: Literal["body", "file", "record"] media_type: str | None = Field(default=None, max_length=256) extension: str | None = Field(default=None, max_length=32) + charset_recoveries: tuple[CharsetRecovery, ...] = () class ExtractionOutcome(ExtractionPart): diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py index e437e90f7d..c1e72f9d2e 100644 --- a/backend/airweave/domains/entities/canonical/gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -10,7 +10,8 @@ from pathlib import Path from airweave.domains.entities.canonical.blob_materializer import read_blob, write_blob -from airweave.domains.entities.canonical.extraction_models import ExtractionPart +from airweave.domains.entities.canonical.extraction_models import CharsetRecovery, ExtractionPart +from airweave.domains.entities.canonical.mime_text import DecodedText, 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 @@ -63,10 +64,10 @@ async def _body( return content -def _decode_text(content: bytes, part: dict) -> str: +def _decode_text(content: bytes, part: dict, path: str) -> DecodedText: header = Message() header["Content-Type"] = _header(part, "Content-Type") or part["mimeType"] - return content.decode(header.get_content_charset() or "utf-8", errors="strict") + return decode_mime_text(content, header.get_content_charset() or "utf-8", path) class GmailProjection: @@ -78,6 +79,7 @@ def __init__(self, record: SourceRecord, storage: StorageBackend, directory: Pat self.storage = storage self.directory = directory self.attachments: list[ProjectionInput] = [] + self.charset_recoveries: list[CharsetRecovery] = [] async def attachment(self, part: dict, path: str) -> None: """Project retained attachment bytes; keep missing bytes explicit on the original.""" @@ -139,7 +141,10 @@ async def render(self, part: dict, path: str) -> str: await self.attachment(part, path) return "" content = await _body(part, path, self.record, self.storage) - text = _decode_text(content, part) + decoded = _decode_text(content, part, path) + if decoded.recovery is not None: + self.charset_recoveries.append(decoded.recovery) + text = decoded.text return text if mime == "text/html" else "
" + html.escape(text) + "
" async def map(self) -> ProjectionInputs: @@ -174,6 +179,7 @@ async def map(self) -> ProjectionInputs: kind="body", media_type="text/html", extension=".html", + charset_recoveries=tuple(self.charset_recoveries), ), entity=entity, ), 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/outlook_projection.py b/backend/airweave/domains/entities/canonical/outlook_projection.py index 03f418909e..ee2b3920a8 100644 --- a/backend/airweave/domains/entities/canonical/outlook_projection.py +++ b/backend/airweave/domains/entities/canonical/outlook_projection.py @@ -13,7 +13,8 @@ from pathlib import Path from airweave.domains.entities.canonical.blob_materializer import read_blob, write_blob -from airweave.domains.entities.canonical.extraction_models import ExtractionPart +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 @@ -114,6 +115,7 @@ def __init__(self, record: SourceRecord, directory: Path, native: OutlookMessage 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.""" @@ -210,7 +212,12 @@ async def render(self, part: EmailMessage, path: str, *, body: bool = True) -> N if _file(part) or (not body and part.get("Content-ID")): await self.attachment(part, path) elif body: - text = _decoded(part).decode(part.get_content_charset() or "us-ascii", errors="strict") + 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) + "
" ) @@ -227,6 +234,7 @@ async def body(self) -> ProjectionInput: kind="body", media_type="text/html", extension=".html", + charset_recoveries=tuple(self.charset_recoveries), ), entity=OutlookMessageEntity( id=native.id, diff --git a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py index 88804cec63..181c418b71 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py +++ b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py @@ -92,6 +92,10 @@ async def test_message_pdf_video_and_missing_part_publish_truthful_coverage( }, completeness="partial", ) + item.payload["payload"]["parts"][0] = part( + "Intact fundraising email body: नमस्ते — café".encode(), + headers=[{"name": "Content-Type", "value": "text/plain; charset=gb2312"}], + ) await capture(database, service, fence, item) store = CanonicalProjectionStore() target = destination() @@ -112,6 +116,9 @@ async def test_message_pdf_video_and_missing_part_publish_truthful_coverage( "unavailable_original", ] assert coverage.parts[1].key == "/payload/parts/1" + recovery = coverage.parts[0].charset_recoveries[0] + assert recovery.source_path == "/payload/parts/0" + assert recovery.from_charset == "gb2312" and recovery.to_charset == "utf-8" assert ( coverage.parts[1].kind == "file" and coverage.parts[1].media_type == "application/pdf" ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py index c5ede3d4aa..80a950964c 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -83,6 +83,24 @@ async def test_plain_charset_body_and_quoted_address_are_preserved(tmp_path): assert source.model_dump() == before +@pytest.mark.asyncio +async def test_mislabeled_utf8_body_is_recovered_without_changing_original(tmp_path): + text = "

नमस्ते — 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( 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_outlook_projection.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py index 6a21e53cb2..5319173d49 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_projection.py @@ -105,6 +105,23 @@ async def test_mixed_related_alternative_preserves_unicode_metadata_and_attachme storage.read_file.assert_awaited_once() +@pytest.mark.asyncio +async def test_mislabeled_utf8_mime_retains_original_and_discloses_recovery(tmp_path): + text = "नमस्ते — café" + raw = ( + b"Content-Type: text/plain; charset=gb2312\r\nContent-Transfer-Encoding: 8bit\r\n\r\n" + + text.encode() + ) + record, storage = source(raw) + before = record.model_dump() + result = await map_outlook(record, storage, tmp_path) + assert text in Path(result.entities[0].local_path).read_text() + recovery = result.parts[0].part.charset_recoveries[0] + assert recovery.source_path == "/mime" and recovery.from_charset == "gb2312" + assert record.model_dump() == before + assert storage.read_file.return_value == raw + + @pytest.mark.asyncio async def test_latin_charset_and_explicit_related_root_are_respected(tmp_path): message = EmailMessage() From c36c7579af0d2dbc64e946f1f038f06d93613eba Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 06:27:47 -0700 Subject: [PATCH 218/266] fix: serialize validated source configuration for JSON persistence --- .../domains/sources/tests/test_validation.py | 28 +++++++++++++++++++ .../airweave/domains/sources/validation.py | 4 +-- 2 files changed, 30 insertions(+), 2 deletions(-) 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/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 From 65559cca015d54c8be5ea45e56ecd545d5b978f1 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 06:45:00 -0700 Subject: [PATCH 219/266] fix: preserve Drive representation omissions in search coverage --- .../entities/canonical/projection_mappers.py | 172 ++++++++++---- .../canonical/tests/test_drive_projection.py | 209 ++++++++++++++++++ .../canonical/tests/test_sheets_capture.py | 9 +- 3 files changed, 345 insertions(+), 45 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_drive_projection.py diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 18649a42d7..fc2a41f2f2 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -10,7 +10,7 @@ from pathlib import Path from tempfile import TemporaryDirectory -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event from airweave.domains.entities.canonical.extraction_models import ExtractionPart @@ -26,6 +26,7 @@ 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): @@ -49,32 +50,118 @@ def _datetime(value: JsonValue) -> datetime | None: 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, -) -> tuple[BaseEntity, ...]: - from airweave.domains.entities.canonical.blob_materializer import read_blob, write_blob +) -> 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 ( - GoogleDriveFolderEntity( - folder_id=record.identity.native_id, - title=_string(data.get("name"), default="Untitled"), - description=_string(data.get("description")), - breadcrumbs=[], - ), + 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): - from airweave.domains.entities.canonical.workspace_docs import read_document - 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" @@ -87,10 +174,7 @@ async def _drive( media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" suffix = ".xlsx" else: - document = await read_document(record, storage) - content = document.text().encode("utf-8") - media_type = "text/plain" - suffix = ".txt" + return await _drive_document(record, storage, directory) else: # Historical Drive records retain their export-only contract. if len(record.blobs) != 1: @@ -100,8 +184,32 @@ async def _drive( 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: - raise ProjectionMappingError("Drive captured bytes have no known file format") + 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, @@ -114,7 +222,7 @@ async def _drive( entity.mime_type = media_type entity.file_type = suffix.lstrip(".") entity.url = entity.web_url - return (entity,) + return _projection_input(0, entity) def _calendar(record: SourceRecord) -> tuple[BaseEntity, ...]: @@ -277,29 +385,6 @@ async def map_record( yield map_native(record) return - if ( - source_name == "google_drive" - and record.identity.record_type == "file" - and record.completeness == "metadata_only" - and not record.blobs - and record.payload.get("mimeType") != "application/vnd.google-apps.folder" - ): - # Capture explicitly omitted bytes. Publish the omission, not invented - # title-only document content; complete records still require their blob. - yield ProjectionInputs( - parts=( - ProjectionInput( - part=ExtractionPart( - part_index=0, - key=record.identity.native_id, - kind="file", - media_type=_string(record.payload.get("mimeType")) or None, - ), - entity=None, - ), - ) - ) - return with TemporaryDirectory(prefix="airweave-projection-") as temporary: directory = Path(temporary) if source_name == "gmail": @@ -322,6 +407,9 @@ async def map_record( yield await map_slack_files(record, _slack(record)[0], storage, directory) return + elif source_name == "google_drive": + yield await _drive(record, storage, directory) + return else: entities = await _map_entities(record, source_name, storage, directory) parts = tuple(_projection_input(index, entity) for index, entity in enumerate(entities)) @@ -364,8 +452,6 @@ async def _map_entities( from airweave.domains.entities.canonical.notion_projection import map_notion entities = map_notion(record) - elif source_name == "google_drive": - entities = await _drive(record, storage, directory) elif source_name == "google_calendar": entities = _calendar(record) elif source_name == "slack": diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py b/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py new file mode 100644 index 0000000000..7f9b8575a7 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py @@ -0,0 +1,209 @@ +"""Drive projection preserves retained representations and explicit content gaps.""" + +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.models import SourceRecord +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.projection_models import ProjectionWork +from airweave.domains.entities.canonical.projector import _select_inputs +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity +from airweave.platform.sources.records.workspace_manifest import ( + DOCS_MIME, + MANIFEST_MIME, + DocsGap, + DocsState, + ExportState, + WorkspaceManifestV1, + canonical_json, +) + + +def original(mime=DOCS_MIME): + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity(record_type="file", native_id="doc"), + revision=1, + payload={"id": "doc", "version": "7", "name": "document", "mimeType": mime}, + 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 retain(source, content, media_type, *, role=None): + digest = hashlib.sha256(content).hexdigest() + return BlobReference( + key=f"canonical/{source.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type=media_type, + role=role, + ) + + +def docs_record(*, native_available): + source = original() + if native_available: + content = canonical_json( + { + "documentId": "doc", + "tabs": [ + { + "tabProperties": {"tabId": "tab"}, + "documentTab": { + "body": { + "content": [ + { + "paragraph": { + "elements": [ + {"textRun": {"content": "retained authored text"}}, + { + "inlineObjectElement": { + "inlineObjectId": "image" + } + }, + ] + } + } + ] + } + }, + } + ], + } + ) + ref = retain(source, content, "application/json") + native = DocsState( + status="complete", + document_blob=ref.sha256, + embedded_media="not_retained", + missing=( + DocsGap( + source_path="/tabs/0/documentTab/inlineObjects/image/embeddedObject", + reason="unsupported_media", + ), + ), + ) + export = ExportState(status="unavailable", reason="export_size_limit") + else: + content = b"retained DOCX bytes for the downstream converter" + ref = retain( + source, + content, + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ) + native = DocsState( + status="unavailable", + embedded_media="not_retained", + missing=(DocsGap(reason="read_size_limit"),), + ) + export = ExportState(status="retained", blob=ref.sha256) + manifest = WorkspaceManifestV1(file_id="doc", drive_version="7", export=export, native=native) + raw = canonical_json(manifest.model_dump(mode="json")) + marked = retain(source, raw, MANIFEST_MIME, role="representation_manifest") + source = source.model_copy(update={"blobs": (ref, marked)}) + contents = {ref.key: content, marked.key: raw} + storage = AsyncMock() + storage.read_file.side_effect = lambda key, **kwargs: contents[key] + return source, storage, content + + +def coverage(result, source): + work = ProjectionWork( + organization_id=uuid4(), record=source, pipeline_version=2, previous_generation=None + ) + return _select_inputs(result, work, "google_drive", uuid4(), lambda _: True)[1] + + +@pytest.mark.asyncio +async def test_docs_text_keeps_missing_media_in_extraction_parts(): + source, storage, _ = docs_record(native_available=True) + before = source.model_dump() + async with map_record(source, "google_drive", storage) as result: + assert len(result.parts) == 2 + assert "retained authored text" in Path(result.entities[0].local_path).read_text() + missing = result.parts[1] + assert missing.entity is None and missing.omission is None + assert missing.part.key == "/tabs/0/documentTab/inlineObjects/image/embeddedObject" + assert missing.part.part_index == 1 + evidence = coverage(result, source) + assert evidence.status == "partial" + assert [part.outcome for part in evidence.parts] == ["indexed", "unavailable_original"] + assert source.model_dump() == before + assert storage.read_file.await_count == 2 + + +@pytest.mark.asyncio +async def test_docs_uses_retained_export_when_native_capture_unavailable(): + source, storage, content = docs_record(native_available=False) + before = source.model_dump() + async with map_record(source, "google_drive", storage) as result: + entity = result.entities[0] + assert Path(entity.local_path).suffix == ".docx" + assert Path(entity.local_path).read_bytes() == content + assert result.parts[1].entity is None and result.parts[1].part.key == "/native" + assert coverage(result, source).status == "partial" + assert source.model_dump() == before and source.completeness == "partial" + assert storage.read_file.await_count == 2 + + +@pytest.mark.asyncio +async def test_unknown_extensionless_original_is_unsupported_not_conversion_failure(): + source = original("application/x-almanac-unknown-binary") + content = b"retained opaque original" + ref = retain(source, content, source.payload["mimeType"]) + source = source.model_copy(update={"blobs": (ref,), "completeness": "complete"}) + storage = AsyncMock() + storage.read_file.return_value = content + before = source.model_dump() + async with map_record(source, "google_drive", storage) as result: + assert len(result.parts) == 1 and not result.entities + assert result.parts[0].omission == "unsupported_format" + assert coverage(result, source).status == "unavailable" + assert source.model_dump() == before + storage.read_file.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_docs_without_native_or_export_publishes_only_unavailable_coverage(): + source = original() + manifest = WorkspaceManifestV1( + file_id="doc", + drive_version="7", + export=ExportState(status="unavailable", reason="export_size_limit"), + native=DocsState( + status="unavailable", + embedded_media="not_retained", + missing=(DocsGap(reason="read_size_limit"),), + ), + ) + raw = canonical_json(manifest.model_dump(mode="json")) + marked = retain(source, raw, MANIFEST_MIME, role="representation_manifest") + source = source.model_copy(update={"blobs": (marked,)}) + before = source.model_dump() + storage = AsyncMock() + storage.read_file.return_value = raw + async with map_record(source, "google_drive", storage) as result: + assert not result.entities + assert len(result.parts) == 1 and result.parts[0].part.part_index == 0 + evidence = coverage(result, source) + assert evidence.status == "unavailable" + assert evidence.parts[0].outcome == "unavailable_original" + assert source.model_dump() == before and source.completeness == "partial" + storage.read_file.assert_awaited_once() diff --git a/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py b/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py index d53169f1de..59ea3fa124 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py +++ b/backend/airweave/domains/entities/canonical/tests/test_sheets_capture.py @@ -114,7 +114,7 @@ def response(request): record.revision + 1, files.storage, ) - entities = await _drive(record, files.storage, tmp_path / "projection") + entities = (await _drive(record, files.storage, tmp_path / "projection")).entities assert len(entities) == 1 from pathlib import Path @@ -135,7 +135,12 @@ def response(request): ) newer = (await capture(database, service, fence, partial)).changes[0].record assert newer.id == record.id and newer.revision == record.revision + 1 - fallback = await _drive(newer, files.storage, tmp_path / "fallback") + 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" From 73ae11288265425ea72c9270dff00381c8ec09c0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 07:05:46 -0700 Subject: [PATCH 220/266] Add bounded optional local OCR and per-file fallback --- .github/workflows/owned-source-store.yml | 3 + Makefile | 8 +- backend/Dockerfile | 7 + backend/airweave/core/config/settings.py | 4 + backend/airweave/core/container/factory.py | 19 +- .../core/container/tests/test_ocr_wiring.py | 39 ++++ backend/airweave/domains/ocr/fallback.py | 37 ++-- backend/airweave/domains/ocr/local.py | 109 +++++++++++ backend/airweave/domains/ocr/local_worker.py | 89 +++++++++ .../domains/ocr/tests/test_fallback.py | 10 +- .../airweave/domains/ocr/tests/test_local.py | 182 ++++++++++++++++++ backend/scripts/smoke_local_ocr.py | 32 +++ docs/owned-sources/local-ocr.md | 51 +++++ 13 files changed, 564 insertions(+), 26 deletions(-) create mode 100644 backend/airweave/core/container/tests/test_ocr_wiring.py create mode 100644 backend/airweave/domains/ocr/local.py create mode 100644 backend/airweave/domains/ocr/local_worker.py create mode 100644 backend/airweave/domains/ocr/tests/test_local.py create mode 100644 backend/scripts/smoke_local_ocr.py create mode 100644 docs/owned-sources/local-ocr.md diff --git a/.github/workflows/owned-source-store.yml b/.github/workflows/owned-source-store.yml index 8b91e8f747..e363ea8c76 100644 --- a/.github/workflows/owned-source-store.yml +++ b/.github/workflows/owned-source-store.yml @@ -98,6 +98,8 @@ jobs: run: make test-auth - name: Verify dependency readiness run: make test-health + - name: Verify OCR execution boundaries and fallback + run: make test-ocr - name: Verify projection recovery and maintenance workflows run: make test-worker - name: Record tested Vespa image @@ -121,3 +123,4 @@ jobs: run: | docker run --rm --entrypoint python almanac-source-store:ci -m compileall -q airweave docker run --rm --entrypoint /bin/sh almanac-source-store:ci -n /app/entrypoint.sh + docker run --rm --entrypoint python almanac-source-store:ci -m scripts.smoke_local_ocr --tessdata /usr/share/tessdata diff --git a/Makefile b/Makefile index 8a46e7c7b7..31c7d5d18e 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ PYTHON := .venv/bin/python PYTEST := .venv/bin/pytest IMAGE ?= almanac-source-store:local -.PHONY: help setup check test-store test-provisioning test-capture test-search test-auth test-health test-worker test-index validate-deploy build +.PHONY: help setup check test-store test-provisioning test-capture test-search test-auth test-health test-worker test-ocr test-index validate-deploy build help: @echo 'setup Install the locked backend and development dependencies' @@ -15,6 +15,7 @@ help: @echo 'test-auth Verify service-key and source authorization boundaries' @echo 'test-health Verify dependency probes and readiness behavior' @echo 'test-worker Verify owned projection, recovery and maintenance workflows' + @echo 'test-ocr Verify local OCR boundaries and per-file fallback' @echo 'test-index Deploy/test schemas on explicitly disposable local Vespa' @echo 'check Run store, capture, search, authentication and health checks' @echo 'validate-deploy Validate the Porter application manifest locally (no deployment)' @@ -55,8 +56,11 @@ test-worker: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py +test-ocr: + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/ocr/tests airweave/core/container/tests/test_ocr_wiring.py + test-index: @test "$$OWNED_VESPA_TEST" = 1 || (echo 'Set OWNED_VESPA_TEST=1 only for a disposable Vespa at localhost:8081/19071; this deploys schemas.' >&2; exit 1) cd $(BACKEND) && $(PYTEST) -q -o log_cli=false tests/integration/test_owned_vespa.py tests/integration/test_canonical_search_prefilters.py airweave/domains/entities/canonical/tests/test_real_vespa_projection.py -check: test-store test-provisioning test-capture test-search test-auth test-health test-worker +check: test-store test-provisioning test-capture test-search test-auth test-health test-worker test-ocr diff --git a/backend/Dockerfile b/backend/Dockerfile index b4b93f363d..3ea8070cdb 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -28,6 +28,13 @@ COPY pyproject.toml poetry.lock ./ RUN --mount=type=cache,target=/opt/poetry-cache \ poetry install --only=main --no-interaction --no-ansi --no-root +# Optional local OCR models: pinned upstream artifacts, verified by BuildKit. +# Inference never downloads weights. Enable explicitly with LOCAL_OCR_TESSDATA_PATH. +ADD --chmod=644 --checksum=sha256:7d4322bd2a7749724879683fc3912cb542f19906c83bcc1a52132556427170b2 https://raw.githubusercontent.com/tesseract-ocr/tessdata_fast/87416418657359cb625c412a48b6e1d6d41c29bd/eng.traineddata /usr/share/tessdata/eng.traineddata +ADD --chmod=644 --checksum=sha256:4c73ffc59d497c186b19d1e90f5d721d678ea6b2e277b719bee4e2af12271825 https://raw.githubusercontent.com/tesseract-ocr/tessdata_fast/87416418657359cb625c412a48b6e1d6d41c29bd/hin.traineddata /usr/share/tessdata/hin.traineddata + +ADD --chmod=644 --checksum=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30 https://raw.githubusercontent.com/tesseract-ocr/tessdata_fast/87416418657359cb625c412a48b6e1d6d41c29bd/LICENSE /usr/share/tessdata/LICENSE + # Copy entrypoint script with proper permissions COPY --chmod=755 entrypoint.sh /app/entrypoint.sh diff --git a/backend/airweave/core/config/settings.py b/backend/airweave/core/config/settings.py index 9b5a520465..6d4d4d9248 100644 --- a/backend/airweave/core/config/settings.py +++ b/backend/airweave/core/config/settings.py @@ -216,6 +216,10 @@ class Settings(BaseSettings): # in-code default in domains/search/config.py. LLM_FALLBACK_CHAIN: Optional[str] = None + # Optional local OCR; models are provisioned at build/setup time, never during inference. + LOCAL_OCR_TESSDATA_PATH: Optional[str] = None + LOCAL_OCR_LANGUAGES: tuple[str, ...] = ("eng",) + # Docling OCR fallback service (None = disabled) DOCLING_BASE_URL: Optional[str] = None diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index f7f6124a7a..7eec94001b 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -103,6 +103,7 @@ ) from airweave.domains.ocr.docling import DoclingOcrAdapter from airweave.domains.ocr.fallback import FallbackOcrProvider +from airweave.domains.ocr.local import LocalOcrProvider from airweave.domains.ocr.mistral.converter import MistralOCR from airweave.domains.ocr.protocols import OcrProvider from airweave.domains.organizations.protocols import UserOrganizationRepositoryProtocol @@ -769,18 +770,28 @@ def _create_ocr_provider( ) -> Optional[OcrProvider]: """Create OCR provider with fallback chain. - Chain order: Mistral (cloud) -> Docling (local service, if configured). + Chain order: local Tesseract (if configured) -> Mistral -> Docling. Docling is only added when DOCLING_BASE_URL is set. Returns None with a warning when no providers are available. """ + providers = [] + if settings.LOCAL_OCR_TESSDATA_PATH: + providers.append( + ( + "local-tesseract", + LocalOcrProvider( + tessdata_path=settings.LOCAL_OCR_TESSDATA_PATH, + languages=settings.LOCAL_OCR_LANGUAGES, + ), + ) + ) try: - mistral_ocr = MistralOCR() + mistral_ocr = MistralOCR() if settings.MISTRAL_API_KEY else None except Exception as e: logger.error(f"Error creating Mistral OCR adapter: {e}") mistral_ocr = None - providers = [] if mistral_ocr: providers.append(("mistral-ocr", mistral_ocr)) @@ -795,7 +806,7 @@ def _create_ocr_provider( if not providers: logger.warning( "No OCR providers available — document processing will be disabled. " - "Set MISTRAL_API_KEY or DOCLING_BASE_URL to enable OCR." + "Set LOCAL_OCR_TESSDATA_PATH, MISTRAL_API_KEY or DOCLING_BASE_URL to enable OCR." ) return None diff --git a/backend/airweave/core/container/tests/test_ocr_wiring.py b/backend/airweave/core/container/tests/test_ocr_wiring.py new file mode 100644 index 0000000000..6701ef5d5c --- /dev/null +++ b/backend/airweave/core/container/tests/test_ocr_wiring.py @@ -0,0 +1,39 @@ +"""Local OCR is opt-in and precedes metered fallback without duplicate processing.""" + +from unittest.mock import patch + +import pytest + +from airweave.adapters.circuit_breaker.fake import FakeCircuitBreaker +from airweave.core.config.settings import Settings +from airweave.core.container.factory import _create_ocr_provider +from airweave.domains.ocr.fakes.provider import FakeOcrProvider + + +def test_no_configured_ocr_remains_disabled(): + settings = Settings.model_construct( + LOCAL_OCR_TESSDATA_PATH=None, MISTRAL_API_KEY=None, DOCLING_BASE_URL=None + ) + assert _create_ocr_provider(FakeCircuitBreaker(), settings) is None + + +@pytest.mark.asyncio +async def test_local_ocr_precedes_cloud_and_receives_configured_languages(): + settings = Settings.model_construct( + LOCAL_OCR_TESSDATA_PATH="/models", + LOCAL_OCR_LANGUAGES=("eng", "hin"), + MISTRAL_API_KEY="test-only", + DOCLING_BASE_URL=None, + ) + local = FakeOcrProvider(default_markdown="local result") + cloud = FakeOcrProvider(default_markdown="cloud result") + with ( + patch( + "airweave.core.container.factory.LocalOcrProvider", return_value=local + ) as constructor, + patch("airweave.core.container.factory.MistralOCR", return_value=cloud), + ): + provider = _create_ocr_provider(FakeCircuitBreaker(), settings) + constructor.assert_called_once_with(tessdata_path="/models", languages=("eng", "hin")) + assert await provider.convert_batch(["scan.pdf"]) == {"scan.pdf": "local result"} + assert cloud.calls == [] 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/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/docs/owned-sources/local-ocr.md b/docs/owned-sources/local-ocr.md new file mode 100644 index 0000000000..fb7b8579fe --- /dev/null +++ b/docs/owned-sources/local-ocr.md @@ -0,0 +1,51 @@ +# Optional local OCR + +The shared OCR provider can use PyMuPDF's bundled Tesseract runtime with explicitly +provisioned language models. It uses the existing conversion interface and does +not introduce a service, database, or Python dependency. Original bytes are +unchanged; recognized text is a derived representation, not an exact transcription. + +## Configuration + +The backend image includes checksum-pinned `tessdata_fast` English and Hindi +models and their Apache 2.0 license under `/usr/share/tessdata`. Enable the adapter: + +```sh +LOCAL_OCR_TESSDATA_PATH=/usr/share/tessdata +LOCAL_OCR_LANGUAGES='["eng","hin"]' +``` + +Without the path, local OCR stays disabled. Missing configured models fail startup. +Inference does not download models. Other languages require provisioning their +model files and explicitly adding their language codes. + +The configured chain is local OCR, then Mistral, then Docling. Only unresolved +files reach the next configured provider; successful files are not processed +again. A configured cloud fallback can incur its normal charges. Local OCR uses +CPU and memory, so zero external inference calls does not mean zero hosting cost. + +## Bounds and failure behavior + +Each file runs sequentially in a disposable process, killed and reaped on timeout +or cancellation. Limits: 120 seconds, 200 MiB input, 200 pages, 20 million rendered +pixels per page at 150 DPI, and 8 MiB UTF-8 output. PDF, PNG, JPEG, BMP, WebP and +single-frame TIFF are supported. Encrypted PDFs and multiframe images are refused. +An over-limit or failed file returns no text to the fallback chain; it does not +publish a truncated prefix. These are processing bounds, not a full OS memory +sandbox. Native parser memory use still depends on the file. + +Full-page OCR is intentional: the installed PyMuPDF partial-OCR path can omit +vector-drawn content. This may recognize native text less faithfully than direct +extraction. Original downloads remain the reference for exact reading. + +## Qualification + +`make test-ocr` checks fallback behavior, configuration and subprocess boundaries. +Set `LOCAL_OCR_TEST_TESSDATA` to a model directory to include real extraction tests. +The image CI runs `python -m scripts.smoke_local_ocr --tessdata /usr/share/tessdata`. + +Local evaluation recovered a retained three-page PDF that previously failed +conversion, without changing its original hash. A synthetic English/Hindi image +recovered the English identifier and Hindi characters but dropped a space between +Hindi words. This demonstrates execution and a real quality limitation; it is not +a multilingual accuracy benchmark or proof of complete provider indexing. From e5785db7cdd39b439e3aa77417ebe73a77448710 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 07:25:03 -0700 Subject: [PATCH 221/266] Reject disabled Composio accounts and auth configurations --- .../auth_provider/providers/composio.py | 33 ++++++++++++- .../tests/test_composio_binding.py | 46 ++++++++++++++++++- docs/owned-sources/local-ocr.md | 5 ++ 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index d5cb6cf586..e336e8a338 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -5,6 +5,7 @@ from uuid import UUID import httpx +from pydantic import BaseModel, Field, StrictBool, ValidationError from airweave.core.credential_sanitizer import ( safe_log_credentials, @@ -13,6 +14,7 @@ from airweave.domains.auth_provider._base import BaseAuthProvider from airweave.domains.auth_provider.exceptions import ( AuthProviderAccountNotFoundError, + AuthProviderConfigError, AuthProviderMissingFieldsError, AuthProviderRateLimitError, AuthProviderTemporaryError, @@ -25,6 +27,29 @@ from airweave.domains.auth_provider.auth_result import AuthResult +class _AuthConfigAvailability(BaseModel): + """Read the optional broker flag; unrelated provider fields remain opaque.""" + + is_disabled: StrictBool = False + + +class _AccountAvailability(BaseModel): + """Missing flags preserve compatibility; present flags must be actual booleans.""" + + is_disabled: StrictBool = False + auth_config: _AuthConfigAvailability = Field(default_factory=_AuthConfigAvailability) + + @classmethod + def from_response(cls, account: dict) -> "_AccountAvailability": + """Reject malformed flags without exposing provider data in validation errors.""" + try: + return cls.model_validate(account) + except ValidationError: + raise AuthProviderConfigError( + "Connected account availability is invalid", provider_name="composio" + ) from None + + @auth_provider( name="Composio", short_name="composio", @@ -105,7 +130,6 @@ async def get_auth_result( ) -> "AuthResult": """Bind managed access without requesting or extracting OAuth tokens.""" from airweave.domains.auth_provider.auth_result import AuthResult - from airweave.domains.auth_provider.exceptions import AuthProviderConfigError from airweave.domains.sources.token_providers.protocol import ( ManagedAuthProvider, ManagedToolAuthProvider, @@ -140,6 +164,7 @@ async def get_auth_result( "https://backend.composio.dev/api/v3/connected_accounts/" + quote(self.account_id, safe=""), ) + availability = _AccountAvailability.from_response(account) if account.get("id") != self.account_id: raise AuthProviderConfigError( "Connected account identity does not match", provider_name="composio" @@ -158,7 +183,11 @@ async def get_auth_result( raise AuthProviderConfigError( "Connected account auth config does not match", provider_name="composio" ) - if account.get("status", "").upper() != "ACTIVE": + if ( + account.get("status", "").upper() != "ACTIVE" + or availability.is_disabled + or availability.auth_config.is_disabled + ): from airweave.domains.auth_provider.exceptions import AuthProviderAuthError raise AuthProviderAuthError( diff --git a/backend/airweave/domains/auth_provider/tests/test_composio_binding.py b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py index fc32a2ede8..d8f3d6c3dc 100644 --- a/backend/airweave/domains/auth_provider/tests/test_composio_binding.py +++ b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py @@ -4,7 +4,10 @@ import pytest -from airweave.domains.auth_provider.exceptions import AuthProviderConfigError +from airweave.domains.auth_provider.exceptions import ( + AuthProviderAuthError, + AuthProviderConfigError, +) from airweave.domains.auth_provider.providers.composio import ComposioAuthProvider @@ -108,3 +111,44 @@ async def test_wispr_still_requires_configured_user(monkeypatch): ) with pytest.raises(AuthProviderConfigError, match="explicit Composio user"): await provider.get_auth_result("wispr", []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "account_disabled,config_disabled,error", + [ + (False, False, None), + (True, False, AuthProviderAuthError), + (False, True, AuthProviderAuthError), + ("false", False, AuthProviderConfigError), + (False, 0, AuthProviderConfigError), + (None, False, AuthProviderConfigError), + ], +) +async def test_managed_binding_validates_disabled_flags( + monkeypatch, account_disabled, config_disabled, error +): + provider = await ComposioAuthProvider.create( + credentials={"api_key": "secret"}, + config={"account_id": "selected", "user_id": "owner", "auth_config_id": "cfg"}, + ) + monkeypatch.setattr( + provider, + "_get_with_auth", + AsyncMock( + return_value={ + "id": "selected", + "user_id": "owner", + "toolkit": {"slug": "wispr_flow_mcp"}, + "status": "ACTIVE", + "is_disabled": account_disabled, + "auth_config": {"id": "cfg", "is_disabled": config_disabled}, + } + ), + ) + if error is None: + result = await provider.get_auth_result("wispr", []) + assert result.managed_auth.connected_account_id == "selected" + else: + with pytest.raises(error): + await provider.get_auth_result("wispr", []) diff --git a/docs/owned-sources/local-ocr.md b/docs/owned-sources/local-ocr.md index fb7b8579fe..1047e4669c 100644 --- a/docs/owned-sources/local-ocr.md +++ b/docs/owned-sources/local-ocr.md @@ -19,6 +19,11 @@ Without the path, local OCR stays disabled. Missing configured models fail start Inference does not download models. Other languages require provisioning their model files and explicitly adding their language codes. +Changing OCR configuration does not automatically rebuild published generations. +Retry failed projections explicitly; use the existing source reprojection/version +operation when replacing already-published text. Do not silently mix a changed +extraction configuration into a claim that an old generation was reprocessed. + The configured chain is local OCR, then Mistral, then Docling. Only unresolved files reach the next configured provider; successful files are not processed again. A configured cloud fallback can incur its normal charges. Local OCR uses From eb07ba7c99364ac57bc2fc3a41394f6a7e1f7a96 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 07:33:01 -0700 Subject: [PATCH 222/266] Retain Wispr notes when transcripts are explicitly absent --- .../entities/canonical/projection_mappers.py | 18 +++++++- backend/airweave/platform/sources/wispr.py | 9 ++++ .../entities/test_projection_mappers.py | 29 +++++++++++++ .../platform/sources/test_wispr_capture.py | 43 +++++++++++++++++++ 4 files changed, 98 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index fc2a41f2f2..370255daa9 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -276,6 +276,22 @@ def _slack(record: SourceRecord) -> tuple[BaseEntity, ...]: 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 + + def _wispr_text(responses: list[JsonValue], field: str) -> str: """Follow exactly the requested ranges; strip server continuation guidance.""" expected = 0 @@ -293,7 +309,7 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: continue if not isinstance(window, dict) or window.get("start_char") != expected or complete: raise ProjectionMappingError("Wispr text ranges overlap or have a gap") - text = _string(response.get(field)) + text = _wispr_range_text(response, field, expected) markers = list( re.finditer( rf"(?m)^\(\.\.\.truncated, \d+ chars remaining; continue with " diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 34e8e08dcd..c1fece6a77 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -269,6 +269,15 @@ async def _body( responses.append({"requested_ranges": arguments, "response": response}) 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: diff --git a/backend/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 78b757da04..6097a4d308 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -284,3 +284,32 @@ async def test_wispr_projection_rejects_ambiguous_stored_continuation(): record("scratchpad_note", payload, native_id="note"), "wispr", AsyncMock() ): pass + + +@pytest.mark.asyncio +async def test_wispr_explicitly_absent_transcript_keeps_available_notes(): + native = { + "id": "meeting", + "content": "available notes", + "has_transcript": False, + "transcript": None, + } + payload = { + "responses": [ + { + "requested_ranges": { + "view_content": {"start_char": 0}, + "view_transcript": {"start_char": 0}, + }, + "response": native, + } + ] + } + async with map_record(record("meeting", payload), "wispr", AsyncMock()) as result: + assert result.entities[0].notes == "available notes" + assert result.entities[0].transcript == "" + assert native["transcript"] is None + native["has_transcript"] = True + with pytest.raises(ProjectionMappingError, match="unavailable or malformed"): + async with map_record(record("meeting", payload), "wispr", AsyncMock()): + pass diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index 4ba6a619c6..062a0bd074 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -423,3 +423,46 @@ async def test_scratchpad_cap_never_substitutes_meeting_start_for_modified_time( await connector.capture_page( CompletedScope(record_type="scratchpad_listing"), ScanContinuation(), files=MagicMock() ) + + +@pytest.mark.asyncio +async def test_explicit_no_transcript_preserves_notes_and_native_null(monkeypatch): + connector = await source() + native = {"id": "m", "content": "available notes", "has_transcript": False, "transcript": None} + execute = AsyncMock(return_value=native) + monkeypatch.setattr(connector, "_execute", execute) + result = await connector._body("meeting", "m") + assert result["responses"][0]["response"] == native + assert execute.await_count == 1 + for ambiguous in ( + {**native, "has_transcript": True}, + {k: v for k, v in native.items() if k != "transcript"}, + ): + execute.return_value = ambiguous + with pytest.raises(ValueError, match="not a string"): + await connector._body("meeting", "m") + + +@pytest.mark.asyncio +async def test_transcript_disappearing_after_first_range_stays_incomplete(monkeypatch): + connector = await source() + monkeypatch.setattr( + connector, + "_execute", + AsyncMock( + side_effect=[ + { + "id": "m", + "content": "notes", + "has_transcript": True, + "transcript": ( + "abc\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)" + ), + }, + {"id": "m", "has_transcript": False, "transcript": None}, + ] + ), + ) + with pytest.raises(ValueError, match="not a string"): + await connector._body("meeting", "m") From 3791d83d994d9f0469612b5f4178f61d8b596b00 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 07:40:42 -0700 Subject: [PATCH 223/266] Load publication visibility with the locked canonical record --- .../entities/canonical/projection_store.py | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index a0d56d68b4..e147245452 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -191,26 +191,28 @@ async def _current(self, db: AsyncSession, work: ProjectionWork) -> Entity | Non ) if sync is None or sync.index_pipeline_version != work.pipeline_version: return None - entity = await db.scalar( - select(Entity) - .where( - Entity.id == work.record.id, - Entity.sync_id == work.record.sync_id, - Entity.organization_id == work.organization_id, + row = ( + await db.execute( + select(Entity, content_is_available()) + .where( + Entity.id == work.record.id, + Entity.sync_id == work.record.sync_id, + Entity.organization_id == work.organization_id, + ) + .with_for_update(of=Entity) + .execution_options(populate_existing=True) ) - .with_for_update() - .execution_options(populate_existing=True) - ) + ).one_or_none() + if row is None: + return None + entity, visible = row if ( - entity is None - or entity.record_revision != work.record.revision + entity.record_revision != work.record.revision or entity.indexed_generation != work.previous_generation ): return None - if entity.deleted_at is None: - visible = await db.scalar(select(content_is_available()).where(Entity.id == entity.id)) - if not visible: - return None + if entity.deleted_at is None and not visible: + return None return entity async def prepare( From 780a3aa4b9cab40885d60679702ffa08d1994b03 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:01:31 -0700 Subject: [PATCH 224/266] Report explicitly judged duplicate crowding in retrieval evaluation --- evaluation/README.md | 33 ++++++++++++- evaluation/retrieval.py | 50 ++++++++++++++++++++ evaluation/tests/test_retrieval.py | 76 ++++++++++++++++++++++++++++++ 3 files changed, 158 insertions(+), 1 deletion(-) diff --git a/evaluation/README.md b/evaluation/README.md index 4f33ac6235..fc15bbb5bf 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -58,6 +58,37 @@ empty-success outcome is reported separately, not averaged into nDCG/recall. An error returning nothing is not a correct abstention. No-answer labels need evidence; merely failing to find an item does not establish absence. +### Duplicate crowding + +Distinct record IDs may contain redundant answers. A session trial returned19 +identical messages in20 keyword slots; ordinary record-level relevance metrics +can reward every copy. Optionally label redundant originals per query: + +```json +"duplicate_groups": [ + {"record_ids": ["session:message-a", "session:message-b"], "assessor": "human"} +] +``` + +Members must be distinct, disjoint across groups, and have explicit equal +relevance judgments. Judge redundancy for the query's intent and filters: +identical text at different dates need not answer a date-specific query equally. +Sharing a session, mailbox or channel alone does not establish redundancy. +Agent-generated labels must use `assessor: "agent"`, even when computed from +exact original text. Keep authoritative message identities intact. + +The per-query `duplicates` report counts labeled records, extra slots after the +first member of each group, and unassigned records in the delivered top20. +Unassigned records are not claimed unique. Without duplicate labels the report +is null, not zero. Counts are descriptive; a failed/empty request is not proof of +good diversity. Existing nDCG/MRR/precision/recall scores use the unchanged ranking; +the evaluator never silently collapses results to improve scores. + +Duplicate labels participate in the dataset fingerprint and assessor provenance. +This schema revision also changes fingerprints for previously parsed datasets; +regenerate the fingerprint and bind runs to the validated dataset deliberately. +Do not copy a hash from a different set of labels to make a comparison pass. + ## Run ```json @@ -109,7 +140,7 @@ the denominator. Inspect per-query and tagged regressions alongside aggregates. Real-corpus collection, request timing instrumentation, reviewed judgments, snippet evaluation and corpus-freezing automation are still separate work. -The four focused synthetic tests verify evaluator behavior, not search quality. +The focused synthetic tests verify evaluator behavior, not search quality. Reference: [ir-measures interfaces and formats](https://ir-measur.es/en/latest/getting-started.html). diff --git a/evaluation/retrieval.py b/evaluation/retrieval.py index 22f4c4f9e3..c2db8f5726 100644 --- a/evaluation/retrieval.py +++ b/evaluation/retrieval.py @@ -35,6 +35,13 @@ class Judgment(Value): assessor: Literal["human", "agent", "synthetic"] +class DuplicateGroup(Value): + """Distinct originals judged redundant for this query's intent and filters.""" + + record_ids: tuple[Identifier, ...] = Field(min_length=2) + assessor: Literal["human", "agent", "synthetic"] + + class Query(Value): """A fixed intent with explicit known-answer or no-answer expectation.""" @@ -44,6 +51,7 @@ class Query(Value): expectation: Literal["relevant_records", "no_answer"] expectation_assessor: Literal["human", "agent", "synthetic"] judgments: tuple[Judgment, ...] + duplicate_groups: tuple[DuplicateGroup, ...] = () @model_validator(mode="after") def validate_judgments(self): @@ -56,6 +64,21 @@ def validate_judgments(self): raise ValueError("Query expectation conflicts with positive judgments") if not self.text.strip() or len(self.tags) != len(set(self.tags)): raise ValueError("Query text must be nonblank and tags unique") + grades = {j.record_id: j.relevance for j in self.judgments} + assigned = set() + for group in self.duplicate_groups: + for record in group.record_ids: + if record in assigned: + raise ValueError( + "Duplicate-group members must be disjoint and unique" + ) + if record not in grades: + raise ValueError( + "Duplicate-group members require relevance judgments" + ) + assigned.add(record) + if len({grades[record] for record in group.record_ids}) != 1: + raise ValueError("Duplicate-group members must have equal relevance") return self @@ -136,6 +159,15 @@ def unique_observations(self): return self +class DuplicateReport(Value): + """Known redundant slots; unassigned records are not declared unique.""" + + cutoff: Literal[20] = 20 + grouped_records: int + duplicate_slots: int + unassigned_records: int + + class QueryReport(Value): """Scores retain failures and label provenance, without copying private query text.""" @@ -149,6 +181,7 @@ class QueryReport(Value): judged_in_top_10: int unjudged_in_top_10: int correctly_empty: bool | None + duplicates: DuplicateReport | None class LatencyReport(Value): @@ -219,6 +252,13 @@ def summarize(dataset: Dataset, run: Run) -> Report: } judged = {j.record_id for j in query.judgments} top = result.record_ids[:10] + membership = { + record: index + for index, group in enumerate(query.duplicate_groups) + for record in group.record_ids + } + duplicate_window = result.record_ids[:20] + assigned = [membership[r] for r in duplicate_window if r in membership] reports.append( QueryReport( query_id=query.id, @@ -229,6 +269,7 @@ def summarize(dataset: Dataset, run: Run) -> Report: sorted( {query.expectation_assessor} | {j.assessor for j in query.judgments} + | {g.assessor for g in query.duplicate_groups} ) ), scores=scores, @@ -238,6 +279,13 @@ def summarize(dataset: Dataset, run: Run) -> Report: correctly_empty=(result.status == "success" and not result.record_ids) if query.expectation == "no_answer" else None, + duplicates=DuplicateReport( + grouped_records=len(assigned), + duplicate_slots=len(assigned) - len(set(assigned)), + unassigned_records=len(duplicate_window) - len(assigned), + ) + if query.duplicate_groups + else None, ) ) answerable = [q for q in reports if q.expectation == "relevant_records"] @@ -265,6 +313,8 @@ def summarize(dataset: Dataset, run: Run) -> Report: "Recall is against judged relevant records, not all connected-account content.", "Unjudged records receive no relevance credit; pool and judge them before conclusions.", "Latency uses nearest-rank percentiles; small samples do not establish tail reliability.", + "Duplicate counts use explicit query-specific labels; unassigned records may also repeat.", + "Relevance metrics score the delivered ranking unchanged, including distinct redundant originals.", ] if any(assessor != "human" for q in reports for assessor in q.assessors): warnings.append( diff --git a/evaluation/tests/test_retrieval.py b/evaluation/tests/test_retrieval.py index 671e177e83..b85bc35a0c 100644 --- a/evaluation/tests/test_retrieval.py +++ b/evaluation/tests/test_retrieval.py @@ -5,6 +5,7 @@ from evaluation.retrieval import ( Dataset, + DuplicateGroup, Judgment, Query, Result, @@ -79,6 +80,7 @@ def test_failed_query_counts_and_unjudged_is_not_labeled_irrelevant(): assert report.queries[0].unjudged_in_top_10 == 1 assert report.queries[2].correctly_empty is False assert any("synthetic" in warning for warning in report.warnings) + assert report.queries[0].duplicates is None def test_latency_conditions_failures_and_tail_are_separate(): @@ -154,3 +156,77 @@ def test_no_answer_success_is_separate_from_answerable_metrics(): report = summarize(dataset, observation) assert report.aggregate == {} assert report.queries[0].correctly_empty is True + + +def test_redundant_originals_remain_in_scores_and_unknown_diversity_stays_unknown(): + query = Query( + id="copies", + text="Find the original note", + expectation="relevant_records", + expectation_assessor="human", + judgments=tuple( + Judgment(record_id=record, relevance=3, assessor="human") + for record in ("a", "b", "c") + ), + duplicate_groups=( + DuplicateGroup(record_ids=("a", "b", "c"), assessor="agent"), + ), + ) + dataset = Dataset(corpus_id="copies", version="v1", queries=(query,)) + observation = Run( + system="copies", + corpus_id=dataset.corpus_id, + dataset_sha256=dataset.fingerprint(), + results=( + Result( + query_id=query.id, + status="partial", + record_ids=("a", "b", "c", "unknown"), + ), + ), + ) + report = summarize(dataset, observation) + row = report.queries[0] + assert row.scores["nDCG@10"] == 1 + assert row.scores["Precision@5"] == 0.6 + assert row.returned == 4 and row.status == "partial" + assert row.duplicates.grouped_records == 3 + assert row.duplicates.duplicate_slots == 2 + assert row.duplicates.unassigned_records == 1 + assert row.assessors == ("agent", "human") + assert any("not exclusively human" in warning for warning in report.warnings) + outside = observation.model_copy( + update={ + "results": ( + Result( + query_id=query.id, + status="success", + record_ids=tuple(f"unassigned-{i}" for i in range(20)) + ("a", "b"), + ), + ) + } + ) + cutoff = summarize(dataset, outside).queries[0].duplicates + assert ( + cutoff.grouped_records, + cutoff.duplicate_slots, + cutoff.unassigned_records, + ) == (0, 0, 20) + unlabeled = query.model_copy(update={"duplicate_groups": ()}) + assert ( + dataset.fingerprint() + != dataset.model_copy(update={"queries": (unlabeled,)}).fingerprint() + ) + + +@pytest.mark.parametrize( + "members", + [(("meeting", "meeting"),), (("meeting", "missing"),), (("meeting", "spam"),)], +) +def test_duplicate_labels_cannot_repeat_invent_or_contradict_judgments(members): + fields = corpus().queries[0].model_dump() + fields["duplicate_groups"] = [ + {"record_ids": records, "assessor": "synthetic"} for records in members + ] + with pytest.raises(ValidationError): + Query.model_validate(fields) From 9d2872b6dedb8089198602a2d307f40ee483a027 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:23:32 -0700 Subject: [PATCH 225/266] Add account-bound Stripe original capture and offline projection --- .../auth_provider/providers/composio.py | 1 + .../tests/test_composio_binding.py | 1 + .../entities/canonical/projection_mappers.py | 4 + .../entities/canonical/stripe_projection.py | 316 ++++++++++++++++++ .../canonical/tests/test_outlook_recovery.py | 36 +- .../canonical/tests/test_stripe_recovery.py | 114 +++++++ .../source_connections/tests/test_update.py | 31 ++ .../domains/source_connections/update.py | 16 +- .../airweave/domains/sync_pipeline/factory.py | 22 +- backend/airweave/domains/syncs/service.py | 6 +- .../domains/syncs/tests/test_service.py | 14 + backend/airweave/platform/configs/config.py | 24 +- backend/airweave/platform/entities/stripe.py | 48 +-- backend/airweave/platform/sources/stripe.py | 68 +++- .../platform/sources/stripe_capture.py | 276 +++++++++++++++ .../entities/test_stripe_projection.py | 277 +++++++++++++++ .../platform/sources/test_stripe_capture.py | 214 ++++++++++++ 17 files changed, 1423 insertions(+), 45 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/stripe_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py create mode 100644 backend/airweave/platform/sources/stripe_capture.py create mode 100644 backend/tests/unit/domains/entities/test_stripe_projection.py create mode 100644 backend/tests/unit/platform/sources/test_stripe_capture.py diff --git a/backend/airweave/domains/auth_provider/providers/composio.py b/backend/airweave/domains/auth_provider/providers/composio.py index e336e8a338..9c672007a5 100644 --- a/backend/airweave/domains/auth_provider/providers/composio.py +++ b/backend/airweave/domains/auth_provider/providers/composio.py @@ -148,6 +148,7 @@ async def get_auth_result( "github": frozenset({"api.github.com"}), "notion": frozenset({"api.notion.com"}), "attio": frozenset({"api.attio.com"}), + "stripe": frozenset({"api.stripe.com"}), } if source_short_name not in hosts and source_short_name != "wispr": raise AuthProviderConfigError( diff --git a/backend/airweave/domains/auth_provider/tests/test_composio_binding.py b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py index d8f3d6c3dc..18251b08cb 100644 --- a/backend/airweave/domains/auth_provider/tests/test_composio_binding.py +++ b/backend/airweave/domains/auth_provider/tests/test_composio_binding.py @@ -65,6 +65,7 @@ async def test_managed_binding_rejects_wrong_or_missing_identity( @pytest.mark.parametrize( "source,user_id", [ + ("stripe", "owner"), ("notion", None), ("notion", "owner"), ("wispr", "owner"), diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 370255daa9..d9bd83b276 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -464,6 +464,10 @@ async def _map_entities( from airweave.domains.entities.canonical.attio_projection import map_attio entities = map_attio(record) + elif source_name == "stripe": + from airweave.domains.entities.canonical.stripe_projection import map_stripe + + entities = map_stripe(record) elif source_name == "notion": from airweave.domains.entities.canonical.notion_projection import map_notion diff --git a/backend/airweave/domains/entities/canonical/stripe_projection.py b/backend/airweave/domains/entities/canonical/stripe_projection.py new file mode 100644 index 0000000000..0ed651531a --- /dev/null +++ b/backend/airweave/domains/entities/canonical/stripe_projection.py @@ -0,0 +1,316 @@ +"""Offline allowlisted Stripe views; amounts retain native currency units. + +No calculations, expanded-object recursion, secret-bearing metadata, or provider +fetches. Event objects remain historical events, never current financial records. +""" + +from datetime import datetime, timezone +from typing import Annotated + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.platform.entities._base import BaseEntity +from airweave.platform.entities.stripe import ( + StripeBalanceEntity, + StripeBalanceTransactionEntity, + StripeChargeEntity, + StripeCustomerEntity, + StripeEventEntity, + StripeInvoiceEntity, + StripePaymentIntentEntity, + StripePaymentMethodEntity, + StripePayoutEntity, + StripeRefundEntity, + StripeSubscriptionEntity, +) + +Timestamp = Annotated[int, Field(ge=0, le=253402300799)] + + +class _Native(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + + +class _Reference(_Native): + id: str + + +class _Money(_Native): + amount: int + currency: str + + +class _Billing(_Native): + name: str | None = None + email: str | None = None + + +class _EventObject(_Native): + id: str | None = None + object: str + + +class _EventData(_Native): + object: _EventObject + + +class _Original(_Native): + id: str | None = None + object: str + created: Timestamp | None = None + livemode: bool | None = None + name: str | None = None + description: str | None = None + currency: str | None = None + status: str | None = None + type: str | None = None + amount: int | None = None + customer: str | _Reference | None = None + invoice: str | _Reference | None = None + charge: str | _Reference | None = None + payment_intent: str | _Reference | None = None + source: str | _Reference | None = None + destination: str | _Reference | None = None + captured: bool | None = None + paid: bool | None = None + refunded: bool | None = None + delinquent: bool | None = None + cancel_at_period_end: bool | None = None + email: str | None = None + phone: str | None = None + invoice_prefix: str | None = None + number: str | None = None + amount_due: int | None = None + amount_paid: int | None = None + amount_remaining: int | None = None + due_date: Timestamp | None = None + arrival_date: Timestamp | None = None + current_period_start: Timestamp | None = None + current_period_end: Timestamp | None = None + canceled_at: Timestamp | None = None + fee: int | None = None + net: int | None = None + reporting_category: str | None = None + method: str | None = None + reason: str | None = None + api_version: str | None = None + pending_webhooks: int | None = None + data: _EventData | None = None + billing_details: _Billing | None = None + available: list[_Money] | None = None + pending: list[_Money] | None = None + instant_available: list[_Money] | None = None + connect_reserved: list[_Money] | None = None + + +def _id(value: str | _Reference | None) -> str | None: + return value.id if isinstance(value, _Reference) else value + + +def _time(value: int | None) -> datetime | None: + return datetime.fromtimestamp(value, timezone.utc) if value is not None else None + + +def map_stripe(record: SourceRecord) -> tuple[BaseEntity, ...]: + """Project one retained root; reject identity/schema drift without logging originals.""" + if ( + record.payload_schema_version != 1 + or record.identity.container_id is not None + or record.parent is not None + or record.deleted_at is not None + or record.content_access != "available" + ): + raise ValueError("Stripe projection requires an available schema-1 root") + try: + native = _Original.model_validate(record.payload) + if native.object != record.identity.record_type: + raise ValueError("Stripe original kind differs from canonical identity") + if native.object == "balance": + if record.identity.native_id != "balance" or native.id is not None: + raise ValueError("Stripe balance identity differs from canonical identity") + return (_balance(record, native),) + if native.id != record.identity.native_id: + raise ValueError("Stripe original ID differs from canonical identity") + return (_entity(record, native),) + except (ValidationError, OverflowError, OSError): + raise ValueError("Stripe projection fields are malformed") from None + + +def _balance(record: SourceRecord, n: _Original) -> StripeBalanceEntity: + if n.available is None or n.pending is None or n.livemode is None: + raise ValueError("Stripe balance requires captured amounts and mode") + return StripeBalanceEntity( + entity_id="balance", + name="Account Balance", + breadcrumbs=[], + balance_id="balance", + balance_name="Account Balance", + snapshot_time=record.observed_at, + available=[item.model_dump() for item in n.available], + pending=[item.model_dump() for item in n.pending], + instant_available=( + [item.model_dump() for item in n.instant_available] + if n.instant_available is not None + else None + ), + connect_reserved=( + [item.model_dump() for item in n.connect_reserved] + if n.connect_reserved is not None + else None + ), + livemode=n.livemode, + ) + + +def _entity(record: SourceRecord, n: _Original) -> BaseEntity: # noqa: C901 + created = _time(n.created) + if created is None: + raise ValueError("Stripe original lacks a captured creation timestamp") + kind = n.object + name = n.name if kind == "customer" else n.number if kind == "invoice" else None + name = name or n.description or f"{kind.replace('_', ' ').title()} {n.id}" + common = { + "entity_id": n.id, + "name": name, + "breadcrumbs": [], + "created_at": created, + "updated_at": record.source_updated_at, + "created_time": created, + } + updated = {"updated_time": record.source_updated_at} + if kind == "balance_transaction": + return StripeBalanceTransactionEntity( + **common, + transaction_id=n.id, + transaction_name=name, + amount=n.amount, + currency=n.currency, + description=n.description, + fee=n.fee, + net=n.net, + reporting_category=n.reporting_category, + source=_id(n.source), + status=n.status, + type=n.type, + ) + if kind == "charge": + return StripeChargeEntity( + **common, + **updated, + charge_id=n.id, + charge_name=name, + amount=n.amount, + currency=n.currency, + description=n.description, + captured=n.captured, + paid=n.paid, + refunded=n.refunded, + customer_id=_id(n.customer), + invoice_id=_id(n.invoice), + ) + if kind == "customer": + return StripeCustomerEntity( + **common, + **updated, + customer_id=n.id, + customer_name=name, + email=n.email, + phone=n.phone, + description=n.description, + currency=n.currency, + delinquent=n.delinquent, + invoice_prefix=n.invoice_prefix, + ) + if kind == "event": + if n.data is None or n.livemode is None or n.type is None: + raise ValueError("Stripe event lacks captured type, mode or object identity") + return StripeEventEntity( + **common, + event_id=n.id, + event_name=n.type, + event_type=n.type, + api_version=n.api_version, + data=n.data.model_dump(exclude_none=True), + livemode=n.livemode, + pending_webhooks=n.pending_webhooks, + ) + if kind == "invoice": + return StripeInvoiceEntity( + **common, + **updated, + invoice_id=n.id, + invoice_name=name, + customer_id=_id(n.customer), + number=n.number, + status=n.status, + amount_due=n.amount_due, + amount_paid=n.amount_paid, + amount_remaining=n.amount_remaining, + due_date=_time(n.due_date), + paid=n.paid, + currency=n.currency, + ) + if kind == "payment_intent": + return StripePaymentIntentEntity( + **common, + **updated, + payment_intent_id=n.id, + payment_intent_name=name, + amount=n.amount, + currency=n.currency, + status=n.status, + description=n.description, + customer_id=_id(n.customer), + ) + if kind == "payment_method": + return StripePaymentMethodEntity( + **common, + payment_method_id=n.id, + payment_method_name=name, + type=n.type, + customer_id=_id(n.customer), + billing_details=n.billing_details.model_dump(exclude_none=True) + if n.billing_details + else {}, + ) + if kind == "payout": + return StripePayoutEntity( + **common, + **updated, + payout_id=n.id, + payout_name=name, + amount=n.amount, + currency=n.currency, + arrival_date=_time(n.arrival_date), + description=n.description, + destination=_id(n.destination), + method=n.method, + status=n.status, + ) + if kind == "refund": + return StripeRefundEntity( + **common, + refund_id=n.id, + refund_name=name, + amount=n.amount, + currency=n.currency, + status=n.status, + reason=n.reason, + charge_id=_id(n.charge), + payment_intent_id=_id(n.payment_intent), + ) + if kind == "subscription": + return StripeSubscriptionEntity( + **common, + **updated, + subscription_id=n.id, + subscription_name=name, + customer_id=_id(n.customer), + status=n.status, + current_period_start=_time(n.current_period_start), + current_period_end=_time(n.current_period_end), + cancel_at_period_end=n.cancel_at_period_end, + canceled_at=_time(n.canceled_at), + ) + raise ValueError("Unsupported Stripe original kind") diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py index a77988a1c3..35ba0da6cb 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py @@ -149,22 +149,43 @@ async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place assert projected.parts[-1].entity is None -async def test_factory_rejects_mode_conflict_before_auth_or_cursor_reset(database, source): +@pytest.mark.parametrize( + "provider,original_config", + [ + ("outlook_mail", {"capture_originals": True, "expected_principal_id": "mailbox"}), + ( + "stripe", + { + "original_capture": { + "expected_account_id": "acct_selected", + "livemode": False, + "api_version": "2025-06-30.basil", + } + }, + ), + ], +) +async def test_factory_rejects_mode_conflict_before_auth_or_cursor_reset( + database, source, provider, original_config +): """A full sync or skip-load cannot reinterpret a persisted Outlook cursor.""" from airweave.domains.sync_pipeline.config import SyncConfig from airweave.domains.sync_pipeline.factory import SyncFactory from airweave.models.sync_cursor import SyncCursor from airweave.platform.sources.outlook_mail import OutlookMailSource + from airweave.platform.sources.stripe import StripeSource _, fence = source factory = object.__new__(SyncFactory) factory._source_lifecycle_service = MagicMock(create=AsyncMock()) factory._source_registry = MagicMock() - factory._source_registry.get.return_value.source_class_ref = OutlookMailSource + factory._source_registry.get.return_value.source_class_ref = ( + OutlookMailSource if provider == "outlook_mail" else StripeSource + ) factory._build_arf_replay_source = AsyncMock(return_value="offline legacy replay") ctx = MagicMock() ctx.organization.id = fence.organization_id - sc = MagicMock(short_name="outlook_mail", config_fields={}) + sc = MagicMock(short_name=provider, config_fields={}) sync = MagicMock(id=fence.sync_id) job = MagicMock(id=fence.job_id) replay = SyncConfig.model_validate({"behavior": {"replay_from_arf": True}}) @@ -177,11 +198,12 @@ async def test_factory_rejects_mode_conflict_before_auth_or_cursor_reset(databas db.add(cursor) cursor.cursor_data = {} await db.flush() - assert await factory._build_source( - db, sync, job, ctx, MagicMock(), sc, False, replay - ) == "offline legacy replay" + assert ( + await factory._build_source(db, sync, job, ctx, MagicMock(), sc, False, replay) + == "offline legacy replay" + ) factory._build_arf_replay_source.reset_mock() - sc.config_fields = {"capture_originals": True, "expected_principal_id": "mailbox"} + sc.config_fields = original_config with pytest.raises(ValueError, match="not mutable ARF"): await factory._build_source(db, sync, job, ctx, MagicMock(), sc, False, replay) cursor.cursor_data = {"delta_link": "legacy continuation"} 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..34a195243a --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py @@ -0,0 +1,114 @@ +"""Stripe source composition and durable recovery; synthetic HTTP, real PostgreSQL.""" + +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 + 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/source_connections/tests/test_update.py b/backend/airweave/domains/source_connections/tests/test_update.py index 75fa6515cd..a7968652df 100644 --- a/backend/airweave/domains/source_connections/tests/test_update.py +++ b/backend/airweave/domains/source_connections/tests/test_update.py @@ -744,3 +744,34 @@ async def test_native_update_cannot_rebind_identity(): 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 c9db0bba1f..99ee86d2dc 100644 --- a/backend/airweave/domains/source_connections/update.py +++ b/backend/airweave/domains/source_connections/update.py @@ -31,7 +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 OutlookMailConfig +from airweave.platform.configs.config import OutlookMailConfig, StripeConfig from airweave.schemas.source_connection import ( AuthenticationMethod, ScheduleConfig, @@ -63,6 +63,19 @@ def _validate_outlook_capture_mode(source: SourceConnection, config: dict[str, A ) +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. @@ -138,6 +151,7 @@ async def update( 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"] diff --git a/backend/airweave/domains/sync_pipeline/factory.py b/backend/airweave/domains/sync_pipeline/factory.py index 437a86b5d9..692ce64ef3 100644 --- a/backend/airweave/domains/sync_pipeline/factory.py +++ b/backend/airweave/domains/sync_pipeline/factory.py @@ -72,7 +72,7 @@ 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 +from airweave.platform.configs.config import OutlookMailConfig, StripeConfig from airweave.platform.sources._base import BaseSource from .entity.pipeline import EntityPipeline @@ -400,15 +400,21 @@ async def _build_source( access_token: Optional[str] = None, ) -> SourceBuildResult: """Build source instance, cursor, file service, and node selections.""" - outlook_capture = False + composed_capture = False if source_connection.short_name == "outlook_mail": - outlook_capture = OutlookMailConfig.model_validate( + composed_capture = OutlookMailConfig.model_validate( source_connection.config_fields or {} ).capture_originals - await self._validate_outlook_capture_mode(db, sync.id, ctx, outlook_capture) + 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 outlook_capture or isinstance(source_class, (CanonicalSource, CanonicalPageSource)): + if composed_capture or isinstance(source_class, (CanonicalSource, CanonicalPageSource)): raise ValueError( "Canonical records must be reindexed from Postgres, not mutable ARF" ) @@ -447,10 +453,10 @@ async def _build_source( source=source, cursor=cursor, files=files, node_selections=node_selections ) - async def _validate_outlook_capture_mode( + 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 Outlook adapter.""" + """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( @@ -484,7 +490,7 @@ async def _validate_outlook_capture_mode( not capture_originals and (canonical or scanned or has_canonical_cursor) ): raise ValueError( - "Outlook capture mode differs from retained state; create a new source" + "Original capture mode differs from retained state; create a new source" ) async def _build_arf_replay_source( diff --git a/backend/airweave/domains/syncs/service.py b/backend/airweave/domains/syncs/service.py index c8041f3e71..f1c81813cb 100644 --- a/backend/airweave/domains/syncs/service.py +++ b/backend/airweave/domains/syncs/service.py @@ -59,7 +59,7 @@ TemporalScheduleServiceProtocol, TemporalWorkflowServiceProtocol, ) -from airweave.platform.configs.config import OutlookMailConfig +from airweave.platform.configs.config import OutlookMailConfig, StripeConfig from airweave.schemas.source_connection import ScheduleConfig from airweave.schemas.sync import SyncCreate from airweave.schemas.sync_job import SyncJobCreate @@ -138,6 +138,10 @@ async def create( if source_entry.short_name == "outlook_mail": # Outlook composes its canonical adapter only for explicit original capture. canonical = OutlookMailConfig.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, diff --git a/backend/airweave/domains/syncs/tests/test_service.py b/backend/airweave/domains/syncs/tests/test_service.py index 010af70462..7b7bb849db 100644 --- a/backend/airweave/domains/syncs/tests/test_service.py +++ b/backend/airweave/domains/syncs/tests/test_service.py @@ -1079,6 +1079,18 @@ async def test_create_no_cron_no_run_immediately(): ("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), ], ) @@ -1108,6 +1120,7 @@ async def test_create_with_cron_calls_temporal_schedule( 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 = { @@ -1115,6 +1128,7 @@ async def test_create_with_cron_calls_temporal_schedule( "gmail": GmailSource, "slack": SlackSource, "outlook_mail": OutlookMailSource, + "stripe": StripeSource, }[source_kind] result = await svc.create( AsyncMock(), diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index c1997c7093..78f15a03d4 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -4,7 +4,15 @@ from typing import Literal, Optional from uuid import UUID -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator +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 @@ -859,10 +867,20 @@ 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]+$") + + 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): 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/sources/stripe.py b/backend/airweave/platform/sources/stripe.py index ffa92f798e..0dcb8fa92c 100644 --- a/backend/airweave/platform/sources/stripe.py +++ b/backend/airweave/platform/sources/stripe.py @@ -20,14 +20,24 @@ import base64 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 @@ -55,6 +65,7 @@ ) 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 +91,8 @@ class StripeSource(BaseSource): including transactions, customers, subscriptions, and account analytics. """ + canonical_record_types = StripeCapture.canonical_record_types + @classmethod async def create( cls, @@ -91,12 +104,60 @@ 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) + response = await self.http_client.get( + "https://api.stripe.com" + path, + params=params, + headers={**headers, **auth_headers}, + timeout=20.0, + ) + if not response.is_success: + # Provider error messages can echo submitted private fields. Retain only + # status and retry timing when translating into existing typed errors. + 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, + ) + return response.json() + @retry( stop=stop_after_attempt(5), retry=retry_if_rate_limit_or_timeout, @@ -551,4 +612,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/tests/unit/domains/entities/test_stripe_projection.py b/backend/tests/unit/domains/entities/test_stripe_projection.py new file mode 100644 index 0000000000..0bf2b848fd --- /dev/null +++ b/backend/tests/unit/domains/entities/test_stripe_projection.py @@ -0,0 +1,277 @@ +"""Offline synthetic Stripe projection; no provider or live-money claims.""" + +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest + +from airweave.domains.entities.canonical.models import SourceRecord +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.stripe_projection import map_stripe +from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder +from airweave.platform.entities.stripe import ( + StripeEventEntity, + StripeInvoiceEntity, + StripePaymentIntentEntity, +) + + +def original(kind, **values): + native = "balance" if kind == "balance" else f"{kind}_123" + payload = {"object": kind, "livemode": False, **values} + if kind != "balance": + payload = {"id": native, "created": 1720000000, **payload} + return SourceRecord( + id=uuid4(), + sync_id=uuid4(), + identity=RecordIdentity(record_type=kind, native_id=native), + revision=1, + payload=payload, + payload_schema_version=1, + capture_hash="fixture", + content_hash=None, + completeness="partial", + observed_at=datetime(2026, 10, 1, tzinfo=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 test_invoice_payment_event_keep_separate_facts_and_ids(): + invoice = original( + "invoice", + number="INV-42", + amount_due=1200, + amount_paid=200, + amount_remaining=1000, + currency="jpy", + status="open", + customer={"id": "cus_1"}, + ) + payment = original( + "payment_intent", + amount=1200, + currency="jpy", + status="requires_action", + client_secret="pi_secret", + ) + event = original( + "event", + type="invoice.paid", + data={ + "object": { + "id": invoice.identity.native_id, + "object": "invoice", + "paid": True, + "amount_paid": 1200, + "metadata": {"secret": "nested_secret"}, + }, + "previous_attributes": {"client_secret": "previous_secret"}, + }, + ) + before = [item.model_dump() for item in (invoice, payment, event)] + inv, pi, evt = [map_stripe(item)[0] for item in (invoice, payment, event)] + assert isinstance(inv, StripeInvoiceEntity) and isinstance(pi, StripePaymentIntentEntity) + assert isinstance(evt, StripeEventEntity) + assert inv.entity_id == invoice.identity.native_id and inv.amount_remaining == 1000 + assert inv.amount_paid == 200 and inv.currency == "jpy" and inv.paid is None + assert pi.entity_id == payment.identity.native_id and pi.status == "requires_action" + assert evt.entity_id == event.identity.native_id and evt.event_type == "invoice.paid" + assert evt.data == {"object": {"id": invoice.identity.native_id, "object": "invoice"}} + assert all(entity.updated_at is None for entity in (inv, pi, evt)) + assert inv.updated_time is None and pi.updated_time is None + assert before == [item.model_dump() for item in (invoice, payment, event)] + + +def test_secrets_absent_from_both_full_entity_and_search_text(): + examples = [ + original( + "payment_intent", amount=500, client_secret="secret_a", metadata={"token": "secret_b"} + ), + original( + "payment_method", + type="card", + card={"number": "secret_c", "cvc": "secret_d", "fingerprint": "secret_e"}, + billing_details={"name": "Buyer", "email": "buyer@example.com", "token": "secret_f"}, + ), + original( + "balance", + available=[{"amount": 500, "currency": "usd", "token": "secret_g"}], + pending=[], + ), + original( + "event", + type="payment_intent.created", + data={ + "object": {"id": "pi_1", "object": "payment_intent", "client_secret": "secret_h"} + }, + ), + ] + builder = TextualRepresentationBuilder() + for record in examples: + entity = map_stripe(record)[0] + assert "secret_" not in entity.model_dump_json() + assert "secret_" not in builder.build_metadata_section(entity, "stripe") + assert "secret_" in str(record.payload) + + +@pytest.mark.parametrize( + "kind", + [ + "balance_transaction", + "charge", + "customer", + "invoice", + "payment_intent", + "payment_method", + "payout", + "refund", + "subscription", + ], +) +def test_resource_breadth_and_unknown_timestamps_flags(kind): + entity = map_stripe(original(kind))[0] + assert entity.entity_id == f"{kind}_123" + assert entity.created_at == datetime.fromtimestamp(1720000000, timezone.utc) + assert entity.updated_at is None + values = entity.model_dump() + for field in ( + "paid", + "captured", + "refunded", + "delinquent", + "cancel_at_period_end", + "updated_time", + ): + if field in values: + assert values[field] is None + + +@pytest.mark.parametrize( + "change", + [ + {"id": "other"}, + {"object": "charge"}, + {"created": "1720000000"}, + {"amount_paid": True}, + {"paid": "false"}, + ], +) +def test_identity_and_typed_field_mismatch_fail_without_payload_echo(change): + record = original("invoice", metadata={"token": "private_secret"}) + record = record.model_copy(update={"payload": {**record.payload, **change}}) + with pytest.raises(ValueError) as error: + map_stripe(record) + assert "private_secret" not in str(error.value) + + +def test_balance_uses_captured_snapshot_and_keeps_currency_buckets_separate(): + record = original( + "balance", + available=[{"amount": 500, "currency": "usd"}, {"amount": 1200, "currency": "jpy"}], + pending=[], + ) + entity = map_stripe(record)[0] + assert entity.snapshot_time == record.observed_at + assert entity.available == [ + {"amount": 500, "currency": "usd"}, + {"amount": 1200, "currency": "jpy"}, + ] + bad = record.model_copy( + update={"identity": RecordIdentity(record_type="balance", native_id="other")} + ) + with pytest.raises(ValueError, match="balance identity"): + map_stripe(bad) + + +def test_native_flags_references_and_financial_amounts_are_not_recomputed(): + charge = map_stripe( + original( + "charge", + amount=1000, + currency="usd", + paid=True, + captured=False, + refunded=False, + customer={"id": "cus_1", "metadata": {"token": "secret_expansion"}}, + ) + )[0] + assert charge.paid is True and charge.captured is False and charge.refunded is False + assert charge.customer_id == "cus_1" and "secret_expansion" not in charge.model_dump_json() + transaction = map_stripe( + original( + "balance_transaction", + amount=-1000, + fee=23, + net=-1023, + currency="usd", + source="ch_1", + ) + )[0] + assert (transaction.amount, transaction.fee, transaction.net) == (-1000, 23, -1023) + assert transaction.source == "ch_1" + refund = map_stripe( + original("refund", amount=200, charge="ch_1", payment_intent={"id": "pi_1"}) + )[0] + assert ( + refund.amount == 200 and refund.charge_id == "ch_1" and refund.payment_intent_id == "pi_1" + ) + + +def test_missing_creation_time_and_wrong_capture_scope_fail(): + record = original("invoice") + for bad in ( + record.model_copy(update={"payload": {**record.payload, "created": None}}), + record.model_copy(update={"payload_schema_version": 2}), + record.model_copy( + update={ + "identity": RecordIdentity( + record_type="invoice", + native_id=record.identity.native_id, + container_id="cus_1", + ) + } + ), + ): + with pytest.raises(ValueError): + map_stripe(bad) + + +def test_legacy_constructors_keep_explicit_flags_and_timestamps(): + from airweave.platform.entities.stripe import StripeChargeEntity, StripeSubscriptionEntity + + charge = StripeChargeEntity.from_api( + {"id": "ch_1", "created": 1720000000, "paid": True, "captured": False, "refunded": False}, + web_url=None, + ) + invoice = StripeInvoiceEntity.from_api( + {"id": "in_1", "created": 1720000000, "paid": True}, + web_url=None, + ) + subscription = StripeSubscriptionEntity.from_api( + {"id": "sub_1", "created": 1720000000, "cancel_at_period_end": False}, + web_url=None, + ) + assert charge.paid is True and charge.captured is False and charge.refunded is False + assert invoice.paid is True and subscription.cancel_at_period_end is False + assert charge.created_time == datetime(2024, 7, 3, 9, 46, 40) + + +@pytest.mark.asyncio +async def test_shared_projection_dispatch_uses_retained_original_without_storage_reads(): + from unittest.mock import AsyncMock + + from airweave.domains.entities.canonical.projection_mappers import map_record + + storage = AsyncMock() + record = original("invoice", number="INV-42", currency="usd", amount_due=1200) + async with map_record(record, "stripe", storage) as projection: + assert len(projection.parts) == 1 + assert isinstance(projection.parts[0].entity, StripeInvoiceEntity) + assert projection.parts[0].entity.invoice_id == record.identity.native_id + assert not storage.mock_calls diff --git a/backend/tests/unit/platform/sources/test_stripe_capture.py b/backend/tests/unit/platform/sources/test_stripe_capture.py new file mode 100644 index 0000000000..6cb9d54bd3 --- /dev/null +++ b/backend/tests/unit/platform/sources/test_stripe_capture.py @@ -0,0 +1,214 @@ +"""Stripe provider interpretation only; no live account or SQL qualification.""" + +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.platform.sources.stripe_capture import ( + StripeCapture, + StripeCaptureConfig, + StripeCaptureError, +) + + +def config(**overrides): + return StripeCaptureConfig( + **{ + "expected_account_id": "acct_selected", + "connected_account_id": "acct_selected", + "livemode": False, + "api_version": "2025-06-30.basil", + **overrides, + } + ) + + +def identity(): + return [ + {"object": "account", "id": "acct_selected"}, + {"object": "balance", "livemode": False, "available": []}, + ] + + +async def page(source, kind, continuation=None): + return await source.capture_page( + CompletedScope(record_type=kind), continuation or ScanContinuation(), files=MagicMock() + ) + + +@pytest.mark.asyncio +async def test_all_resource_families_preserve_native_json_and_context(): + read = AsyncMock(side_effect=identity()) + source = await StripeCapture.create(read, config()) + assert len(source.canonical_record_types) == 11 + assert all( + source.capture_cycle_configuration.policy(kind) == "discovery_only" + for kind in source.canonical_record_types + ) + for kind in source.canonical_record_types: + native = {"object": kind, "livemode": False, "unknown": {"nested": [1, None]}} + if kind != "balance": + native.update(id="obj_one", created=123) + if kind == "payment_method": + native.update(type="us_bank_account", us_bank_account={"last4": "6789"}) + read.side_effect = None + read.return_value = ( + native + if kind == "balance" + else { + "object": "list", + "data": [native], + "has_more": False, + "url": "https://untrusted.invalid/do-not-follow", + } + ) + result = await page(source, kind) + assert result.final and result.records[0].payload == native + assert result.records[0].completeness == "partial" + assert result.records[0].parent is None + assert result.records[0].kind == "upsert" + assert result.records[0].source_updated_at is None + if kind == "subscription": + assert read.call_args.kwargs["params"]["status"] == "all" + if kind == "payment_method": + assert "type" not in read.call_args.kwargs["params"] + assert result.records[0].payload["type"] == "us_bank_account" + for call in read.call_args_list: + assert call.kwargs["headers"] == { + "Stripe-Version": "2025-06-30.basil", + "Stripe-Account": "acct_selected", + } + assert call.args[0].startswith("/v1/") + + +@pytest.mark.asyncio +async def test_failed_page_retries_same_cursor_and_binding_cannot_change(): + raw = {"id": "cus_one", "object": "customer", "livemode": False} + read = AsyncMock(side_effect=identity() + [{"object": "list", "data": [raw], "has_more": True}]) + source = await StripeCapture.create(read, config()) + first = await page(source, "customer") + failure = RuntimeError("transport failed") + read.side_effect = failure + with pytest.raises(RuntimeError) as error: + await page(source, "customer", first.continuation) + assert error.value is failure + assert read.call_args.kwargs["params"]["starting_after"] == "cus_one" + read.side_effect = None + read.return_value = {"object": "list", "data": [], "has_more": False} + final = await page(source, "customer", first.continuation) + assert final.final and final.records == () + calls = read.await_count + with pytest.raises(StripeCaptureError, match="binding"): + await page(source, "event", first.continuation) + other = await StripeCapture.create( + AsyncMock(side_effect=identity()), config(api_version="2024-06-20") + ) + with pytest.raises(StripeCaptureError, match="binding"): + await page(other, "customer", first.continuation) + assert read.await_count == calls + + +@pytest.mark.asyncio +async def test_failed_reattest_resets_prior_identity_and_mode(): + read = AsyncMock(side_effect=identity()) + source = await StripeCapture.create(read, config()) + read.side_effect = None + read.return_value = {"object": "account", "id": "acct_other"} + with pytest.raises(StripeCaptureError, match="account"): + await source.verify_principal() + calls = read.await_count + with pytest.raises(StripeCaptureError, match="not attested"): + await page(source, "customer") + assert read.await_count == calls + read.side_effect = [identity()[0], {"object": "balance", "livemode": True}] + with pytest.raises(StripeCaptureError, match="mode"): + await source.verify_principal() + + +@pytest.mark.asyncio +async def test_malformed_page_never_becomes_empty_success(): + read = AsyncMock(side_effect=identity()) + source = await StripeCapture.create(read, config()) + read.side_effect = None + obj = {"id": "cus_one", "object": "customer", "livemode": False} + for response in ( + {"object": "list", "data": [], "has_more": True}, + {"object": "list", "data": [obj, obj], "has_more": False}, + {"object": "list", "data": [{**obj, "object": "event"}], "has_more": False}, + {"object": "list", "data": [{**obj, "livemode": True}], "has_more": False}, + {"object": "list", "data": [{**obj, "livemode": None}], "has_more": False}, + {"object": "list", "data": [], "has_more": "false"}, + {"object": "list", "has_more": False}, + ): + read.return_value = response + with pytest.raises(StripeCaptureError): + await page(source, "customer") + + +@pytest.mark.asyncio +async def test_event_snapshot_never_updates_embedded_object_or_confirms_absence(): + event = { + "id": "evt_one", + "object": "event", + "api_version": "2019-02-19", + "created": 123, + "livemode": False, + "type": "customer.deleted", + "data": {"object": {"id": "cus_deleted", "object": "customer", "deleted": True}}, + } + read = AsyncMock( + side_effect=identity() + [{"object": "list", "data": [event], "has_more": False}] + ) + source = await StripeCapture.create(read, config()) + result = await page(source, "event") + assert len(result.records) == 1 + assert result.records[0].identity.native_id == "evt_one" + assert result.records[0].payload == event + assert result.records[0].kind == "upsert" + with pytest.raises(StripeCaptureError, match="cannot confirm absence"): + await source.confirm_absent(MagicMock()) + + +@pytest.mark.asyncio +async def test_managed_source_wires_capture_without_exposing_credentials(): + import httpx + + from airweave.domains.sources.exceptions import SourceEntityForbiddenError, SourceError + from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider + from airweave.platform.configs.config import StripeConfig + from airweave.platform.sources.stripe import StripeSource + + auth = ManagedAuthProvider( + api_key="test-only", + connected_account_id="ca_selected", + allowed_hosts=frozenset({"api.stripe.com"}), + ) + client = AsyncMock() + client.get.side_effect = [httpx.Response(200, json=value) for value in identity()] + source = await StripeSource.create( + auth=auth, + logger=MagicMock(), + http_client=client, + config=StripeConfig(original_capture=config()), + ) + assert isinstance(source.capture_page_source, StripeCapture) + for call in client.get.call_args_list: + assert call.args[0].startswith("https://api.stripe.com/v1/") + assert call.kwargs["headers"] == { + "Stripe-Version": "2025-06-30.basil", + "Stripe-Account": "acct_selected", + } + client.get.side_effect = None + client.get.return_value = httpx.Response(403, json={"error": {"message": "private-secret"}}) + with pytest.raises(SourceEntityForbiddenError) as exc: + await page(source.capture_page_source, "customer") + assert "private-secret" not in str(exc.value) + with pytest.raises(SourceError, match="explicitly bound"): + await StripeSource.create( + auth=auth, + logger=MagicMock(), + http_client=client, + config=StripeConfig(), + ) From 07979dd2c5ad868332582a7cd68cdd448cc31fe2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:29:42 -0700 Subject: [PATCH 226/266] Bound Stripe capture response bodies through streamed reads --- .../canonical/tests/test_stripe_recovery.py | 8 +++ backend/airweave/platform/sources/stripe.py | 51 +++++++++++----- .../platform/sources/test_stripe_capture.py | 61 +++++++++++++++++++ 3 files changed, 104 insertions(+), 16 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py index 34a195243a..90ad830f82 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_stripe_recovery.py @@ -1,5 +1,6 @@ """Stripe source composition and durable recovery; synthetic HTTP, real PostgreSQL.""" +from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 @@ -49,6 +50,13 @@ async def get(url, **kwargs): 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", diff --git a/backend/airweave/platform/sources/stripe.py b/backend/airweave/platform/sources/stripe.py index 0dcb8fa92c..c172c41748 100644 --- a/backend/airweave/platform/sources/stripe.py +++ b/backend/airweave/platform/sources/stripe.py @@ -19,6 +19,8 @@ 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 @@ -92,6 +94,7 @@ class StripeSource(BaseSource): """ canonical_record_types = StripeCapture.canonical_record_types + MAX_CAPTURE_RESPONSE_BYTES = 16 * 1024 * 1024 @classmethod async def create( @@ -137,26 +140,42 @@ async def _read_original( auth_headers = {"Authorization": f"Bearer {credentials.api_key}"} else: auth_headers = await authorization_headers(self.auth) - response = await self.http_client.get( + async with self.http_client.stream( + "GET", "https://api.stripe.com" + path, params=params, headers={**headers, **auth_headers}, timeout=20.0, - ) - if not response.is_success: - # Provider error messages can echo submitted private fields. Retain only - # status and retry timing when translating into existing typed errors. - 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, - ) - return response.json() + 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), diff --git a/backend/tests/unit/platform/sources/test_stripe_capture.py b/backend/tests/unit/platform/sources/test_stripe_capture.py index 6cb9d54bd3..1f9d1a3349 100644 --- a/backend/tests/unit/platform/sources/test_stripe_capture.py +++ b/backend/tests/unit/platform/sources/test_stripe_capture.py @@ -1,5 +1,6 @@ """Stripe provider interpretation only; no live account or SQL qualification.""" +from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock import pytest @@ -187,6 +188,13 @@ async def test_managed_source_wires_capture_without_exposing_credentials(): ) client = AsyncMock() client.get.side_effect = [httpx.Response(200, json=value) for value in identity()] + + @asynccontextmanager + async def stream(method, url, **kwargs): + assert method == "GET" and kwargs["follow_redirects"] is False + yield await client.get(url, **kwargs) + + client.stream = MagicMock(side_effect=stream) source = await StripeSource.create( auth=auth, logger=MagicMock(), @@ -212,3 +220,56 @@ async def test_managed_source_wires_capture_without_exposing_credentials(): http_client=client, config=StripeConfig(), ) + + +@pytest.mark.parametrize("status", [200, 403]) +def test_stream_limit_and_error_status_close_response_without_reading_private_tail(status): + import asyncio + + # Own this loop so the standard runner also finalizes HTTPX's nested async + # generators before closing it; no timing sleep or warning suppression. + asyncio.run(_stream_failure_scenario(status)) + + +async def _stream_failure_scenario(status): + from uuid import uuid4 + + import httpx + + from airweave.domains.sources.exceptions import SourceError + from airweave.domains.sources.token_providers.protocol import ManagedAuthProvider + from airweave.platform.http_client.airweave_client import AirweaveHttpClient + from airweave.platform.sources.stripe import StripeSource + + class Body(httpx.AsyncByteStream): + reads = 0 + closed = False + + async def __aiter__(self): + for _ in range(3): + self.reads += 1 + yield b"x" * (64 * 1024) + raise AssertionError("Private tail must not be consumed") + + async def aclose(self): + self.closed = True + + body = Body() + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(status, stream=body)) + ) as client: + source = StripeSource( + auth=ManagedAuthProvider( + api_key="fixture", + connected_account_id="fixture", + allowed_hosts=frozenset({"api.stripe.com"}), + ), + logger=MagicMock(), + http_client=AirweaveHttpClient(client, uuid4(), "stripe"), + ) + source.MAX_CAPTURE_RESPONSE_BYTES = 64 * 1024 + with pytest.raises(SourceError) as exc: + await source._read_original("/v1/customers", params={}, headers={}) + assert body.reads == (2 if status == 200 else 0) + assert body.closed + assert "xxxxx" not in str(exc.value) From 31a18dd114ab5a89bad9c6d7c839d0f34e26c8a4 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:35:26 -0700 Subject: [PATCH 227/266] Provision Stripe captures without allowing reconnect scope changes --- .../domains/owned_provisioning/models.py | 17 ++- .../domains/owned_provisioning/store.py | 10 ++ .../tests/test_provisioning.py | 108 ++++++++++++++++++ backend/airweave/platform/configs/config.py | 9 ++ 4 files changed, 143 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index e79430178d..7455af8b8c 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -11,7 +11,7 @@ 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"] + provider: Literal["gmail", "google_calendar", "google_drive", "slack", "outlook_mail", "stripe"] 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) @@ -27,6 +27,8 @@ def native_principal(self): """Slack workspace membership requires its native user, not broker user_id.""" if (self.provider == "slack") != (self.expected_user_identity is not None): raise ValueError("Only Slack requires an expected native user identity") + if self.provider == "stripe": + self.source_config() # Validate mode/version before admission, not during delivery. return self @field_validator("cron") @@ -39,6 +41,13 @@ def valid_schedule(cls, value: str) -> str: def source_config(self) -> dict[str, JsonValue]: """Expected identity cannot be overridden inside unstructured provider config.""" + 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 == "outlook_mail": return { **self.config, @@ -108,6 +117,7 @@ def native_principal(provider: str, config: dict) -> tuple[str, str | None]: GoogleDriveConfig, OutlookMailConfig, SlackConfig, + StripeConfig, ) user = None @@ -123,6 +133,11 @@ def native_principal(provider: str, config: dict) -> tuple[str, str | None]: 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 diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py index dbbacb0fcb..b71a6db268 100644 --- a/backend/airweave/domains/owned_provisioning/store.py +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -20,6 +20,7 @@ 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, @@ -210,6 +211,15 @@ async def _configure( ) .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", + ) source.config_fields = config source.readable_auth_provider_id = spec.auth_provider source.auth_provider_config = spec.auth_config() diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index 1e201e5900..a373db3263 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -569,3 +569,111 @@ async def fail_after_admission(db, obj_in, ctx, uow=None): 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, + ) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 78f15a03d4..caf384cf1b 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -876,6 +876,15 @@ class StripeCaptureConfig(BaseModel): 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): """Legacy entity ingestion, or explicitly bound original capture on a fresh source.""" From 06b9a3b22515d2e9e5fabc81acd3ab3a893449e6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:48:04 -0700 Subject: [PATCH 228/266] Rerank validated owned-source candidates through the existing model port --- backend/airweave/core/container/factory.py | 4 +- .../canonical/tests/test_owned_ranking.py | 111 ++++++++++ backend/airweave/domains/search/owned.py | 134 +++++++++++- .../airweave/domains/search/owned_models.py | 20 ++ .../search/tests/test_owned_ranking.py | 196 ++++++++++++++++++ 5 files changed, 453 insertions(+), 12 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_owned_ranking.py create mode 100644 backend/airweave/domains/search/tests/test_owned_ranking.py diff --git a/backend/airweave/core/container/factory.py b/backend/airweave/core/container/factory.py index 7eec94001b..d89b73fe7f 100644 --- a/backend/airweave/core/container/factory.py +++ b/backend/airweave/core/container/factory.py @@ -1406,7 +1406,9 @@ def _create_search_services( ) return { - "owned_search": OwnedSearchService(executor, source_registry), + "owned_search": OwnedSearchService( + executor, source_registry, reranker=reranker, tokenizer=tokenizer + ), "instant_search": instant_search, "classic_search": classic_search, "agentic_search": agentic_search, diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_ranking.py b/backend/airweave/domains/entities/canonical/tests/test_owned_ranking.py new file mode 100644 index 0000000000..ff2dc9758f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_ranking.py @@ -0,0 +1,111 @@ +"""Real SQL disclosure and final publication fencing around a fake reranker.""" + +from unittest.mock import AsyncMock + +import pytest +from sqlalchemy import event + +from airweave.api import deps +from airweave.core.protocols.reranker import RerankerResult +from airweave.domains.entities.canonical.tests.helpers import capture, observation +from airweave.domains.entities.canonical.tests.test_owned_search import ( # noqa: F401 + http_search, # noqa: F811 + indexed, # noqa: F811 +) +from airweave.domains.entities.canonical.tests.test_search_visibility import hit +from airweave.domains.search.types import SearchResults + + +class QuarterCharacterTokenizer: + """Fake count sufficient to expose full-chunk versus display-excerpt behavior.""" + + def count_tokens(self, text): + return (len(text) + 3) // 4 + + +@pytest.mark.parametrize("edit_stage", [None, "before", "during", "before_revoke", "during_revoke"]) +async def test_retained_match_disclosure_and_final_gate( + database, + source, + indexed, # noqa: F811 + http_search, # noqa: F811 + edit_stage, +): + fence, locator, _ = indexed + client, vector, _, _, _ = http_search + capture_service, _ = source + candidate = hit(fence, locator.encode()).model_copy( + update={ + "textual_representation": "x" * 2500 + "MATCH_AFTER_DISPLAY_BOUNDARY", + "raw_source_fields": {"secret": "RAW_MUST_NOT_TRANSFER"}, + } + ) + vector.seed_results(SearchResults(results=[candidate])) + service = client._transport.app.dependency_overrides[deps.get_container]().owned_search + service._tokenizer = QuarterCharacterTokenizer() + checked_out = 0 + + def checkout(*args): + nonlocal checked_out + checked_out += 1 + + def checkin(*args): + nonlocal checked_out + checked_out -= 1 + + async with database() as db: + engine = db.bind.sync_engine + event.listen(engine, "checkout", checkout) + event.listen(engine, "checkin", checkin) + + async def edit(): + changed = ( + observation(kind="delete", removal_reason="access_revoked") + if edit_stage and edit_stage.endswith("revoke") + else observation(payload={"changed": True}) + ) + await capture(database, capture_service, fence, changed) + + if edit_stage in {"before", "before_revoke"}: + enrich = service._enrich + + async def enrich_then_edit(*args): + result = await enrich(*args) + await edit() + return result + + service._enrich = enrich_then_edit + + async def rerank(query, documents, top_n): + assert checked_out == 0 + assert "MATCH_AFTER_DISPLAY_BOUNDARY" in documents[0] + assert "RAW_MUST_NOT_TRANSFER" not in documents[0] + if edit_stage in {"during", "during_revoke"}: + await edit() + return [RerankerResult(index=0, relevance_score=0.9)] + + service._reranker = type("Port", (), {})() + service._reranker.rerank = AsyncMock(side_effect=rerank) + try: + response = await client.post( + "/sync/search", + json={ + "query": "budget", + "sync_ids": [str(fence.sync_id)], + "mode": "keyword", + }, + ) + finally: + event.remove(engine, "checkout", checkout) + event.remove(engine, "checkin", checkin) + assert response.status_code == 200, response.text + body = response.json() + assert len(body["items"]) == int(edit_stage is None) + assert service._reranker.rerank.await_count == int( + edit_stage not in {"before", "before_revoke"} + ) + if edit_stage is None: + assert body["ranking"]["method"] == "shared_rerank" + assert "MATCH_AFTER_DISPLAY_BOUNDARY" not in body["items"][0]["excerpts"][0] + else: + assert body["excluded_candidates"] == 1 diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index a739fc561f..7204dbcb82 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -1,5 +1,7 @@ """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 @@ -10,6 +12,8 @@ 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 @@ -22,6 +26,7 @@ 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, OwnedSearchHit, OwnedSearchRequest, @@ -70,12 +75,21 @@ class _EnrichmentRecord(BaseModel): class OwnedSearchService: - """No provider lifecycle, planner, answer generator or reranker dependency.""" + """Canonical retrieval and optional shared ranking, without source provider calls.""" - def __init__(self, executor: SearchPlanExecutorProtocol, registry: SourceRegistryProtocol): + 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, @@ -98,7 +112,7 @@ async def search( }, ) scope_snapshot = self._scope_identity(scopes) - hits, scores, exclusions, postfiltered = {}, {}, 0, 0 + hits, scores, matched_text, exclusions, postfiltered = {}, {}, {}, 0, 0 engine_partial, full = False, False plan = SearchPlan( query=SearchQuery(primary=request.query), @@ -121,7 +135,7 @@ async def search( collection_readable_id=readable_id, indexed_only=True, ) - group_hits, group_scores, rejected, filtered = await self._enrich( + group_hits, group_scores, group_text, rejected, filtered = await self._enrich( db, ctx, request, sync_ids, scopes, results ) engine_partial |= results.engine_partial @@ -129,24 +143,39 @@ async def search( 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: - 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") eligible = await self._final_publications(db, ctx, request, scores, scope_snapshot) exclusions += len(hits.keys() - eligible) - ranked = sorted(hits.keys() & eligible, key=lambda key: (-scores[key][0], str(key))) + 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] return OwnedSearchResponse( items=tuple(hits[key] for key in ranked[: request.limit]), + ranking=ranking, sources=sources, candidate_window_full=full, engine_partial=engine_partial, excluded_candidates=exclusions, postfilter_excluded=postfiltered, - retrieval_incomplete=engine_partial + retrieval_incomplete=ranking.shortlist_truncated + or engine_partial or full or exclusions > 0 or postfiltered > 0 @@ -160,6 +189,82 @@ async def search( ), ) + 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( @@ -276,11 +381,17 @@ async def _enrich( sync_ids: list[UUID], scopes: dict[UUID, _SourceScope], results: SearchResults, - ) -> tuple[dict[UUID, OwnedSearchHit], dict[UUID, tuple[float, ProjectionLocator]], int, int]: + ) -> 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, exclusions, postfiltered = {}, {}, 0, 0 + 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( @@ -376,11 +487,12 @@ async def _enrich( ), ) scores[row.id] = (1 / (60 + rank), locator) + matched_text[row.id] = result.textual_representation 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, exclusions, postfiltered + return hits, scores, matched_text, exclusions, postfiltered @staticmethod def _scope_identity(scopes: dict[UUID, _SourceScope]) -> set[tuple]: diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py index 62ce497f61..8851a36cc6 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -72,10 +72,30 @@ class OwnedSearchCoverage(BaseModel): 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" 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..d053192ccb --- /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") 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" From 4861594a8489485cb83e72eb07aec0988fbfd627 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 08:59:27 -0700 Subject: [PATCH 229/266] Group eligible conversation matches while preserving exact originals --- .../canonical/tests/test_owned_grouping.py | 125 ++++++++++++++++++ backend/airweave/domains/search/owned.py | 89 ++++++++++++- .../airweave/domains/search/owned_models.py | 17 ++- .../search/tests/test_owned_grouping.py | 103 +++++++++++++++ .../search/tests/test_owned_ranking.py | 2 +- 5 files changed, 332 insertions(+), 4 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py create mode 100644 backend/airweave/domains/search/tests/test_owned_grouping.py diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py b/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py new file mode 100644 index 0000000000..65b1d5f562 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py @@ -0,0 +1,125 @@ +"""Real SQL final publication checks precede grouping and result limits.""" + +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from sqlalchemy import select, update + +from airweave.api import deps +from airweave.core.protocols.reranker import RerankerResult +from airweave.domains.entities.canonical.projection_models import ProjectionLocator +from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared +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.collection import Collection +from airweave.models.source_connection import SourceConnection + + +@pytest.mark.parametrize("withdraw", [False, True, "parent"]) +async def test_group_after_final_gate_promotes_survivor_and_limits_cards( + database, + source, + indexed, # noqa: F811 + http_search, # noqa: F811 + withdraw, +): + fence, _, connection = indexed + capture_service, _ = source + client, vector, _, _, _ = http_search + messages = [ + observation( + identity=RecordIdentity(record_type="message", native_id=f"m{i}"), + payload={"threadId": "shared" if i < 3 else "other"}, + ) + for i in range(4) + ] + if withdraw == "parent": + parent = RecordIdentity(record_type="session", native_id="session") + await capture(database, capture_service, fence, observation(identity=parent)) + messages = [ + message.model_copy( + update={ + "identity": message.identity.model_copy( + update={"container_id": parent.native_id} + ), + "parent": parent, + } + ) + for message in messages + ] + async with database() as db: + await db.execute( + update(SourceConnection) + .where(SourceConnection.id == connection.id) + .values(short_name="almanac") + ) + await db.commit() + await capture(database, capture_service, fence, *messages) + store = CanonicalProjectionStore() + async with database() as db: + collection_id = await db.scalar( + select(Collection.id).where(Collection.readable_id == connection.readable_collection_id) + ) + pending = await store.pending(db, fence.organization_id, fence.sync_id) + candidates = [] + for work in sorted(pending, key=lambda item: item.record.identity.native_id): + 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()) + if withdraw == "parent": + candidate.airweave_system_metadata.source_name = "almanac" + candidate.name = work.record.identity.native_id + candidate.textual_representation = "exact " + candidate.name + candidates.append(candidate) + vector.seed_results(SearchResults(results=candidates)) + service = client._transport.app.dependency_overrides[deps.get_container]().owned_search + service._tokenizer = type("Tokenizer", (), {"count_tokens": lambda self, text: len(text)})() + + async def rerank(query, documents, top_n): + if withdraw: + await capture( + database, + capture_service, + fence, + observation( + identity=parent if withdraw == "parent" else messages[0].identity, + kind="delete", + removal_reason="access_revoked", + ), + ) + return [RerankerResult(i, float(top_n - i)) for i in range(top_n)] + + service._reranker = type("Reranker", (), {"rerank": AsyncMock(side_effect=rerank)})() + response = await client.post( + "/sync/search", + json={"query": "exact", "sync_ids": [str(fence.sync_id)], "mode": "keyword", "limit": 2}, + ) + assert response.status_code == 200, response.text + items = response.json()["items"] + if withdraw == "parent": + assert items == [] + return + assert len(items) == 2 + first = items[0] + assert first["identity"]["native_id"] == ("m1" if withdraw else "m0") + assert first["title"] == first["identity"]["native_id"] + assert first["excerpts"] == ["exact " + first["title"]] + assert first["group"]["matched_records"] == (2 if withdraw else 3) + assert first["group"]["native_id"] == "shared" + assert items[1]["group"]["native_id"] == "other" + assert all(m["identity"]["native_id"] != "m0" for m in first["group"]["additional_matches"]) diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 7204dbcb82..3d7cfc057f 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -28,7 +28,9 @@ from airweave.domains.search.owned_models import ( OwnedRanking, OwnedSearchCoverage, + OwnedSearchGroup, OwnedSearchHit, + OwnedSearchMatch, OwnedSearchRequest, OwnedSearchResponse, ) @@ -67,6 +69,9 @@ class _EnrichmentRecord(BaseModel): 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 @@ -166,8 +171,9 @@ async def search( 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(hits[key] for key in ranked[: request.limit]), + items=tuple(grouped[: request.limit]), ranking=ranking, sources=sources, candidate_window_full=full, @@ -179,7 +185,7 @@ async def search( or full or exclusions > 0 or postfiltered > 0 - or len(ranked) > request.limit + or len(grouped) > request.limit or any( row.pending_records or row.partially_indexed_records @@ -189,6 +195,81 @@ async def search( ), ) + @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"})) + 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, @@ -404,6 +485,9 @@ async def _enrich( 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, @@ -486,6 +570,7 @@ async def _enrich( 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 excerpt = result.textual_representation[:2000] diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py index 8851a36cc6..817b060a65 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -41,7 +41,7 @@ def validate_scope(self): return self -class OwnedSearchHit(BaseModel): +class OwnedSearchMatch(BaseModel): """Canonical identity plus excerpts; no native payload or generation IDs.""" record_id: UUID @@ -60,6 +60,21 @@ class OwnedSearchHit(BaseModel): 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.""" 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..f888fb1274 --- /dev/null +++ b/backend/airweave/domains/search/tests/test_owned_grouping.py @@ -0,0 +1,103 @@ +"""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.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, 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() + + +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 index d053192ccb..bce98eae4c 100644 --- a/backend/airweave/domains/search/tests/test_owned_ranking.py +++ b/backend/airweave/domains/search/tests/test_owned_ranking.py @@ -103,7 +103,7 @@ async def test_global_orchestration_rechecks_before_and_after_without_holding_se keys = [UUID(int=1), UUID(int=2)] syncs = [uuid4(), uuid4()] collections = [uuid4(), uuid4()] - hits = {key: SimpleNamespace(title="Title") for key in keys} + hits = {key: SimpleNamespace(title="Title", group=None) for key in keys} active = 0 @asynccontextmanager From 908d2ab47519c33d7ffe01ebee0e9eca4364b65f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:07:16 -0700 Subject: [PATCH 230/266] Preserve extraction evidence when converting grouped matches --- backend/airweave/domains/search/owned.py | 4 +++- .../domains/search/tests/test_owned_grouping.py | 10 ++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 3d7cfc057f..882b51de1c 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -253,7 +253,9 @@ def _group_hits(ranked: list[OwnedSearchHit]) -> list[OwnedSearchHit]: assert previous is not None additional = previous.additional_matches if len(additional) < 3: - exact = OwnedSearchMatch.model_validate(hit.model_dump(exclude={"group"})) + exact = OwnedSearchMatch.model_validate( + hit.model_dump(exclude={"group"}, round_trip=True) + ) additional = ( *additional, exact.model_copy(update={"excerpts": exact.excerpts[:1]}), diff --git a/backend/airweave/domains/search/tests/test_owned_grouping.py b/backend/airweave/domains/search/tests/test_owned_grouping.py index f888fb1274..69cc4e7a27 100644 --- a/backend/airweave/domains/search/tests/test_owned_grouping.py +++ b/backend/airweave/domains/search/tests/test_owned_grouping.py @@ -6,6 +6,10 @@ 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 @@ -54,6 +58,10 @@ 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:]]) @@ -70,6 +78,8 @@ def test_source_scoped_group_keeps_exact_anchors_and_observed_count(): ) 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(): From b7484e289eece009bbe83efadabc96023e280381 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:24:23 -0700 Subject: [PATCH 231/266] test: cover scoped retrieval across all Vespa schemas --- backend/tests/integration/test_owned_vespa.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index ff746f7127..0824db1381 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -5,6 +5,7 @@ """ import json +import math import os from datetime import datetime, timezone from uuid import uuid4 @@ -18,6 +19,7 @@ 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 @@ -122,6 +124,99 @@ async def test_real_schema_retrieval_collection_isolation_and_delete(): 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, + ) + 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" From 86414dc7d398091b9e9bbb516e8278ad5174da2a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:28:38 -0700 Subject: [PATCH 232/266] Preserve scoped semantic candidates with standard filtered HNSW traversal --- .../domains/search/adapters/vector_db/vespa_client.py | 4 ++++ backend/tests/integration/test_owned_vespa.py | 8 ++++++++ 2 files changed, 12 insertions(+) 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 7bc6840f50..ed58eda12a 100644 --- a/backend/airweave/domains/search/adapters/vector_db/vespa_client.py +++ b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py @@ -384,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} diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index 0824db1381..c2491e6960 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -202,6 +202,14 @@ async def test_all_schemas_retrieve_varied_vectors_with_collection_and_sync_scop ), 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 From d4f2a6853e3af31252d577da7029a61b69cafd87 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:37:13 -0700 Subject: [PATCH 233/266] Distinguish conversation ranking from displayed-original relevance --- evaluation/README.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/evaluation/README.md b/evaluation/README.md index fc15bbb5bf..85a74f5768 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -191,3 +191,25 @@ payloads, keys, file paths or destination URLs. Exit codes: 0 completed capture; or inconsistent import progress (including cross-record preflight rejection); 130 interrupted; 1 unexpected failure. Failed commands may have committed earlier pages: a nonzero exit is not evidence of rollback. + +## Grouped conversation results + +Judge the unit the user actually sees. A conversation card can be relevant even +when its displayed message is not the answer. Keep these two experiments separate: + +- **Card ranking:** stable conversation identity for a session/thread card, and + original record identity for an ungrouped result. Judge the conversation as a + whole. Use a separately fingerprinted dataset tagged `unit:card`. +- **Displayed-original ranking:** representative original followed by the explicitly + returned additional matches, preserving card order and within-card order. Judge + those exact original identities in a separate dataset tagged `unit:original`. + Do not add unreturned members merely because they belong to a returned group. + This is a ranking of inspectable originals, not the number of visible cards. + +Both use the existing evaluator; no second scoring implementation is needed. +A card hit does not prove passage recall, and `matched_records` is not a list of +answers or evidence that every member is relevant. Keep query filters, corpus +revisions, model settings and assessment provenance fixed for both experiments. +The one-session trial in the product worklog is diagnostic, not a judged benchmark: +it found the session in every mode, but the expected passage only in keyword and +hybrid mode. More diverse sessions and independently judged queries are required. From 20774dfa6363170de59c2984d614e307ed6eb8d2 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:50:58 -0700 Subject: [PATCH 234/266] Capture Outlook folder deltas with mailbox-owned recovery checkpoints --- .../canonical/tests/test_outlook_recovery.py | 134 +++++++- .../platform/sources/outlook_graph.py | 33 ++ .../platform/sources/outlook_mail_capture.py | 297 ++++++++++++++---- .../platform/sources/outlook_mail_models.py | 79 ++++- .../sources/tests/test_outlook_identity.py | 27 ++ .../tests/test_outlook_mail_capture.py | 240 +++++++++++++- 6 files changed, 708 insertions(+), 102 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py index 35ba0da6cb..424aef19f5 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_recovery.py @@ -9,13 +9,15 @@ from airweave.adapters.storage.filesystem import FilesystemBackend from airweave.domains.entities.canonical.blob_materializer import read_blob -from airweave.domains.entities.canonical.cycle_models import CompleteCycle +from airweave.domains.entities.canonical.cycle_models import CompleteCycle, TerminalCheckpoint from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.requests import CompletedScope from airweave.domains.sources.token_providers.static import StaticTokenProvider from airweave.domains.storage.file_service import FileService from airweave.domains.sync_pipeline.canonical_scan import CanonicalScanDriver from airweave.models.capture_scan import CaptureScan from airweave.models.entity import Entity +from airweave.models.sync_job import SyncJob from airweave.platform.configs.config import OutlookMailConfig from airweave.platform.sources.outlook_graph import OutlookGraphClient from airweave.platform.sources.outlook_mail_capture import OutlookMailCapture @@ -34,6 +36,21 @@ def __init__(self): self.fail_next_page = True self.calls = [] self.moved = False + self.incremental = False + self.fail_incremental = True + + def incremental_page(self, request): + """One new delta round interrupted after the first folder.""" + assert request.url.params.get("$deltatoken") == "one" + if "/sent/" in request.url.path and self.fail_incremental: + raise httpx.ConnectError("Synthetic incremental interruption", request=request) + return httpx.Response( + 200, + json={ + "value": [{"id": "a", "@removed": {"reason": "deleted"}}], + "@odata.deltaLink": str(request.url.copy_with(query=b"$deltatoken=two")), + }, + ) async def handle(self, request): path = request.url.path @@ -41,17 +58,38 @@ async def handle(self, request): return httpx.Response(200, json={"id": "mailbox"}) assert request.headers["Prefer"] == 'IdType="ImmutableId"' self.calls.append((path, str(request.url.query))) - if path == "/v1.0/me/messages": + if path == "/v1.0/me/mailFolders": + return httpx.Response(200, json={"value": [{"id": "inbox"}, {"id": "sent"}]}) + if path.endswith("/childFolders"): + return httpx.Response(200, json={"value": []}) + if self.incremental and path.endswith("/messages/delta"): + return self.incremental_page(request) + if path == "/v1.0/me/mailFolders/sent/messages/delta": + return httpx.Response( + 200, + json={ + "value": [], + "@odata.deltaLink": BASE + "/mailFolders/sent/messages/delta?$deltatoken=one", + }, + ) + if path == "/v1.0/me/mailFolders/inbox/messages/delta": if request.url.params.get("$skiptoken") == "second": if self.fail_next_page: raise httpx.ConnectError("Synthetic page interruption", request=request) self.moved = True - return httpx.Response(200, json={"value": [{"id": "a"}, {"id": "b"}]}) + return httpx.Response( + 200, + json={ + "value": [{"id": "a"}, {"id": "b"}], + "@odata.deltaLink": BASE + + "/mailFolders/inbox/messages/delta?$deltatoken=one", + }, + ) return httpx.Response( 200, json={ "value": [{"id": "a"}], - "@odata.nextLink": BASE + "/messages?$skiptoken=second", + "@odata.nextLink": BASE + "/mailFolders/inbox/messages/delta?$skiptoken=second", }, ) if path.endswith("/$value"): @@ -62,8 +100,12 @@ async def handle(self, request): 200, json={ "id": identity, - "changeKey": "moved" if self.moved and identity == "a" else "original", - "parentFolderId": "sent" if self.moved and identity == "a" else "inbox", + "changeKey": "returned" + if self.incremental + else ("moved" if self.moved and identity == "a" else "original"), + "parentFolderId": "sent" + if self.moved and identity == "a" and not self.incremental + else "inbox", "subject": "Original", "body": {"contentType": "text", "content": "Original retained bytes"}, "createdDateTime": "2026-09-01T00:00:00Z", @@ -93,6 +135,24 @@ async def run_capture(database, service, fence, mailbox, storage): await files.cleanup_sync_directory(MagicMock()) +async def complete(database, service, fence, cycle): + """Match production terminal publication using the exact committed root scan.""" + async with database() as db: + scan = await service.read_scan(db, fence, CompletedScope(record_type="message")) + return await service.complete_cycle( + db, + CompleteCycle( + fence=fence, + expected=cycle.version, + terminal_checkpoint=TerminalCheckpoint( + record_type="message", + expected=scan.version, + checkpoint=scan.provider_checkpoint, + ), + ), + ) + + async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place( database, source, tmp_path ): @@ -120,15 +180,16 @@ async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place mailbox.fail_next_page = False split = len(mailbox.calls) after = await run_capture(database, service, renewed, mailbox, storage) - async with database() as db: - after = await service.complete_cycle( - db, CompleteCycle(fence=renewed, expected=after.version) - ) + after = await complete(database, service, renewed, after) assert after.phase == "complete" assert after.last_full_capture.discovery == "incomplete" - assert after.promoted_checkpoint is None + assert after.promoted_checkpoint is not None assert after.version.cycle_id == before.version.cycle_id - listing = [query for path, query in mailbox.calls[split:] if path == "/v1.0/me/messages"] + listing = [ + query + for path, query in mailbox.calls[split:] + if path == "/v1.0/me/mailFolders/inbox/messages/delta" + ] assert len(listing) == 1 and "second" in listing[0] async with database() as db: rows = list((await db.scalars(select(Entity))).all()) @@ -141,6 +202,55 @@ async def test_resumed_mailbox_page_preserves_original_and_updates_move_in_place assert retained.parent is None and retained.identity.container_id is None assert retained.completeness == "partial" and len(retained.blobs) == 1 assert await read_blob(retained, retained.blobs[0], storage) == MIME + # A new job consumes the promoted folder tokens, never re-enumerates messages. + async with database() as db: + previous_job = await db.get(SyncJob, renewed.job_id) + previous_job.status = "completed" + job_id = uuid4() + db.add( + SyncJob( + id=job_id, + organization_id=fence.organization_id, + sync_id=fence.sync_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, + ) + mailbox.incremental = True + with pytest.raises(httpx.RequestError): + await run_capture(database, service, next_fence, mailbox, storage) + async with database() as db: + partial = await service.read_cycle(db, next_fence) + assert partial.mode == "changes" and partial.phase == "active" + assert partial.promoted_checkpoint == after.promoted_checkpoint + rows = list((await db.scalars(select(Entity))).all()) + assert len(rows) == 2 and all(row.deleted_at is None for row in rows) + mailbox.fail_incremental = False + split = len(mailbox.calls) + changed = await run_capture(database, service, next_fence, mailbox, storage) + changed = await complete(database, service, next_fence, changed) + async with database() as db: + rows = list((await db.scalars(select(Entity))).all()) + moved = next(row for row in rows if row.native_id == "a") + assert moved.id == initial_id and moved.deleted_at is None + assert moved.source_payload["parentFolderId"] == "inbox" + retained = await service.store.read(db, fence.organization_id, fence.sync_id, moved.id) + assert await read_blob(retained, retained.blobs[0], storage) == MIME + assert changed.mode == "changes" + assert all( + "two" in value + for value in changed.promoted_checkpoint.checkpoint.value["folder_links"].values() + ) + delta_requests = [path for path, _ in mailbox.calls[split:] if path.endswith("/delta")] + assert delta_requests == ["/v1.0/me/mailFolders/sent/messages/delta"] # Actual mapper dispatch consumes retained bytes after the provider client has closed. async with map_record(retained, "outlook_mail", storage) as projected: assert len(projected.entities) == 1 diff --git a/backend/airweave/platform/sources/outlook_graph.py b/backend/airweave/platform/sources/outlook_graph.py index 6f9931ba40..a24091ca51 100644 --- a/backend/airweave/platform/sources/outlook_graph.py +++ b/backend/airweave/platform/sources/outlook_graph.py @@ -22,6 +22,24 @@ 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.""" @@ -78,6 +96,8 @@ async def _request( ) -> 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: @@ -143,6 +163,19 @@ async def get( 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()) diff --git a/backend/airweave/platform/sources/outlook_mail_capture.py b/backend/airweave/platform/sources/outlook_mail_capture.py index dc800985bc..f6b0b27cd9 100644 --- a/backend/airweave/platform/sources/outlook_mail_capture.py +++ b/backend/airweave/platform/sources/outlook_mail_capture.py @@ -1,30 +1,46 @@ -"""Opt-in mailbox-owned full Outlook capture; no legacy cursor or activation changes.""" +"""Mailbox-owned Outlook originals with bounded per-folder delta progress.""" import hashlib import json from datetime import datetime, timezone -from urllib.parse import quote, urlsplit +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.cycle_models import ( + CaptureCycle, + CycleConfiguration, + ProviderCheckpoint, +) from airweave.domains.entities.canonical.models import SourceRecord -from airweave.domains.entities.canonical.page_source import CapturePage, CapturePlan +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, OutlookGraphClient +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, - OutlookMessagePage, ) BASE = "https://graph.microsoft.com/v1.0/me" @@ -34,7 +50,7 @@ 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. Incremental folder deltas remain a required, unimplemented follow-up. + qualified. Folder removals are signals, never proof of mailbox deletion. """ canonical_record_types = ("message",) @@ -60,7 +76,7 @@ def __init__( self.fingerprint = hashlib.sha256( json.dumps( { - "version": 1, + "version": 2, "id_type": "ImmutableId", "config": self.config.model_dump(mode="json"), "included_ids": sorted(included_ids), @@ -121,45 +137,217 @@ def capture_cycle_configuration(self) -> CycleConfiguration: ) async def prepare_cycle(self, previous: CaptureCycle | None) -> CapturePlan: - """Every cycle is explicitly full; no delta completion is implied.""" + """Only a compatible promoted mailbox checkpoint authorizes changes mode.""" await self.graph.verify_principal() self._require_principal() - return CapturePlan() + 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 a fresh scan to its persisted principal, selection and ID format.""" + """Bind progress to the persisted principal, selection and immutable ID format.""" self._require_principal() if cycle.configuration != self.capture_cycle_configuration: - raise OutlookBoundaryError( - "Outlook capture configuration changed", source_short_name="outlook_mail" + 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 ) - return self._continuation(OutlookMailContinuation(fingerprint=self.fingerprint)) + ) @staticmethod - def _continuation(state: OutlookMailContinuation) -> ScanContinuation: - return ScanContinuation(value=state.model_dump(mode="json")) + 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 OutlookBoundaryError( - "Invalid Outlook capture continuation", source_short_name="outlook_mail" - ) from None + raise self._error("Invalid Outlook capture continuation") from None if state.fingerprint != self.fingerprint: - raise OutlookBoundaryError( - "Outlook continuation belongs to another capture scope", - source_short_name="outlook_mail", - ) + raise self._error("Outlook continuation belongs to another capture scope") return state - def _list_url(self, url: str) -> None: + @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) - if urlsplit(url).path != "/v1.0/me/messages": - raise OutlookBoundaryError( - "Outlook continuation changed the message collection", - source_short_name="outlook_mail", + 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, @@ -169,54 +357,27 @@ async def capture_page( files: FileService, parent: SourceRecord | None = None, ) -> CapturePage: - """Capture one complete message; durable pending IDs avoid refetching inventory.""" + """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 OutlookBoundaryError( - "Outlook requires the mailbox message root", source_short_name="outlook_mail" - ) + 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: - if state.started and state.next_link is None: - return CapturePage(records=(), continuation=continuation, final=True) - url = state.next_link or f"{BASE}/messages" - self._list_url(url) - raw = await self.graph.get( - url, - params=None if state.started else {"$select": "id", "$top": 25}, - immutable_ids=True, - ) - try: - page = OutlookMessagePage.model_validate(raw) - except ValidationError: - raise OutlookBoundaryError( - "Invalid Outlook message inventory", source_short_name="outlook_mail" - ) from None - if page.next_link: - self._list_url(page.next_link) - if page.next_link == url: - raise OutlookBoundaryError( - "Outlook continuation did not advance", source_short_name="outlook_mail" - ) - state = OutlookMailContinuation( - fingerprint=self.fingerprint, - started=True, - pending_ids=tuple(dict.fromkeys(item.id for item in page.value)), - next_link=page.next_link, - ) - self._continuation(state) # Validate the byte bound before hydrating any originals. + state = await self._delta_page(state) records = () if state.pending_ids: observation = await self.message(state.pending_ids[0], files=files) - # New out-of-selection objects are not stored as synthetic tombstones. - # Existing objects are checked by the shared exact omission phase. - records = (observation,) if observation.kind == "upsert" else () + # 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:]}) - return CapturePage( - records=records, - continuation=self._continuation(state), - final=not state.pending_ids and state.next_link is None, - ) + 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: diff --git a/backend/airweave/platform/sources/outlook_mail_models.py b/backend/airweave/platform/sources/outlook_mail_models.py index b5f34fa475..05d2a88572 100644 --- a/backend/airweave/platform/sources/outlook_mail_models.py +++ b/backend/airweave/platform/sources/outlook_mail_models.py @@ -70,23 +70,78 @@ class OutlookMessagePage(BaseModel): next_link: str | None = Field(default=None, alias="@odata.nextLink", min_length=1) -class OutlookMailContinuation(BaseModel): - """Only pending native IDs and one native next link; no mailbox-sized folder map.""" +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[1] = 1 + version: Literal[2] = 2 fingerprint: str = Field(pattern=r"^[a-f0-9]{64}$") - started: bool = False + 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 = None + next_link: str | None = Field(default=None, min_length=1) @model_validator(mode="after") def coherent_progress(self): - """An initial cursor cannot masquerade as an acknowledged native page.""" - if len(set(self.pending_ids)) != len(self.pending_ids) or any( - not x for x in self.pending_ids - ): - raise ValueError("Pending Outlook identities must be nonempty and unique") - if not self.started and (self.pending_ids or self.next_link is not None): - raise ValueError("Initial Outlook continuation cannot contain progress") + """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/tests/test_outlook_identity.py b/backend/airweave/platform/sources/tests/test_outlook_identity.py index 134bb200fb..c746b331ba 100644 --- a/backend/airweave/platform/sources/tests/test_outlook_identity.py +++ b/backend/airweave/platform/sources/tests/test_outlook_identity.py @@ -290,3 +290,30 @@ def graph(request): 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) diff --git a/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py b/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py index bc4b8796e8..3279531c4e 100644 --- a/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py +++ b/backend/airweave/platform/sources/tests/test_outlook_mail_capture.py @@ -61,6 +61,25 @@ def initial(source): ) +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 = [] @@ -71,9 +90,16 @@ 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 == "/v1.0/me/messages": + 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"}]} + 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: @@ -85,7 +111,7 @@ def graph(request): async with httpx.AsyncClient(transport=httpx.MockTransport(graph)) as client: source = await capture(client) retained = files() - first = await source.capture_page(ROOT, initial(source), files=retained) + 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" @@ -101,7 +127,8 @@ def graph(request): 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/messages") == 1 + assert requested.count("/v1.0/me/mailFolders/folder-b/messages/delta") == 1 + assert second.provider_checkpoint is not None @pytest.mark.asyncio @@ -167,12 +194,15 @@ def graph(request): @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/messages?$skip=next" + 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: @@ -182,11 +212,14 @@ def graph(request): 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 message collection"): - await source.capture_page(ROOT, initial(source), files=files()) - next_link = "https://graph.microsoft.com/v1.0/me/messages?token=" + "x" * 66000 - with pytest.raises(ValueError, match="64 KiB"): - await source.capture_page(ROOT, initial(source), files=files()) + 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 @@ -237,3 +270,190 @@ def graph(request): ) 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 From 191739d6c1d0f85c283ce894f5e659c5c5fe0866 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 09:58:10 -0700 Subject: [PATCH 235/266] Validate omitted Slack messages without treating them as channels --- .../tests/test_slack_file_recovery.py | 59 ++++++++++++++++++- backend/airweave/platform/sources/slack.py | 35 +++++++++++ .../platform/sources/test_slack_capture.py | 33 +++++++++++ 3 files changed, 125 insertions(+), 2 deletions(-) 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 index 2bbe09a8f2..0ada62b9ad 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py @@ -34,7 +34,9 @@ } -async def attempt(database, source, storage, calls, *, number, fail_last=False): +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( @@ -44,7 +46,9 @@ async def attempt(database, source, storage, calls, *, number, fail_last=False): expected_team_id="T1", expected_user_id="U1", capture_files=True ) connector._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") - connector._get = AsyncMock(side_effect=[ROOT, {"messages": [MESSAGE]}]) + connector._get = AsyncMock(side_effect=responses or [ROOT, {"messages": [MESSAGE]}]) + 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): @@ -112,3 +116,54 @@ async def test_last_file_retry_preserves_committed_children_and_inventory_hides_ 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 diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 09f0c24098..152354a493 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -567,6 +567,11 @@ async def _capture_page_response( 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}) @@ -576,6 +581,36 @@ async def confirm_absent(self, record: SourceRecord) -> None: 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.""" + 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")) + if any(message.get("ts") == identity.native_id for message in messages): + raise ValueError("Slack omitted an accessible prior message; retry enumeration") + raise ValueError("Slack message omission remains unconfirmed; capture is incomplete") + # ------------------------------------------------------------------ # Federated search # ------------------------------------------------------------------ diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 4e900a7c97..5f24f3a333 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -227,3 +227,36 @@ async def test_conversation_dates_use_distinct_native_units(): assert first.source_updated_at.isoformat() == "2023-07-21T18:56:43.820000+00:00" assert malformed.source_created_at is malformed.source_updated_at is None assert absent.source_created_at is absent.source_updated_at is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply", [False, True]) +async def test_message_omission_reads_exact_native_message_and_never_confirms_empty(reply): + connector = await source() + payload = {"ts": "2", **({"thread_ts": "1"} if reply else {})} + record = SourceRecord.model_construct( + identity=RecordIdentity(record_type="message", native_id="2", container_id="C1"), + parent=RecordIdentity(record_type="channel", native_id="C1"), + payload=payload, + ) + connector._get = AsyncMock(return_value={"messages": [payload]}) + with pytest.raises(ValueError, match="accessible prior message"): + await connector.confirm_absent(record) + operation = "conversations.replies" if reply else "conversations.history" + connector._get.assert_awaited_once_with( + f"https://slack.com/api/{operation}", + { + "channel": "C1", + "oldest": "2", + "latest": "2", + "inclusive": "true", + "limit": 1, + **({"ts": "1"} if reply else {}), + }, + ) + connector._get = AsyncMock(return_value={"messages": []}) + with pytest.raises(ValueError, match="remains unconfirmed"): + await connector.confirm_absent(record) + connector._get = AsyncMock(side_effect=SlackApiError("channel_not_found")) + with pytest.raises(SlackApiError): + await connector.confirm_absent(record) From 32d56ebd098890277aa64d9b37583d52d9389114 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:22:01 -0700 Subject: [PATCH 236/266] Preserve Slack recovery with fenced exact attachment-owner validation --- .../entities/canonical/cycle_models.py | 21 ++ .../domains/entities/canonical/cycle_store.py | 61 +++- .../domains/entities/canonical/scan_models.py | 12 + .../domains/entities/canonical/scan_store.py | 128 ++++++++- .../domains/entities/canonical/service.py | 6 + .../entities/canonical/tests/conftest.py | 1 + .../tests/test_exact_parent_validation.py | 271 ++++++++++++++++++ .../tests/test_slack_file_recovery.py | 5 +- .../canonical/tests/test_slack_recovery.py | 47 ++- .../sync_pipeline/canonical_capture.py | 1 + .../domains/sync_pipeline/canonical_scan.py | 45 ++- backend/airweave/models/capture_scan.py | 2 + backend/airweave/platform/sources/slack.py | 58 +++- .../versions/0012_scan_parent_verification.py | 28 ++ backend/tests/live/canonical_capture.py | 1 + .../platform/sources/test_slack_capture.py | 32 +++ 16 files changed, 669 insertions(+), 50 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_exact_parent_validation.py create mode 100644 backend/alembic/versions/0012_scan_parent_verification.py diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 38b7507ef9..3715f42c71 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -47,6 +47,7 @@ class CycleConfiguration(BaseModel): completion_policies: dict[RecordKind, CompletionPolicy] = Field(default_factory=dict) known_object_validation: tuple[RecordKind, ...] = () scope_changes: tuple[RecordKind, ...] = () + exact_parent_validation: tuple[RecordKind, ...] = () membership: Literal["retained", "observed"] = "retained" def digest(self) -> str: @@ -56,6 +57,8 @@ def digest(self) -> str: value.pop("membership") # Preserve existing provider checkpoint digests. if not self.known_object_validation: value.pop("known_object_validation") # Preserve existing default configuration digests. + if not self.exact_parent_validation: + value.pop("exact_parent_validation") if not self.scope_changes: value.pop("scope_changes") value["completion_policies"] = {kind: self.policy(kind) for kind in self.parents} @@ -75,6 +78,7 @@ def from_source( completion_policies: dict[RecordKind, CompletionPolicy] | None = None, known_object_validation: tuple[RecordKind, ...] = (), scope_changes: tuple[RecordKind, ...] = (), + exact_parent_validation: tuple[RecordKind, ...] = (), ) -> "CycleConfiguration": """Snapshot the existing source declaration; no second topology registry.""" if len(set(record_types)) != len(record_types): @@ -91,6 +95,7 @@ def from_source( completion_policies=completion_policies or {}, known_object_validation=known_object_validation, scope_changes=scope_changes, + exact_parent_validation=exact_parent_validation, ) @model_validator(mode="before") @@ -146,11 +151,27 @@ def reachable_types(self) -> "CycleConfiguration": raise ValueError("Capture kinds must be reachable from a root") return self + @model_validator(mode="after") + def valid_exact_parent_kinds(self) -> "CycleConfiguration": + """Root inventories always revalidate discovery before descendant work.""" + if len(set(self.exact_parent_validation)) != len(self.exact_parent_validation) or any( + kind not in self.parents or None in self.parents[kind] or not self.children_of(kind) + for kind in self.exact_parent_validation + ): + raise ValueError("Exact parent validation requires unique nonroot owner kinds") + if self.exact_parent_validation and self.scope_changes: + raise ValueError("Exact parent validation does not support mixed scope plans") + return self + @property def root_record_types(self) -> tuple[str, ...]: """Kinds permitting records without a parent.""" return tuple(kind for kind, parents in self.parents.items() if None in parents) + def fresh_inventory(self, kind: str) -> bool: + """Exact owner reads replace attempt-wide inventory freshness only by opt-in.""" + return bool(self.children_of(kind)) and kind not in self.exact_parent_validation + def children_of(self, kind: str) -> tuple[str, ...]: """Scopes required for a visible record of this kind.""" return tuple(child for child, parents in self.parents.items() if kind in parents) diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index f123514887..f3e3ebf313 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -3,7 +3,7 @@ from datetime import datetime, timezone from uuid import UUID, uuid4 -from sqlalchemy import and_, exists, select +from sqlalchemy import and_, exists, or_, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import aliased @@ -192,7 +192,24 @@ def inventory_complete(subject, fence: WriterFence, state: CaptureCycle): CaptureScan.parent_visibility_epoch.is_not_distinct_from( subject.parent_visibility_epoch ), - CaptureScan.membership_attempt_id == fence.attempt_id, + or_( + subject.entity_definition_short_name.in_( + state.configuration.exact_parent_validation + ), + CaptureScan.membership_attempt_id == fence.attempt_id, + ), + or_( + subject.parent_record_type.is_(None), + subject.parent_record_type.not_in(state.configuration.exact_parent_validation), + and_( + CaptureScan.parent_verified_attempt_id == fence.attempt_id, + CaptureScan.parent_verified_revision + == select(owner.record_revision) + .where(owner.id == parent_id) + .correlate(subject) + .scalar_subquery(), + ), + ), CaptureScan.phase == "complete", ) ) @@ -222,7 +239,7 @@ async def root_ready(db: AsyncSession, fence: WriterFence, state: CaptureCycle) CaptureScan.cycle_id == state.version.cycle_id, CaptureScan.phase == "complete", ] - if state.configuration.children_of(kind): + if state.configuration.fresh_inventory(kind): predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) if await db.scalar(select(CaptureScan.id).where(*predicates).limit(1)) is None: return False @@ -267,7 +284,12 @@ async def scope_owner( async def attest_scope( - db: AsyncSession, fence: WriterFence, state: CaptureCycle, scope: CompletedScope + db: AsyncSession, + fence: WriterFence, + state: CaptureCycle, + scope: CompletedScope, + *, + require_owner_receipt: bool = True, ) -> bool: """A scope's full owner chain must be visible and freshly enumerated.""" parent = await scope_owner(db, fence, state, scope) @@ -285,6 +307,24 @@ async def attest_scope( ) if ready is None or not await root_ready(db, fence, state): raise CycleConflict("Refresh and complete ancestor membership before child work") + if ( + require_owner_receipt + and parent.entity_definition_short_name in state.configuration.exact_parent_validation + ): + verified = await db.scalar( + select(CaptureScan.id).where( + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.parent_record_id == parent.id, + CaptureScan.record_type == scope.record_type, + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.parent_visibility_epoch == parent.visibility_epoch, + CaptureScan.parent_verified_attempt_id == fence.attempt_id, + CaptureScan.parent_verified_revision == parent.record_revision, + ) + ) + if verified is None: + raise CycleConflict("Scope requires fresh exact parent verification") return False @@ -299,8 +339,17 @@ def child_scope_complete(fence: WriterFence, state: CaptureCycle, record_type: s CaptureScan.cycle_id == state.version.cycle_id, CaptureScan.phase == "complete", ] - if state.configuration.children_of(record_type): + if state.configuration.fresh_inventory(record_type): predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) + predicates.append( + or_( + Entity.entity_definition_short_name.not_in(state.configuration.exact_parent_validation), + and_( + CaptureScan.parent_verified_attempt_id == fence.attempt_id, + CaptureScan.parent_verified_revision == Entity.record_revision, + ), + ) + ) return exists(select(CaptureScan.id).where(*predicates)) @@ -317,7 +366,7 @@ async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) CaptureScan.cycle_id == cycle_id, CaptureScan.phase == "complete", ] - if state.configuration.children_of(kind): + if state.configuration.fresh_inventory(kind): predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) if await db.scalar(select(CaptureScan.id).where(*predicates).limit(1)) is None: return ScopeWork(record_type=kind) diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index cb3090f0eb..cb1e9a63ee 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -51,6 +51,7 @@ class BeginScan(BaseModel): plan: ScopePlan | None = None expected_parent_epoch: int | None = Field(default=None, ge=1) expected_parent_revision: int | None = Field(default=None, ge=1) + exact_parent_observation: CaptureRecord | None = None continuation: ScanContinuation = Field(default_factory=ScanContinuation) @@ -67,6 +68,8 @@ class ScanState(BaseModel): started_at: AwareDatetime completed_at: AwareDatetime | None parent_visibility_epoch: int | None = None + parent_verified_attempt_id: UUID | None = None + parent_verified_revision: int | None = None membership_attempt_id: UUID | None = None completion_policy: CompletionPolicy = "exhaustive" mode: Literal["full", "changes"] = "full" @@ -74,6 +77,15 @@ class ScanState(BaseModel): execution: ScopeExecution | None = None +class ScanAdmission(BaseModel): + """An unavailable refreshed owner is committed without admitting child work.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + state: ScanState | None + parent: SourceRecord | None + capture: CaptureResult + + class CommitScanPage(BaseModel): """Records and next state become durable together under one writer fence.""" diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 355d9f3b0c..45798282d4 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -39,6 +39,7 @@ CommitOmission, CommitScanPage, ReconcileScan, + ScanAdmission, ScanContinuation, ScanResult, ScanState, @@ -52,6 +53,7 @@ from airweave.domains.entities.canonical.store import ( CanonicalRecordStore, CanonicalStoreError, + content_is_available, source_record, ) from airweave.models.capture_scan import CaptureScan @@ -102,6 +104,8 @@ def scan_state(row: CaptureScan, parent: Entity | None = None) -> ScanState: completed_at=row.completed_at, parent_visibility_epoch=row.parent_visibility_epoch, membership_attempt_id=row.membership_attempt_id, + parent_verified_attempt_id=row.parent_verified_attempt_id, + parent_verified_revision=row.parent_verified_revision, execution=ScopeExecution.model_validate(row.execution_state) if row.execution_state else None, @@ -210,12 +214,28 @@ def validate_scope_changes( def attest_planned_owner(cycle: CaptureCycle, request: BeginScan, parent: Entity | None) -> None: """The post-I/O plan must still describe the captured owner used to select it.""" if request.expected_parent_epoch != (parent.visibility_epoch if parent else None) or ( - cycle.mode == "mixed" + (cycle.mode == "mixed" or request.exact_parent_observation is not None) and request.expected_parent_revision != (parent.record_revision if parent else None) ): raise ScanConflict("Scope owner changed during plan selection") +def preserve_owner_verification( + row: CaptureScan, request: BeginScan, parent: Entity | None +) -> None: + """A cursor restart preserves only a receipt still bound to this exact owner.""" + keep_verification = bool( + parent + and row.cycle_id == request.cycle_id + and row.parent_visibility_epoch == parent.visibility_epoch + and row.parent_verified_attempt_id == request.fence.attempt_id + and row.parent_verified_revision == parent.record_revision + ) + if not keep_verification: + row.parent_verified_attempt_id = None + row.parent_verified_revision = None + + class LockedPageCapture(Protocol): """Internal composition: capture the entire page under the caller's writer lock. @@ -308,14 +328,107 @@ def _check_parent_epoch(row: CaptureScan, parent: Entity | None) -> None: if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): raise ScanConflict("Scope owner changed; restart its sweep with current CAS") + async def admit(self, db: AsyncSession, request: BeginScan) -> ScanAdmission: + """Commit fresh owner state and its child receipt under the existing writer fence.""" + sync = await self.records._fenced_sync(db, request.fence) + cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) + await attest_scope( + db, + request.fence, + cycle, + request.scope, + require_owner_receipt=request.exact_parent_observation is None, + ) + parent = await scope_owner(db, request.fence, cycle, request.scope) + observation = request.exact_parent_observation + empty = CaptureResult(changes=(), sequence=sync.observed_change_sequence, unchanged=0) + if observation is None: + state = await self._begin(db, request, context=(cursor, cycle, parent)) + return ScanAdmission( + state=state, parent=source_record(parent) if parent else None, capture=empty + ) + if ( + parent is None + or parent.entity_definition_short_name + not in cycle.configuration.exact_parent_validation + ): + raise ScanConflict("Scope does not accept exact parent verification") + attest_planned_owner(cycle, request, parent) + if request.fingerprint != cycle.configuration.fingerprint: + raise CycleConflict("Scan configuration differs from the active cycle") + row = await self._row(db, request.fence, request.scope) + if row is None: + if request.expected is not None: + raise ScanConflict("Verification expected an existing scope") + elif request.expected is None: + raise ScanConflict("Verification requires the current scope version") + else: + self._expect(row, request.expected, row.cycle_id) + previous = source_record(parent) + if ( + observation.identity != previous.identity + or observation.parent != previous.parent + or observation.allow_reparent + or observation.removal_reason == "absent" + ): + raise ScanConflict("Exact parent observation changed identity or claimed absence") + captured = await self.records._capture_locked( + db, + sync, + CaptureBatch(fence=request.fence, records=(observation,)), + mark_seen=False, + ) + await db.refresh(parent) + available = await db.scalar(select(content_is_available()).where(Entity.id == parent.id)) + if parent.deleted_at is not None or not available: + return ScanAdmission(state=None, parent=source_record(parent), capture=captured) + updated = request.model_copy( + update={ + "expected_parent_epoch": parent.visibility_epoch, + "expected_parent_revision": parent.record_revision, + "restart": request.restart + or bool(row and row.parent_visibility_epoch != parent.visibility_epoch), + } + ) + await self._begin(db, updated, context=(cursor, cycle, parent)) + row = await self._row(db, request.fence, request.scope) + row.parent_verified_attempt_id = request.fence.attempt_id + row.parent_verified_revision = parent.record_revision + row.revision += 1 + await db.flush() + return ScanAdmission( + state=await self._state(db, row), parent=source_record(parent), capture=captured + ) + async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: - """Resume unchanged scans; restart and next-cycle transitions require exact CAS.""" + """Ordinary admission cannot accept uncommitted owner observations.""" + if request.exact_parent_observation is not None: + raise ScanConflict("Exact parent observations require typed scan admission") + return await self._begin(db, request) + + async def _begin_context( + self, + db: AsyncSession, + request: BeginScan, + ) -> tuple[SyncCursor, CaptureCycle, Entity | None]: + """Ordinary entry obtains the same authority checks as exact admission.""" await self.records._fenced_sync(db, request.fence) cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) await attest_scope(db, request.fence, cycle, request.scope) parent = await scope_owner(db, request.fence, cycle, request.scope) + return cursor, cycle, parent + + async def _begin( + self, + db: AsyncSession, + request: BeginScan, + *, + context: tuple[SyncCursor, CaptureCycle, Entity | None] | None = None, + ) -> ScanState: + """Resume unchanged scans; reuse admission checks held under the same writer lock.""" + cursor, cycle, parent = context or await self._begin_context(db, request) attest_planned_owner(cycle, request, parent) - inventory = bool(cycle.configuration.children_of(request.scope.record_type)) + inventory = cycle.configuration.fresh_inventory(request.scope.record_type) if request.fingerprint != cycle.configuration.fingerprint: raise CycleConflict("Scan configuration differs from the active cycle") row = await self._row(db, request.fence, request.scope) @@ -361,6 +474,7 @@ async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: raise ScanConflict("Changed scope configuration requires an explicit restart") row.revision += 1 begin_execution(row, cycle, request, parent.visibility_epoch if parent else None) + preserve_owner_verification(row, request, parent) row.parent_visibility_epoch = parent.visibility_epoch if parent else None row.membership_attempt_id = request.fence.attempt_id if inventory else None row.cycle_id = request.cycle_id @@ -404,7 +518,7 @@ async def page( if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): raise ScanConflict("Scope owner changed; restart from its current epoch") if ( - cycle.configuration.children_of(row.record_type) + cycle.configuration.fresh_inventory(row.record_type) and row.membership_attempt_id != request.fence.attempt_id ): raise CycleConflict("Membership belongs to an earlier writer attempt") @@ -525,7 +639,7 @@ async def missing( if ( row.phase != "reconciling" or ( - cycle.configuration.children_of(row.record_type) + cycle.configuration.fresh_inventory(row.record_type) and row.membership_attempt_id != fence.attempt_id ) or row.parent_visibility_epoch != (parent.visibility_epoch if parent else None) @@ -616,7 +730,7 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul if row.parent_visibility_epoch != (parent.visibility_epoch if parent else None): raise ScanConflict("Scope owner changed; restart from its current epoch") if ( - cycle.configuration.children_of(row.record_type) + cycle.configuration.fresh_inventory(row.record_type) and row.membership_attempt_id != request.fence.attempt_id ): raise CycleConflict("Membership belongs to an earlier writer attempt") @@ -648,7 +762,7 @@ async def reconcile(self, db: AsyncSession, request: ReconcileScan) -> ScanResul scope=(await self._state(db, row)).scope, removal_reason=( "scope_removed" - if cycle.configuration.children_of(row.record_type) + if cycle.configuration.fresh_inventory(row.record_type) else request.removal_reason ), observed_at=request.observed_at, diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index 5a48ca8283..cccc879f52 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -36,6 +36,7 @@ CommitOmission, CommitScanPage, ReconcileScan, + ScanAdmission, ScanResult, ScanState, ) @@ -159,6 +160,11 @@ async def begin_scan(self, db: AsyncSession, request: BeginScan) -> ScanState: async with UnitOfWork(db): return await self.scans.begin(db, request) + async def admit_scan(self, db: AsyncSession, request: BeginScan) -> ScanAdmission: + """Exact owner capture and child admission commit or roll back together.""" + async with UnitOfWork(db): + return await self.scans.admit(db, request) + async def commit_scan_page(self, db: AsyncSession, request: CommitScanPage) -> ScanResult: """Commit original records and their continuation atomically.""" async with UnitOfWork(db): diff --git a/backend/airweave/domains/entities/canonical/tests/conftest.py b/backend/airweave/domains/entities/canonical/tests/conftest.py index 37328d3238..b77c43f40d 100644 --- a/backend/airweave/domains/entities/canonical/tests/conftest.py +++ b/backend/airweave/domains/entities/canonical/tests/conftest.py @@ -281,6 +281,7 @@ async def database(request): await connection.run_sync(migrate, "0009_extraction_coverage.py") await connection.run_sync(migrate, "0010_owned_provisioning.py") await connection.run_sync(migrate, "0011_retained_projection_text.py") + await connection.run_sync(migrate, "0012_scan_parent_verification.py") yield async_sessionmaker(engine, expire_on_commit=False) finally: await engine.dispose() diff --git a/backend/airweave/domains/entities/canonical/tests/test_exact_parent_validation.py b/backend/airweave/domains/entities/canonical/tests/test_exact_parent_validation.py new file mode 100644 index 0000000000..a83af6669f --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_exact_parent_validation.py @@ -0,0 +1,271 @@ +"""Fenced exact owner reads authorize child work without resetting discovery.""" + +from uuid import uuid4 + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import BeginCycle, CycleConfiguration +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 +from airweave.domains.entities.canonical.scan_store import ScanConflict +from airweave.domains.entities.canonical.tests.helpers import capture +from airweave.domains.entities.canonical.tests.test_cycles import finish, page +from airweave.domains.entities.canonical.tests.test_forest_scans import collect, item, work +from airweave.models.entity import Entity + +CONFIG = CycleConfiguration( + fingerprint="a" * 64, + parents={"channel": (None,), "message": ("channel",), "file": ("message",)}, + exact_parent_validation=("message",), +) + + +async def setup(database, source): + service, fence = source + async with database() as db: + cycle = await service.begin_cycle(db, BeginCycle(fence=fence, configuration=CONFIG)) + channel = item("channel", "C") + message = item("message", "M", channel.identity) + await collect(database, service, fence, cycle, CompletedScope(record_type="channel"), channel) + await collect( + database, + service, + fence, + cycle, + CompletedScope( + record_type="message", + container_id=message.identity.container_id, + parent=channel.identity, + ), + message, + ) + selected = await work(database, service, fence, cycle) + owner = selected.parent + scope = CompletedScope( + record_type="file", + container_id=parent_container_key(message.identity), + parent=message.identity, + ) + request = BeginScan( + fence=fence, + scope=scope, + cycle_id=cycle.version.cycle_id, + fingerprint=CONFIG.fingerprint, + expected_parent_epoch=selected.parent_visibility_epoch, + expected_parent_revision=owner.revision, + exact_parent_observation=message, + ) + return cycle, message, owner, request + + +async def test_receipt_preserves_sighting_and_completed_children(database, source): + service, fence = source + cycle, message, owner, request = await setup(database, source) + async with database() as db: + seen = await db.scalar(select(Entity.last_seen_run_id).where(Entity.id == owner.id)) + admitted = await service.admit_scan(db, request) + assert admitted.parent == owner + assert admitted.state.parent_verified_attempt_id == fence.attempt_id + state = await page(database, service, fence, admitted.state, final=True) + state = await finish(database, service, fence, state) + async with database() as db: + assert await db.scalar(select(Entity.last_seen_run_id).where(Entity.id == owner.id)) == seen + repeated = await service.admit_scan( + db, + request.model_copy( + update={ + "expected": state.version, + "exact_parent_observation": None, + } + ), + ) + assert repeated.state.phase == "complete" + assert await work(database, service, fence, cycle) is None + + +async def test_missing_receipt_and_stale_post_io_revision_fail(database, source): + service, fence = source + _, message, _, request = await setup(database, source) + with pytest.raises(CycleConflict, match="fresh exact"): + async with database() as db: + await service.begin_scan( + db, request.model_copy(update={"exact_parent_observation": None}) + ) + await capture( + database, service, fence, message.model_copy(update={"payload": {"id": "M", "new": True}}) + ) + with pytest.raises(ScanConflict, match="owner changed"): + async with database() as db: + await service.admit_scan(db, request) + + +async def test_unavailable_observation_commits_without_child_admission(database, source): + service, fence = source + _, message, owner, request = await setup(database, source) + unavailable = message.model_copy(update={"kind": "delete", "removal_reason": "scope_removed"}) + async with database() as db: + result = await service.admit_scan( + db, request.model_copy(update={"exact_parent_observation": unavailable}) + ) + assert result.state is None + assert result.parent.deleted_at is not None + async with database() as db: + assert (await db.get(Entity, owner.id)).deleted_at is not None + + +async def test_new_writer_needs_root_refresh_before_exact_admission(database, source): + service, fence = source + _, _, _, request = await setup(database, source) + 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, + ) + with pytest.raises(CycleConflict): + async with database() as db: + await service.admit_scan(db, request.model_copy(update={"fence": newer})) + + +async def test_owner_edit_invalidates_receipt_and_restarts_inventory(database, source): + service, fence = source + _, message, owner, request = await setup(database, source) + async with database() as db: + admitted = await service.admit_scan(db, request) + state = admitted.state + changed = message.model_copy( + update={ + "payload": {"id": "M", "files": ["new"]}, + "descendant_visibility_fields": ("files",), + } + ) + await capture(database, service, fence, changed) + with pytest.raises(CycleConflict, match="fresh exact"): + await page(database, service, fence, state, final=True) + async with database() as db: + current = await db.get(Entity, owner.id) + refreshed = await service.admit_scan( + db, + request.model_copy( + update={ + "expected": state.version, + "expected_parent_epoch": current.visibility_epoch, + "expected_parent_revision": current.record_revision, + "exact_parent_observation": changed, + } + ), + ) + assert refreshed.state.version.sweep_id != state.version.sweep_id + assert refreshed.parent.payload["files"] == ["new"] + assert refreshed.state.parent_verified_revision == refreshed.parent.revision + # Native cursor expiry may restart the sweep, retaining this still-valid receipt. + async with database() as db: + restarted = await service.begin_scan( + db, + request.model_copy( + update={ + "expected": refreshed.state.version, + "restart": True, + "expected_parent_epoch": refreshed.state.parent_visibility_epoch, + "expected_parent_revision": refreshed.parent.revision, + "exact_parent_observation": None, + } + ), + ) + await page(database, service, fence, restarted, final=True) + + +async def test_stale_writer_cannot_apply_observation(database, source): + from airweave.domains.entities.canonical.store import StaleWriter + + service, fence = source + _, message, owner, request = await setup(database, source) + async with database() as db: + await service.activate_writer( + db, + fence.organization_id, + fence.sync_id, + fence.job_id, + attempt_id=uuid4(), + attempt_number=2, + ) + with pytest.raises(StaleWriter): + async with database() as db: + await service.admit_scan( + db, + request.model_copy( + update={ + "exact_parent_observation": message.model_copy( + update={"payload": {"bad": True}} + ), + } + ), + ) + async with database() as db: + current = await db.get(Entity, owner.id) + assert current.source_payload == message.payload + assert current.record_revision == owner.revision + + +async def test_new_writer_reverifies_completed_children_without_recapture(database, source): + service, fence = source + cycle, message, _, request = await setup(database, source) + async with database() as db: + admitted = await service.admit_scan(db, request) + original = item("file", "F", message.identity) + state = await page(database, service, fence, admitted.state, original, final=True) + state = await finish(database, service, fence, state) + async with database() as db: + file_before = await db.scalar(select(Entity).where(Entity.native_id == "F")) + revision, payload = file_before.record_revision, file_before.source_payload + newer = await service.activate_writer( + db, + fence.organization_id, + fence.sync_id, + fence.job_id, + attempt_id=uuid4(), + attempt_number=2, + ) + assert (await work(database, service, newer, cycle)).record_type == "channel" + root_scope = CompletedScope(record_type="channel") + async with database() as db: + old_root = await service.read_scan(db, newer, root_scope) + root = await service.begin_scan( + db, + BeginScan( + fence=newer, + scope=root_scope, + cycle_id=cycle.version.cycle_id, + fingerprint=CONFIG.fingerprint, + expected=old_root.version, + restart=True, + ), + ) + root = await page(database, service, newer, root, item("channel", "C"), final=True) + await finish(database, service, newer, root) + pending = await work(database, service, newer, cycle) + assert pending.record_type == "file" # Completed history survives; receipt does not. + async with database() as db: + refreshed = await service.admit_scan( + db, + request.model_copy( + update={ + "fence": newer, + "expected": state.version, + "expected_parent_revision": pending.parent.revision, + "expected_parent_epoch": pending.parent_visibility_epoch, + } + ), + ) + assert refreshed.state.phase == "complete" + assert refreshed.state.version.sweep_id == state.version.sweep_id + assert refreshed.state.parent_verified_attempt_id == newer.attempt_id + assert await work(database, service, newer, cycle) is None + async with database() as db: + current = await db.scalar(select(Entity).where(Entity.native_id == "F")) + assert (current.record_revision, current.source_payload) == (revision, payload) 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 index 0ada62b9ad..3a3b81c94a 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py @@ -46,7 +46,10 @@ async def attempt( expected_team_id="T1", expected_user_id="U1", capture_files=True ) connector._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") - connector._get = AsyncMock(side_effect=responses or [ROOT, {"messages": [MESSAGE]}]) + # Fresh attempts enumerate history and then revalidate the attachment owner. + # A retry reuses completed history and only refreshes the owner. + defaults = [ROOT, *([{"messages": [MESSAGE]}] if number == 1 else []), {"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) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index b5a122f503..e3f82f53c2 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -24,14 +24,16 @@ REPLIES = {"messages": [MESSAGE, {"ts": "1.1", "thread_ts": "1", "text": "Synthetic reply"}]} -async def runner(database, source, responses, attempt=1, service=None): +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") + 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) @@ -68,10 +70,13 @@ async def saved(database): 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 + database, source, capture_files ): - first, _, _ = await runner(database, source, [ROOT, HISTORY, asyncio.CancelledError()]) + 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) @@ -79,12 +84,32 @@ async def test_cancel_after_history_resumes_pending_reply_without_repeating_hist 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) + second, connector, _ = await runner( + database, source, [ROOT, REPLIES], attempt=2, capture_files=capture_files + ) + if capture_files: + initial = iter([ROOT, REPLIES]) + + async def respond(url, params): + if "oldest" not in params: + return next(initial) + assert params["oldest"] == params["latest"] + return { + "messages": [ + message for message in REPLIES["messages"] if message["ts"] == params["oldest"] + ] + } + + connector._get = AsyncMock(side_effect=respond) await run(second) - assert [call.args[0].split("/")[-1] for call in connector._get.call_args_list] == [ - "conversations.list", - "conversations.replies", - ] + operations = [call.args[0].split("/")[-1] for call in connector._get.call_args_list] + assert operations[:2] == ["conversations.list", "conversations.replies"] + assert len(operations) == (4 if capture_files else 2) + if capture_files: + # Exact owner reads must not become another unbounded history traversal. + for call in connector._get.call_args_list[2:]: + assert call.args[1]["oldest"] == call.args[1]["latest"] + assert call.args[1]["limit"] == 1 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 @@ -117,7 +142,9 @@ async def test_expired_reply_cursor_restarts_whole_scope_and_never_partial_recon "messages": [{"ts": "1.old", "thread_ts": "1"}], "response_metadata": {"next_cursor": "expired"}, } - first, _, _ = await runner(database, source, [ROOT, HISTORY, first_reply, asyncio.CancelledError()]) + first, _, _ = await runner( + database, source, [ROOT, HISTORY, first_reply, asyncio.CancelledError()] + ) with pytest.raises(asyncio.CancelledError): await run(first) rows, _, scans = await saved(database) diff --git a/backend/airweave/domains/sync_pipeline/canonical_capture.py b/backend/airweave/domains/sync_pipeline/canonical_capture.py index d475eeff47..bf8c34faf5 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_capture.py +++ b/backend/airweave/domains/sync_pipeline/canonical_capture.py @@ -75,6 +75,7 @@ def __init__( 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") diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index 14baf7bc3d..d3110492ab 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -33,6 +33,7 @@ CommitOmission, CommitScanPage, ReconcileScan, + ScanAdmission, ScanContinuation, ScanState, ) @@ -207,7 +208,11 @@ async def scan( parent_epoch: int | None = None, ) -> None: """Only a provider invalid-cursor response permits one explicit sweep restart.""" - state = await self.begin_scope(cycle, scope, parent, parent_epoch, refresh_membership) + admission = await self.begin_scope(cycle, scope, parent, parent_epoch) + if admission.state is None: + return + state = admission.state + parent = admission.parent restarts = 0 while state.phase == "collecting": await self.check_limits() @@ -254,8 +259,7 @@ async def begin_scope( scope: CompletedScope, parent: SourceRecord | None, parent_epoch: int | None, - refresh_membership: bool, - ) -> ScanState: + ) -> 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) @@ -263,11 +267,16 @@ async def begin_scope( previous and previous.cycle_id == cycle.version.cycle_id and ( - (refresh_membership and previous.membership_attempt_id != self.fence.attempt_id) + ( + 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": @@ -287,9 +296,23 @@ async def begin_scope( 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: - return await self.service.begin_scan( + admission = await self.service.admit_scan( db, BeginScan( fence=self.fence, @@ -301,12 +324,14 @@ async def begin_scope( continuation=initial, plan=plan, expected_parent_epoch=parent_epoch, - expected_parent_revision=parent.revision - if parent and cycle.mode == "mixed" - else None, + 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, @@ -339,9 +364,7 @@ async def restart_scope( continuation=initial, plan=plan, expected_parent_epoch=state.parent_visibility_epoch, - expected_parent_revision=parent.revision - if parent and cycle.mode == "mixed" - else None, + expected_parent_revision=parent.revision if parent else None, ), ) diff --git a/backend/airweave/models/capture_scan.py b/backend/airweave/models/capture_scan.py index c30dabb0ac..fcae496086 100644 --- a/backend/airweave/models/capture_scan.py +++ b/backend/airweave/models/capture_scan.py @@ -31,6 +31,8 @@ class CaptureScan(OrganizationBase): 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] diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 152354a493..511309d587 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -259,6 +259,7 @@ def capture_cycle_configuration(self) -> CycleConfiguration: 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 @@ -362,20 +363,9 @@ async def capture_page( 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": - for record in page.records: - self._file_inventory(record.payload) page = page.model_copy( update={ - "records": tuple( - record.model_copy( - update={ - "payload_schema_version": 2, - "descendant_visibility_fields": ("files",), - "completeness": "complete", - } - ) - for record in page.records - ) + "records": tuple(self._attachment_owner(record) for record in page.records) } ) return page @@ -383,6 +373,31 @@ async def capture_page( # 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 @@ -583,6 +598,13 @@ async def confirm_absent(self, record: SourceRecord) -> None: 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 @@ -607,9 +629,15 @@ async def _confirm_message_omission(self, record: SourceRecord) -> None: 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")) - if any(message.get("ts") == identity.native_id for message in messages): - raise ValueError("Slack omitted an accessible prior message; retry enumeration") - raise ValueError("Slack message omission remains unconfirmed; capture is incomplete") + 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 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/tests/live/canonical_capture.py b/backend/tests/live/canonical_capture.py index 191ab4485c..ca5e1cc175 100644 --- a/backend/tests/live/canonical_capture.py +++ b/backend/tests/live/canonical_capture.py @@ -321,6 +321,7 @@ async def main(): "0009_extraction_coverage.py", "0010_owned_provisioning.py", "0011_retained_projection_text.py", + "0012_scan_parent_verification.py", ): await connection.run_sync(migrate, migration) with TemporaryDirectory(prefix="airweave-private-live-") as directory: diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index 5f24f3a333..cf5bc0c1c8 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -260,3 +260,35 @@ async def test_message_omission_reads_exact_native_message_and_never_confirms_em connector._get = AsyncMock(side_effect=SlackApiError("channel_not_found")) with pytest.raises(SlackApiError): await connector.confirm_absent(record) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply", [False, True]) +async def test_exact_attachment_owner_refresh_preserves_native_inventory_and_parent(reply): + connector = await source() + connector.slack_config = SlackConfig( + expected_team_id="T1", expected_user_id="U1", capture_files=True + ) + identity = RecordIdentity(record_type="message", native_id="2", container_id="C1") + parent = RecordIdentity(record_type="channel", native_id="C1") + payload = {"ts": "2", **({"thread_ts": "1"} if reply else {})} + record = SourceRecord.model_construct(identity=identity, parent=parent, payload=payload) + fresh = {**payload, "files": [{"id": "F2"}], "unknown_native_field": "retained"} + connector._get = AsyncMock(return_value={"messages": [fresh]}) + result = await connector.refresh_known(record, files=MagicMock()) + assert result.identity == identity and result.parent == parent + assert result.payload == fresh and result.payload_schema_version == 2 + assert result.descendant_visibility_fields == ("files",) + assert connector.capture_cycle_configuration.exact_parent_validation == ("message",) + operation = "conversations.replies" if reply else "conversations.history" + assert connector._get.call_args.args[0].endswith(operation) + assert connector._get.call_args.args[1]["oldest"] == "2" + assert connector._get.call_args.args[1]["latest"] == "2" + for messages, error in [ + ([], "remains unconfirmed"), + ([{**fresh, "thread_ts": "wrong"}], "thread identity"), + ([fresh, fresh], "duplicate identities"), + ]: + connector._get = AsyncMock(return_value={"messages": messages}) + with pytest.raises(ValueError, match=error): + await connector.refresh_known(record, files=MagicMock()) From 3dbaedff5d9b7f36f9c3d2c9a8941f650ed42177 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:28:22 -0700 Subject: [PATCH 237/266] Reject unversioned Wispr multi-range captures --- backend/airweave/platform/sources/wispr.py | 7 +++++ .../platform/sources/test_wispr_capture.py | 26 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/backend/airweave/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index c1fece6a77..00191b9626 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -234,6 +234,12 @@ def _next_offset(text: JsonValue, field: str) -> int | None: 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") + async def _body( self, kind: Literal["meeting", "scratchpad_note"], identity: str ) -> dict[str, JsonValue]: @@ -285,6 +291,7 @@ async def _body( 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") diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index 062a0bd074..735b794d4f 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -453,6 +453,7 @@ async def test_transcript_disappearing_after_first_range_stays_incomplete(monkey side_effect=[ { "id": "m", + "modified_at": "2026-10-01T00:00:00Z", "content": "notes", "has_transcript": True, "transcript": ( @@ -460,9 +461,32 @@ async def test_transcript_disappearing_after_first_range_stays_incomplete(monkey "continue with view_transcript.start_char=3...)" ), }, - {"id": "m", "has_transcript": False, "transcript": None}, + { + "id": "m", + "modified_at": "2026-10-01T00:00:00Z", + "has_transcript": False, + "transcript": None, + }, ] ), ) with pytest.raises(ValueError, match="not a string"): await connector._body("meeting", "m") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [{}, {"modified_at": None}, {"modified_at": ""}]) +async def test_multirange_body_requires_version_before_following_continuation(version): + connector = await source() + connector._execute = AsyncMock( + return_value={ + "id": "m", + **version, + "content": "notes", + "transcript": "abc\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)", + } + ) + with pytest.raises(ValueError, match="no usable version"): + await connector._body("meeting", "m") + connector._execute.assert_awaited_once() From 6d805a6618a2be675994c6fb325372ce2f6e4553 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:34:04 -0700 Subject: [PATCH 238/266] Bound Wispr streamed responses and retained body archives --- .../canonical/tests/test_wispr_recovery.py | 57 ++++++++ backend/airweave/platform/sources/wispr.py | 85 +++++++++++- .../platform/sources/test_wispr_capture.py | 127 ++++++++++++++++++ 3 files changed, 262 insertions(+), 7 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py index 0598513304..78a3b75eea 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_wispr_recovery.py @@ -278,3 +278,60 @@ async def test_active_meeting_only_cycle_requires_explicit_restart_for_scratchpa 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/platform/sources/wispr.py b/backend/airweave/platform/sources/wispr.py index 00191b9626..c703856b21 100644 --- a/backend/airweave/platform/sources/wispr.py +++ b/backend/airweave/platform/sources/wispr.py @@ -9,6 +9,7 @@ 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 @@ -21,15 +22,22 @@ RecordIdentity, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation -from airweave.domains.sources.exceptions import SourceServerError +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.""" @@ -162,13 +170,60 @@ async def create( return instance async def _post(self, path: str, body: dict) -> dict: - response = await self.http_client.post( + 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()}, + headers={ + "x-api-key": self._tool_auth.api_key.get_secret_value(), + "Accept-Encoding": "identity", + }, json=body, - ) - ComposioTransport._check_proxy_response(response, response.request) - return response.json() + ) 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: @@ -240,6 +295,18 @@ def _require_range_version(version: JsonValue) -> None: 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]: @@ -257,6 +324,7 @@ async def _body( {"content": 0}, ) responses: list[JsonValue] = [] + archive_bytes = len(b'{"responses":[]}') for _ in range(100): arguments = { id_field: identity, @@ -272,7 +340,9 @@ async def _body( "modified_at" ): raise ValueError("Wispr body changed during paginated capture") - responses.append({"requested_ranges": arguments, "response": response}) + 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 ( @@ -342,6 +412,7 @@ async def capture_page( 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=( diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index 735b794d4f..399aef0154 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -490,3 +490,130 @@ async def test_multirange_body_requires_version_before_following_continuation(ve with pytest.raises(ValueError, match="no usable version"): await connector._body("meeting", "m") connector._execute.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [200, 400, 401, 429, 500]) +async def test_streamed_response_limit_closes_success_and_error(monkeypatch, status): + import httpx + + from airweave.domains.storage import FileSkippedException + from airweave.platform.sources import wispr + + class Stream(httpx.AsyncByteStream): + closed = False + + def __init__(self): + self.iterator = self.chunks() + + async def chunks(self): + yield b"x" * 16 + yield b"never needed" + + def __aiter__(self): + return self.iterator + + async def aclose(self): + await self.iterator.aclose() + self.closed = True + + stream = Stream() + monkeypatch.setattr(wispr, "MAX_RESPONSE_BYTES", 8) + monkeypatch.setattr(wispr, "MAX_ERROR_BYTES", 8) + connector = await source() + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(status, stream=stream)) + ) as client: + connector._http_client = client + with pytest.raises(Exception) as caught: + await connector._post("session", {}) + assert not isinstance(caught.value, FileSkippedException) + assert stream.closed + expected = { + 200: "SourceError", + 400: "ComposioProxyError", + 401: "AuthProviderAuthError", + 429: "AuthProviderRateLimitError", + 500: "AuthProviderServerError", + } + assert type(caught.value).__name__ == expected[status] + + +@pytest.mark.asyncio +async def test_nonidentity_encoding_is_rejected_and_valid_unicode_preserved(): + import json + + import httpx + + from airweave.domains.sources.exceptions import SourceError + + connector = await source() + value = {"data": {"text": "你好 café 🎵"}, "error": None} + + def respond(request): + assert request.headers["accept-encoding"] == "identity" + return httpx.Response(200, content=json.dumps(value, ensure_ascii=False).encode()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + connector._http_client = client + assert await connector._post("session", {}) == value + async with httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, headers={"content-encoding": "br"}, content=b"") + ) + ) as client: + connector._http_client = client + with pytest.raises(SourceError, match="identity encoding"): + await connector._post("session", {}) + + +@pytest.mark.asyncio +async def test_aggregate_archive_limit_counts_native_metadata(monkeypatch): + from airweave.domains.sources.exceptions import SourceError + from airweave.platform.sources import wispr + + connector = await source() + first = { + "id": "m", + "modified_at": "same", + "content": "你好", + "native_metadata": "x" * 250, + "transcript": ( + "(...truncated, 3 chars remaining; continue with view_transcript.start_char=3...)" + ), + } + second = { + "id": "m", + "modified_at": "same", + "transcript": "🎵 café", + "native_metadata": "y" * 250, + } + execute = AsyncMock(side_effect=[first, second]) + monkeypatch.setattr(connector, "_execute", execute) + monkeypatch.setattr(wispr, "MAX_BODY_BYTES", 700) + with pytest.raises(SourceError, match="body exceeds byte limit"): + await connector._body("meeting", "m") + assert execute.await_count == 2 + execute.side_effect = [first, second] + monkeypatch.setattr(wispr, "MAX_BODY_BYTES", 2000) + captured = await connector._body("meeting", "m") + assert [part["response"] for part in captured["responses"]] == [first, second] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status, expected", [(401, "AuthProviderAuthError"), (429, "AuthProviderRateLimitError")] +) +async def test_encoded_error_preserves_header_status_failure(status, expected): + import httpx + + connector = await source() + async with httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response(status, headers={"content-encoding": "br"}, content=b"") + ) + ) as client: + connector._http_client = client + with pytest.raises(Exception) as caught: + await connector._post("session", {}) + assert type(caught.value).__name__ == expected From fae431a1c13274b9a058394a7919b4816927c550 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:39:36 -0700 Subject: [PATCH 239/266] Provision owned Linear and Attio sources with fixed workspace scope --- .../domains/owned_provisioning/models.py | 59 ++++++++- .../domains/owned_provisioning/store.py | 10 ++ .../tests/test_provisioning.py | 120 ++++++++++++++++++ 3 files changed, 185 insertions(+), 4 deletions(-) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 7455af8b8c..0b081833de 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -4,14 +4,31 @@ from uuid import UUID from croniter import croniter -from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator +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", "stripe"] + provider: Literal[ + "gmail", + "google_calendar", + "google_drive", + "slack", + "outlook_mail", + "stripe", + "linear", + "attio", + ] 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) @@ -27,10 +44,18 @@ def native_principal(self): """Slack workspace membership requires its native user, not broker user_id.""" if (self.provider == "slack") != (self.expected_user_identity is not None): raise ValueError("Only Slack requires an expected native user identity") - if self.provider == "stripe": - self.source_config() # Validate mode/version before admission, not during delivery. + if self.provider in {"stripe", "linear", "attio"}: + self.source_config() # Validate native scope before admission, not during delivery. return self + @field_validator("expected_identity") + @classmethod + def workspace_identity(cls, value: str, info: ValidationInfo) -> str: + """Workspace UUIDs have one canonical spelling across admission and reconnect.""" + if info.data.get("provider") in {"linear", "attio"}: + return str(UUID(value)) + return value + @field_validator("cron") @classmethod def valid_schedule(cls, value: str) -> str: @@ -41,6 +66,17 @@ def valid_schedule(cls, value: str) -> str: def source_config(self) -> dict[str, JsonValue]: """Expected identity cannot be overridden inside unstructured provider config.""" + 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 @@ -109,8 +145,23 @@ class ProvisionedSource(BaseModel): 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 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, diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py index b71a6db268..3c24e32fbc 100644 --- a/backend/airweave/domains/owned_provisioning/store.py +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -220,6 +220,16 @@ async def _configure( detail="Reconnect changes Stripe account mode, context or API version; " "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() diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index a373db3263..e903dad390 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -677,3 +677,123 @@ def test_stripe_provisioning_requires_explicit_typed_mode_and_version(config): 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 * * * *", + ) From 077367642e60001fcefe2375dd99d372d420f4f6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:49:45 -0700 Subject: [PATCH 240/266] Qualify GitHub native identity and standalone provisioning --- backend/airweave/api/examples.py | 15 +-- .../canonical/tests/test_github_recovery.py | 1 + .../domains/owned_provisioning/models.py | 26 +++- .../domains/owned_provisioning/store.py | 11 ++ .../tests/test_provisioning.py | 118 ++++++++++++++++++ backend/airweave/platform/configs/config.py | 5 + backend/airweave/platform/sources/github.py | 10 +- .../platform/sources/test_github_capture.py | 46 ++++++- 8 files changed, 217 insertions(+), 15 deletions(-) diff --git a/backend/airweave/api/examples.py b/backend/airweave/api/examples.py index 632e9c99e2..180be34e0b 100644 --- a/backend/airweave/api/examples.py +++ b/backend/airweave/api/examples.py @@ -183,9 +183,12 @@ "short_name": "github", "collection": "engineering-docs-ab123", "config_fields": { - "repo_name": "company/main-repo", - "include_issues": True, - "include_pull_requests": True, + "expected_user_id": 42, + "repositories": [ + {"repository_id": 100, "owner_id": 200, "full_name": "company/main-repo"} + ], + "include_code": True, + "include_conversations": True, }, "auth_fields": { "personal_access_token": "ghp_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx", @@ -238,12 +241,6 @@ "type": "string", "secret": True, }, - { - "name": "repo_name", - "title": "Repository Name", - "description": "Full repository name in format 'owner/repo'", - "type": "string", - }, ] }, }, diff --git a/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py index 34d6a16b3b..f10ed9813e 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_github_recovery.py @@ -137,6 +137,7 @@ async def store_blob(content, *, media_type=None): http_client=client, config=GitHubConfig.model_validate( { + "expected_user_id": 42, "repositories": [{"repository_id": 100, "owner_id": 200, "full_name": "team/repo"}], "include_code": native.code, } diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 0b081833de..5eb92c8c9f 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -28,6 +28,7 @@ class ManagedSource(BaseModel): "stripe", "linear", "attio", + "github", ] expected_identity: str = Field(min_length=1, max_length=512) expected_user_identity: str | None = Field(default=None, min_length=1, max_length=512) @@ -44,14 +45,18 @@ def native_principal(self): """Slack workspace membership requires its native user, not broker user_id.""" if (self.provider == "slack") != (self.expected_user_identity is not None): raise ValueError("Only Slack requires an expected native user identity") - if self.provider in {"stripe", "linear", "attio"}: + if self.provider in {"stripe", "linear", "attio", "github"}: self.source_config() # Validate native scope before admission, not during delivery. return self @field_validator("expected_identity") @classmethod - def workspace_identity(cls, value: str, info: ValidationInfo) -> str: - """Workspace UUIDs have one canonical spelling across admission and reconnect.""" + 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"}: return str(UUID(value)) return value @@ -66,6 +71,17 @@ def valid_schedule(cls, value: str) -> str: def source_config(self) -> dict[str, JsonValue]: """Expected identity cannot be overridden inside unstructured provider config.""" + 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 @@ -155,6 +171,10 @@ def _workspace_principal(provider: str, config: dict) -> str: def native_principal(provider: str, config: dict) -> tuple[str, str | None]: """Read the original attested principal from the protected source config.""" + 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) diff --git a/backend/airweave/domains/owned_provisioning/store.py b/backend/airweave/domains/owned_provisioning/store.py index 3c24e32fbc..6becbd27a5 100644 --- a/backend/airweave/domains/owned_provisioning/store.py +++ b/backend/airweave/domains/owned_provisioning/store.py @@ -220,6 +220,17 @@ async def _configure( 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 diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index e903dad390..6660c6004d 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -797,3 +797,121 @@ def test_workspace_provisioning_rejects_invalid_intent(provider, identity, confi 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 * * * *", + ) diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index caf384cf1b..e4c53a74b5 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -203,6 +203,11 @@ 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 diff --git a/backend/airweave/platform/sources/github.py b/backend/airweave/platform/sources/github.py index ce69f91dfc..80d1410707 100644 --- a/backend/airweave/platform/sources/github.py +++ b/backend/airweave/platform/sources/github.py @@ -28,6 +28,7 @@ parent_container_key, ) 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 @@ -127,7 +128,14 @@ async def create( async def validate(self) -> None: """Verify authentication independently from permission to every selected repository.""" - Identity.model_validate(await self.reader.object("/user")) + 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, + ) def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: """Use one bounded, stable parent identity policy at every hierarchy depth.""" diff --git a/backend/tests/unit/platform/sources/test_github_capture.py b/backend/tests/unit/platform/sources/test_github_capture.py index a08984a375..71bdfc2819 100644 --- a/backend/tests/unit/platform/sources/test_github_capture.py +++ b/backend/tests/unit/platform/sources/test_github_capture.py @@ -30,7 +30,10 @@ } COMMIT, TREE = "1" * 40, "2" * 40 BRANCH = {"name": "main", "commit": {"sha": COMMIT}} -CONFIG = {"repositories": [{"repository_id": 100, "owner_id": 200, "full_name": "team/repo"}]} +CONFIG = { + "expected_user_id": 42, + "repositories": [{"repository_id": 100, "owner_id": 200, "full_name": "team/repo"}], +} def response(value=None, *, status=200, headers=None, content=None): @@ -350,7 +353,7 @@ def test_legacy_and_ambiguous_selections_rejected(): with pytest.raises(ValidationError, match="Refresh GitHub"): GitHubConfig(repo_name="team/repo") with pytest.raises(ValidationError): - GitHubConfig.model_validate({"repositories": CONFIG["repositories"] * 2}) + GitHubConfig.model_validate({**CONFIG, "repositories": CONFIG["repositories"] * 2}) @pytest.mark.asyncio @@ -629,3 +632,42 @@ async def test_transferred_issue_withdraws_old_scope_without_reading_new_comment assert error.value.removal_reason == "scope_removed" await capture.confirm_absent(issue) assert all("/comments" not in call.args[0] for call in client.get.call_args_list) + + +@pytest.mark.asyncio +async def test_validate_attests_native_user_and_identity_binds_capture_configuration(): + connector, client = await source(response({"id": 42, "login": "renamed-user"})) + await connector.validate() + assert client.get.call_args.args[0] == "https://api.github.com/user" + other, _ = await source(config={**CONFIG, "expected_user_id": 43}) + assert ( + connector.capture_cycle_configuration.fingerprint + != other.capture_cycle_configuration.fingerprint + ) + + +@pytest.mark.asyncio +async def test_validate_rejects_different_native_user_before_repository_reads(): + from airweave.domains.sources.exceptions import SourceAuthError + + connector, client = await source(response({"id": 43, "login": "same-display-label"})) + with pytest.raises(SourceAuthError, match="does not match"): + await connector.validate() + assert client.get.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", [{}, {"id": None}, {"id": "42"}, {"id": True}, {"id": 0}]) +async def test_validate_rejects_missing_or_malformed_native_identity(identity): + connector, _ = await source(response(identity)) + with pytest.raises(ValidationError): + await connector.validate() + + +@pytest.mark.parametrize("identity", [None, "42", True, 0, -1, 1.5]) +def test_config_requires_strict_positive_native_user_id(identity): + config = {**CONFIG, "expected_user_id": identity} + if identity is None: + config.pop("expected_user_id") + with pytest.raises(ValidationError): + GitHubConfig.model_validate(config) From 14dcea8e790db019a4c949a50c962348fad2565d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 10:55:09 -0700 Subject: [PATCH 241/266] Bind Notion provisioning and capture to native workspace and bot --- .../canonical/tests/test_notion_recovery.py | 2 +- .../domains/owned_provisioning/models.py | 34 ++++++-- .../tests/test_provisioning.py | 81 +++++++++++++++++ backend/airweave/platform/configs/config.py | 6 +- backend/airweave/platform/sources/notion.py | 28 +++++- .../platform/sources/test_notion_capture.py | 86 ++++++++++++++++++- 6 files changed, 225 insertions(+), 12 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py index d463093478..8ba8e657fa 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_recovery.py @@ -100,7 +100,7 @@ async def runner(database, source, fixture, attempt=1, files=None): auth=StaticTokenProvider("fixture"), logger=MagicMock(), http_client=client, - config=NotionConfig(), + config=NotionConfig(expected_workspace_id=UUID(int=100), expected_bot_id=UUID(int=101)), ) ctx, _, runtime, bus = components(database, source) pipeline = CanonicalCapturePipeline( diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 5eb92c8c9f..191123dc71 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -29,6 +29,7 @@ class ManagedSource(BaseModel): "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) @@ -42,10 +43,10 @@ class ManagedSource(BaseModel): @model_validator(mode="after") def native_principal(self): - """Slack workspace membership requires its native user, not broker user_id.""" - if (self.provider == "slack") != (self.expected_user_identity is not None): - raise ValueError("Only Slack requires an expected native user identity") - if self.provider in {"stripe", "linear", "attio", "github"}: + """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 @@ -57,7 +58,15 @@ def canonical_identity(cls, value: str, info: ValidationInfo) -> str: 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"}: + 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 @@ -71,6 +80,16 @@ def valid_schedule(cls, value: str) -> str: 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 @@ -171,6 +190,11 @@ def _workspace_principal(provider: str, config: dict) -> str: 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 diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index 6660c6004d..aff9921c83 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -915,3 +915,84 @@ def test_github_provisioning_rejects_invalid_principal_or_missing_selection(iden 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/platform/configs/config.py b/backend/airweave/platform/configs/config.py index e4c53a74b5..60a4c8dfa8 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -534,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): diff --git a/backend/airweave/platform/sources/notion.py b/backend/airweave/platform/sources/notion.py index f9ec3ab78f..040a4924da 100644 --- a/backend/airweave/platform/sources/notion.py +++ b/backend/airweave/platform/sources/notion.py @@ -35,7 +35,7 @@ RecordIdentity, ) from airweave.domains.entities.canonical.scan_models import ScanContinuation -from airweave.domains.sources.exceptions import SourceError +from airweave.domains.sources.exceptions import SourceAuthError, SourceError from airweave.domains.sources.token_providers.protocol import ( SourceAuthProvider, authorization_headers, @@ -105,6 +105,16 @@ 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 @@ -232,8 +242,10 @@ async def create( ) -> NotionSource: """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"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(), @@ -282,7 +294,17 @@ async def _request( async def validate(self) -> None: """Verify identity independently of content discovery capabilities.""" payload = await self._request("users/me") - _Id.model_validate(payload) + 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, + ) async def _retrieve(self, kind: str, native_id: str) -> dict[str, JsonValue] | None: native_id = str(UUID(native_id)) diff --git a/backend/tests/unit/platform/sources/test_notion_capture.py b/backend/tests/unit/platform/sources/test_notion_capture.py index 84600cd936..b1df3a8410 100644 --- a/backend/tests/unit/platform/sources/test_notion_capture.py +++ b/backend/tests/unit/platform/sources/test_notion_capture.py @@ -64,7 +64,7 @@ async def source(*, gets=(), posts=()): ), logger=MagicMock(), http_client=client, - config=NotionConfig(), + config=NotionConfig(expected_workspace_id=UUID(int=100), expected_bot_id=UUID(int=101)), ) return result, client @@ -550,3 +550,87 @@ async def test_property_archive_roundtrips_unicode_and_empty_value(property_file assert archive["responses"] == [value] assert record.payload["name"] == 'Notes "日本語"' assert archive["property_id"] == "p%25" + + +def principal(**overrides): + return { + "object": "user", + "type": "bot", + "id": str(UUID(int=101)), + "bot": {"workspace_id": str(UUID(int=100))}, + **overrides, + } + + +@pytest.mark.asyncio +async def test_validate_exact_workspace_bot_pair_and_fingerprint(): + connector, client = await source(gets=[principal(name="Editable label")]) + await connector.validate() + assert client.get.call_args.args[0] == "https://api.notion.com/v1/users/me" + for field in ("expected_workspace_id", "expected_bot_id"): + changed = await NotionSource.create( + auth=connector.auth, + logger=MagicMock(), + http_client=AsyncMock(), + config=connector.config.model_copy(update={field: UUID(int=102)}), + ) + assert ( + changed.capture_cycle_configuration.fingerprint + != connector.capture_cycle_configuration.fingerprint + ) + assert set(connector.capture_cycle_configuration.completion_policies.values()) == { + "discovery_with_validation" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [principal(id=str(UUID(int=102))), principal(bot={"workspace_id": str(UUID(int=102))})], +) +async def test_validate_rejects_other_workspace_or_bot(payload): + from airweave.domains.sources.exceptions import SourceAuthError + + connector, client = await source(gets=[payload]) + with pytest.raises(SourceAuthError, match="does not match") as caught: + await connector.validate() + assert caught.value.status_code == 200 + assert client.get.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + {}, + principal(object="page"), + principal(type="person"), + principal(id="not-a-uuid"), + principal(bot={}), + principal(bot={"workspace_id": None}), + principal(bot={"workspace_id": "name"}), + ], +) +async def test_validate_rejects_missing_or_malformed_bot_identity(payload): + connector, _ = await source(gets=[payload]) + with pytest.raises(ValidationError): + await connector.validate() + + +@pytest.mark.parametrize( + "config", + [ + {}, + {"expected_workspace_id": str(UUID(int=100))}, + {"expected_workspace_id": "workspace name", "expected_bot_id": str(UUID(int=101))}, + {"expected_workspace_id": str(UUID(int=100)), "expected_bot_id": "bot name"}, + { + "expected_workspace_id": str(UUID(int=100)), + "expected_bot_id": str(UUID(int=101)), + "extra": True, + }, + ], +) +def test_notion_config_requires_only_verified_uuid_pair(config): + with pytest.raises(ValidationError): + NotionConfig.model_validate(config) From a1dd4f6fb7c8448a4c119f61ed35c3c270d3010a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 11:11:31 -0700 Subject: [PATCH 242/266] Project retained Notion property archives without provider fetches --- .../canonical/notion_property_projection.py | 300 +++++++++++++++++ .../entities/canonical/projection_mappers.py | 9 +- .../tests/test_notion_property_projection.py | 312 ++++++++++++++++++ backend/airweave/platform/entities/notion.py | 4 +- 4 files changed, 623 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/notion_property_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py 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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index d9bd83b276..ae9b9f77ae 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -385,7 +385,7 @@ def excluded_from_search(record: SourceRecord, source_name: str) -> bool: @asynccontextmanager -async def map_record( +async def map_record( # noqa: C901 -- explicit provider dispatch keeps mapper ownership visible record: SourceRecord, source_name: str, storage: StorageBackend, @@ -423,6 +423,13 @@ async def map_record( yield await map_slack_files(record, _slack(record)[0], storage, directory) return + elif source_name == "notion" and record.identity.record_type == "page_property": + from airweave.domains.entities.canonical.notion_property_projection import ( + map_notion_property, + ) + + yield await map_notion_property(record, storage) + return elif source_name == "google_drive": yield await _drive(record, storage, directory) return diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py new file mode 100644 index 0000000000..81a00e4e43 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py @@ -0,0 +1,312 @@ +"""Offline native properties, bounded reads and retained identity.""" + +import hashlib +import json +from unittest.mock import AsyncMock + +import pytest + +from airweave.domains.entities.canonical.projection_mappers import ( + ProjectionMappingError, + map_record, +) +from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity +from airweave.domains.entities.canonical.tests.test_notion_projection import PAGE, STAMP, original + + +def item(kind, value): + return {"object": "property_item", "id": "a%3Ab", "type": kind, kind: value} + + +def page(kind, values, more=False, calculation=None): + metadata = {"id": "a%3Ab", "type": kind, "next_url": "next" if more else None} + if calculation is not None: + metadata[kind] = calculation + return { + "object": "list", + "type": "property_item", + "results": values, + "property_item": metadata, + "has_more": more, + "next_cursor": "next" if more else None, + } + + +def fixture(kind, responses, status="available"): + record = original() + header = { + "format_version": 1, + "page_id": PAGE, + "property_id": "a%3Ab", + "notion_version": "2026-03-11", + "page_last_edited_time": STAMP, + } + content = json.dumps({**header, "responses": responses}, ensure_ascii=False).encode() + digest = hashlib.sha256(content).hexdigest() + blob = BlobReference( + key=f"canonical/{record.sync_id}/blobs/sha256/{digest}", + sha256=digest, + size_bytes=len(content), + media_type="application/json", + ) + record = record.model_copy( + update={ + "identity": RecordIdentity( + record_type="page_property", native_id="a%3Ab", container_id=PAGE + ), + "parent": RecordIdentity(record_type="page", native_id=PAGE), + "payload_schema_version": 2, + "payload": { + **header, + "name": "Résumé", + "property": {"id": "a%3Ab", "type": kind}, + "response_count": len(responses), + "value_status": status, + }, + "blobs": (blob,), + } + ) + return record, content + + +async def project(record, content): + storage = AsyncMock() + storage.read_file.return_value = content + async with map_record(record, "notion", storage) as result: + return result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind,value,expected", + [ + ("number", 12, "12"), + ("number", 9007199254740993, "9007199254740993"), + ("checkbox", False, "false"), + ("select", {"name": "準備"}, "準備"), + ("date", {"start": "2026-10-01"}, "2026-10-01"), + ("people", {"id": PAGE, "name": "Mira"}, f"Mira ({PAGE})"), + ("relation", {"id": PAGE}, PAGE), + ( + "files", + [{"name": "Report.pdf", "file": {"url": "https://signed.example/secret"}}], + "Report.pdf", + ), + ("formula", {"type": "string", "string": "計画"}, "計画"), + ], +) +async def test_native_values(kind, value, expected): + responses = ( + [page(kind, [item(kind, value)])] if kind in {"people", "relation"} else [item(kind, value)] + ) + record, content = fixture(kind, responses) + result = await project(record, content) + entity = result.entities[0] + assert entity.text == expected and entity.native_id == "a%3Ab" + assert entity.created_at is None and entity.updated_at is None + assert result.parts[0].native_body is None + assert "exclude linked bodies" in entity.content_coverage + + +@pytest.mark.asyncio +async def test_paginated_text_and_terminal_rollup(): + record, content = fixture( + "rich_text", + [ + page("rich_text", [item("rich_text", {"type": "text", "plain_text": "你好 "})], True), + page("rich_text", [item("rich_text", {"type": "mention", "plain_text": "Mira"})]), + ], + ) + assert (await project(record, content)).entities[0].text == "你好 Mira" + record, content = fixture( + "rollup", + [ + page("rollup", [], True, {"type": "incomplete", "function": "sum"}), + page("rollup", [], calculation={"type": "number", "number": 9, "function": "sum"}), + ], + ) + assert (await project(record, content)).entities[0].text == "9" + + +@pytest.mark.asyncio +async def test_unsupported_and_incomplete_distinct(): + record, content = fixture("formula", [item("formula", {"type": "unsupported"})], "unsupported") + result = await project(record, content) + assert not result.entities and result.parts[0].omission == "unsupported_format" + record, content = fixture("rollup", [page("rollup", [], calculation={"type": "incomplete"})]) + with pytest.raises(ProjectionMappingError, match="incomplete"): + await project(record, content) + + +@pytest.mark.asyncio +async def test_authority_hash_and_pagination_fail_closed(): + record, content = fixture("number", [item("number", 1)]) + for change in [ + {"parent": None}, + {"payload_schema_version": 1}, + {"payload": {**record.payload, "response_count": 2}}, + {"payload": {**record.payload, "page_last_edited_time": "2026-10-02T00:00:00Z"}}, + ]: + with pytest.raises(ValueError): + await project(record.model_copy(update=change), content) + with pytest.raises(ValueError): + await project(record, content + b" ") + record, content = fixture("rich_text", [page("rich_text", [], True)]) + with pytest.raises(ProjectionMappingError, match="pagination"): + await project(record, content) + + +@pytest.mark.asyncio +async def test_budget_before_storage(): + record, _ = fixture("number", [item("number", 1)]) + record = record.model_copy( + update={"blobs": (record.blobs[0].model_copy(update={"size_bytes": 16 * 1024 * 1024 + 1}),)} + ) + storage = AsyncMock() + with pytest.raises(ProjectionMappingError, match="projection budget"): + async with map_record(record, "notion", storage): + pass + storage.read_file.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind,responses", + [ + ("future_type", [item("future_type", None)]), + ("number", [page("number", [])]), + ("people", [item("people", {"id": PAGE})]), + ("checkbox", [item("checkbox", None)]), + ], +) +async def test_unknown_or_invalid_native_shape_fails(kind, responses): + record, content = fixture(kind, responses) + with pytest.raises(ValueError): + await project(record, content) + + +@pytest.mark.integration +async def test_property_publication_obeys_current_parent(database, source): + from uuid import uuid4 + + from sqlalchemy import select + + 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.tests.helpers import ( + capture, + observation, + publish_prepared, + ) + from airweave.models import Entity, Sync + + service, fence = source + record, content = fixture("number", [item("number", 1)]) + blob = record.blobs[0].model_copy( + update={"key": f"canonical/{fence.sync_id}/blobs/sha256/{record.blobs[0].sha256}"} + ) + root = observation(identity=record.parent, payload={"id": PAGE}) + await capture( + database, + service, + fence, + root, + observation( + identity=record.identity, + parent=record.parent, + payload=record.payload, + payload_schema_version=2, + blobs=(blob,), + completeness="partial", + ), + ) + store = CanonicalProjectionStore() + async with database() as db: + work = next( + w + for w in await store.pending(db, fence.organization_id, fence.sync_id) + if w.record.identity.record_type == "page_property" + ) + assert (await project(work.record, content)).entities[0].text == "1" + generation = uuid4() + async with database() as db: + assert await publish_prepared(store, db, work, generation, 1) + locator = ProjectionLocator( + record_id=work.record.id, + revision=work.record.revision, + pipeline_version=work.pipeline_version, + generation=generation, + part_index=0, + ) + await capture( + database, + service, + fence, + root.model_copy(update={"kind": "delete", "removal_reason": "access_revoked"}), + ) + 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(), 1) + + +@pytest.mark.integration +async def test_unsupported_property_publishes_explicit_extraction(database, source, tmp_path): + from sqlalchemy import select + + from airweave.adapters.storage.filesystem import FilesystemBackend + from airweave.domains.entities.canonical.projection_store import ( + CanonicalProjectionStore, + current_extraction, + ) + from airweave.domains.entities.canonical.tests.helpers import capture, observation + from airweave.domains.entities.canonical.tests.test_extraction_coverage import ( + destination, + logger, + projector, + ) + from airweave.models import Entity + + service, fence = source + record, content = fixture("formula", [item("formula", {"type": "unsupported"})], "unsupported") + blob = record.blobs[0].model_copy( + update={"key": f"canonical/{fence.sync_id}/blobs/sha256/{record.blobs[0].sha256}"} + ) + storage = FilesystemBackend(tmp_path) + await storage.write_file(blob.key, content) + await capture( + database, + service, + fence, + observation(identity=record.parent, payload={"id": PAGE}), + observation( + identity=record.identity, + parent=record.parent, + payload=record.payload, + payload_schema_version=2, + blobs=(blob,), + completeness="partial", + ), + ) + async with database() as db: + work = next( + w + for w in await CanonicalProjectionStore().pending( + db, fence.organization_id, fence.sync_id + ) + if w.record.identity.record_type == "page_property" + ) + assert await projector(database, storage).project_one(work, "notion", destination(), logger) + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.id == work.record.id)) + coverage = await current_extraction( + db, fence.organization_id, fence.sync_id, row.id, row.record_revision + ) + assert row.indexed_chunk_count == 0 + assert coverage.parts[0].outcome == "unsupported" + assert coverage.status == "unavailable" diff --git a/backend/airweave/platform/entities/notion.py b/backend/airweave/platform/entities/notion.py index 06d911481d..d8f6c99eae 100644 --- a/backend/airweave/platform/entities/notion.py +++ b/backend/airweave/platform/entities/notion.py @@ -561,7 +561,9 @@ def web_url(self) -> str: class NotionOriginalEntity(BaseEntity): """Searchable metadata or block-local text from one retained Notion original.""" - native_id: str = AirweaveField(..., description="Native Notion UUID.", is_entity_id=True) + 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 From 572a19c32a4efff24a40d10da1deec291e4dcafa Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 11:52:54 -0700 Subject: [PATCH 243/266] Add fenced native record withdrawal and explicit access renewal --- .../api/v1/endpoints/native_imports.py | 48 ++++ .../domains/native_ingestion/README.md | 52 ++++- .../domains/native_ingestion/access_models.py | 42 ++++ .../domains/native_ingestion/access_store.py | 155 +++++++++++++ .../native_ingestion/import_service.py | 30 +++ .../domains/native_ingestion/store.py | 45 +++- .../native_ingestion/tests/test_access_api.py | 209 ++++++++++++++++++ 7 files changed, 568 insertions(+), 13 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/access_models.py create mode 100644 backend/airweave/domains/native_ingestion/access_store.py create mode 100644 backend/airweave/domains/native_ingestion/tests/test_access_api.py diff --git a/backend/airweave/api/v1/endpoints/native_imports.py b/backend/airweave/api/v1/endpoints/native_imports.py index aab69ff2b9..597966ef1c 100644 --- a/backend/airweave/api/v1/endpoints/native_imports.py +++ b/backend/airweave/api/v1/endpoints/native_imports.py @@ -13,6 +13,7 @@ from airweave.core.container import Container from airweave.db.session import get_db from airweave.domains.entities.canonical.store import CanonicalStoreError, WriterBusy +from airweave.domains.native_ingestion.access_models import NativeAccessChange, NativeRecordAccess from airweave.domains.native_ingestion.errors import NativeAdmissionError, NativeImportNotFound from airweave.domains.native_ingestion.import_models import NativeImportState, StartNativeImport from airweave.domains.native_ingestion.page_models import CommitNativePage, NativePageAck @@ -178,3 +179,50 @@ async def cancel_native_import( raise HTTPException(404, {"code": error.code, "message": str(error)}) from error except CanonicalStoreError as error: raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.get( + "/{source_id}/imports/{request_key}/records/{record_id}/access", + response_model=NativeRecordAccess, +) +async def read_access( + source_id: UUID, + request_key: RequestKey, + record_id: UUID, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeRecordAccess: + """Read access metadata under the authenticated native source binding.""" + try: + return await container.native_imports.read_access( + db, ctx.organization.id, source_id, request_key, record_id + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error + + +@router.post( + "/{source_id}/imports/{request_key}/records/{record_id}/access", + response_model=NativeRecordAccess, +) +async def change_access( + source_id: UUID, + request_key: RequestKey, + record_id: UUID, + request: NativeAccessChange, + db: AsyncSession = Depends(get_db), + ctx: ApiContext = Depends(backend_actor), + container: Container = Depends(deps.get_container), +) -> NativeRecordAccess: + """Accept explicit fresh source evidence; never restore from an old page retry.""" + try: + return await container.native_imports.change_access( + db, ctx.organization.id, source_id, request_key, record_id, request + ) + except NativeImportNotFound as error: + raise HTTPException(404, {"code": error.code, "message": str(error)}) from error + except CanonicalStoreError as error: + raise HTTPException(409, {"code": error.code, "message": str(error)}) from error diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 7a8d6c16eb..2b5de8cc3e 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -35,8 +35,8 @@ 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; explicit renewal needs a separately -designed access-authority operation. Parent access gates remain canonical-owned. +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 @@ -66,8 +66,8 @@ 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. Explicit access renewal and authoritative publisher -completeness remain pending. Native jobs cannot be cancelled through the generic +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 @@ -115,8 +115,9 @@ 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 access -renewal remains a separate unfinished capability. +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 @@ -172,3 +173,42 @@ that import's cycle ID; once reused, the API reports unavailable/superseded inst 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. 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/import_service.py b/backend/airweave/domains/native_ingestion/import_service.py index 97b0094b6b..a060ecc8da 100644 --- a/backend/airweave/domains/native_ingestion/import_service.py +++ b/backend/airweave/domains/native_ingestion/import_service.py @@ -5,6 +5,8 @@ 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 @@ -25,6 +27,7 @@ class NativeImports: 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)) @@ -109,3 +112,30 @@ async def finish( 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/store.py b/backend/airweave/domains/native_ingestion/store.py index 575d0a3d26..02062a1476 100644 --- a/backend/airweave/domains/native_ingestion/store.py +++ b/backend/airweave/domains/native_ingestion/store.py @@ -41,7 +41,7 @@ def compare_versions(previous: NativeVersion, incoming: NativeVersion) -> int: return -1 if lower else 1 if higher else 0 -def _admit(entity: Entity, snapshot: NativeSnapshot) -> bool: +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) @@ -55,9 +55,9 @@ def _admit(entity: Entity, snapshot: NativeSnapshot) -> bool: if comparison == 0: if previous != snapshot: raise NativeAdmissionError("Native snapshot version has conflicting content") - # A retry acknowledges retention, never restores local visibility. - return False - if entity.removal_reason in ("access_revoked", "scope_removed", "absent"): + # 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 @@ -110,6 +110,7 @@ def _admit_snapshots( request: IngestNativeBatch, *, sightings: bool, + revalidated_ids: tuple[UUID, ...] = (), ) -> NativeAdmission: """Apply version and membership rules to one preloaded bounded batch.""" accepted = [] @@ -121,7 +122,7 @@ def _admit_snapshots( 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): + 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: @@ -157,7 +158,13 @@ async def ingest(self, db: AsyncSession, request: IngestNativeBatch) -> CaptureR ) async def admit_locked( - self, db: AsyncSession, sync: Sync, request: IngestNativeBatch, *, sightings: bool = False + 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 ( @@ -205,8 +212,32 @@ async def admit_locked( ) 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 + 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( 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"] From 1784e39fcad12e62c54c6d40708378f13833032b Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 12:13:02 -0700 Subject: [PATCH 244/266] Preserve native filenames on retained source blobs --- .../domains/entities/canonical/requests.py | 6 ++++ .../platform/sources/gmail_capture.py | 10 +++++- .../sources/records/google_drive_content.py | 12 +++++++ .../platform/sources/slack_content.py | 3 +- .../sources/tests/test_gmail_capture.py | 2 ++ .../records/test_google_drive_content.py | 31 +++++++++++++++---- .../platform/sources/test_slack_content.py | 2 ++ 7 files changed, 58 insertions(+), 8 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/requests.py b/backend/airweave/domains/entities/canonical/requests.py index 1674938642..5ea1ed87a3 100644 --- a/backend/airweave/domains/entities/canonical/requests.py +++ b/backend/airweave/domains/entities/canonical/requests.py @@ -65,6 +65,10 @@ class BlobReference(BaseModel): sha256: str = Field(pattern=r"^[a-f0-9]{64}$") size_bytes: int = Field(ge=0) media_type: str | None = None + filename: str | None = Field( + default=None, + description="Provider filename or format-specific export label; never a storage path", + ) source_path: str | None = Field( default=None, description="RFC 6901 JSON Pointer into the unchanged native payload" ) @@ -82,6 +86,8 @@ def manifest_is_not_native_location(self) -> "BlobReference": def historical_descriptor(self, handler): """Preserve historical hashes without dropping unrelated nullable fields.""" value = handler(self) + if self.filename is None: + value.pop("filename", None) if self.role is None: value.pop("role", None) return value diff --git a/backend/airweave/platform/sources/gmail_capture.py b/backend/airweave/platform/sources/gmail_capture.py index 322d70ac89..b0d39d0332 100644 --- a/backend/airweave/platform/sources/gmail_capture.py +++ b/backend/airweave/platform/sources/gmail_capture.py @@ -184,7 +184,15 @@ async def capture_bodies( except FileSkippedException: complete = False continue - blobs.append(blob.model_copy(update={"source_path": path})) + 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: diff --git a/backend/airweave/platform/sources/records/google_drive_content.py b/backend/airweave/platform/sources/records/google_drive_content.py index 876d29cfc4..58b7e2fdff 100644 --- a/backend/airweave/platform/sources/records/google_drive_content.py +++ b/backend/airweave/platform/sources/records/google_drive_content.py @@ -47,6 +47,7 @@ async def capture_file_content( 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( @@ -77,6 +78,17 @@ async def capture_file_content( ) +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 diff --git a/backend/airweave/platform/sources/slack_content.py b/backend/airweave/platform/sources/slack_content.py index f49d83c728..70ffab8aa4 100644 --- a/backend/airweave/platform/sources/slack_content.py +++ b/backend/airweave/platform/sources/slack_content.py @@ -31,6 +31,7 @@ class SlackFileMetadata(BaseModel): 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 @@ -146,7 +147,7 @@ async def _download( ) 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, None + return blob.model_copy(update={"filename": file.name or None}), None async def _enrich( diff --git a/backend/airweave/platform/sources/tests/test_gmail_capture.py b/backend/airweave/platform/sources/tests/test_gmail_capture.py index 62ada414ef..ac72942306 100644 --- a/backend/airweave/platform/sources/tests/test_gmail_capture.py +++ b/backend/airweave/platform/sources/tests/test_gmail_capture.py @@ -177,6 +177,7 @@ async def test_external_mime_bytes_are_stored_before_complete_record_without_pay "parts": [ { "mimeType": "text/plain", + "filename": "नोट्स.txt", "body": { "attachmentId": "body-id", "size": 3, @@ -193,6 +194,7 @@ async def test_external_mime_bytes_are_stored_before_complete_record_without_pay 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") diff --git a/backend/tests/unit/platform/sources/records/test_google_drive_content.py b/backend/tests/unit/platform/sources/records/test_google_drive_content.py index 41a6eaa910..b4d25f303a 100644 --- a/backend/tests/unit/platform/sources/records/test_google_drive_content.py +++ b/backend/tests/unit/platform/sources/records/test_google_drive_content.py @@ -35,6 +35,7 @@ async def test_complete_body_keeps_original_metadata_and_blob(inputs): ) result = await capture_file_content(record, **inputs) assert result.completeness == "complete" and result.blobs[0].size_bytes == 3 + assert result.blobs[0].filename == "notes" assert result.payload == record.payload and result.content_hash == "a" * 64 assert "file%2Fid" in inputs["files"].capture_canonical_url.call_args.kwargs["url"] @@ -64,15 +65,18 @@ async def test_storage_failure_is_not_silently_a_partial_success(inputs): async def test_native_spreadsheet_keeps_existing_export_only_format(inputs): record = file_record( - {"id": "doc", "mimeType": "application/vnd.google-apps.spreadsheet", "version": "7"} + { + "id": "doc", + "mimeType": "application/vnd.google-apps.spreadsheet", + "version": "7", + "name": "बजट", + } ) - await capture_file_content(record, **inputs) + result = await capture_file_content(record, **inputs) + assert result.blobs[0].filename == "बजट.xlsx" call = inputs["files"].capture_canonical_url.call_args.kwargs assert "/export?" in call["url"] - assert ( - call["media_type"] - == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" - ) + assert call["media_type"] == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" @pytest.mark.parametrize( @@ -99,3 +103,18 @@ async def test_provider_export_limit_is_narrow_and_preserves_native_payload(inpu assert result is record and result.completeness == "metadata_only" and not result.blobs assert result.payload == record.payload inputs["get"].assert_not_awaited() + + +def test_unnamed_blob_preserves_historical_descriptor_and_named_blob_roundtrips(): + old = { + "key": "blob", + "sha256": "a" * 64, + "size_bytes": 3, + "media_type": None, + "source_path": None, + } + blob = BlobReference.model_validate(old) + assert blob.model_dump(mode="json") == old + named = blob.model_copy(update={"filename": "../बजट.pdf"}) + assert BlobReference.model_validate_json(named.model_dump_json()).filename == "../बजट.pdf" + assert named.key == blob.key and named.sha256 == blob.sha256 diff --git a/backend/tests/unit/platform/sources/test_slack_content.py b/backend/tests/unit/platform/sources/test_slack_content.py index 6f9ffbb501..b1ae8e48a4 100644 --- a/backend/tests/unit/platform/sources/test_slack_content.py +++ b/backend/tests/unit/platform/sources/test_slack_content.py @@ -81,6 +81,8 @@ async def test_file_enrichment_preserves_raw_child_and_root_blob(): item = original({"id": "F1"}) result = await capture_slack_file(source, item, files) assert result.payload == {"id": "F1"} and result.completeness == "complete" + assert result.blobs[0].filename == "note.txt" + assert result.blobs[1].filename is None assert [b.source_path for b in result.blobs] == ["", None] manifest = SlackFileManifest.model_validate_json(saved[-1]) assert manifest.files[0].file["name"] == "note.txt" From a0f9687531c41c821cb2f2d5b39bd072e0c39df5 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 12:27:41 -0700 Subject: [PATCH 245/266] Index selected structured fields from native knowledge originals --- .../domains/native_ingestion/README.md | 19 ++- .../native_ingestion/knowledge_fields.py | 123 ++++++++++++++++ .../domains/native_ingestion/projection.py | 3 + .../native_ingestion/tests/test_projection.py | 131 ++++++++++++++++++ 4 files changed, 274 insertions(+), 2 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/knowledge_fields.py diff --git a/backend/airweave/domains/native_ingestion/README.md b/backend/airweave/domains/native_ingestion/README.md index 2b5de8cc3e..caf1db3b7a 100644 --- a/backend/airweave/domains/native_ingestion/README.md +++ b/backend/airweave/domains/native_ingestion/README.md @@ -77,8 +77,11 @@ 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. Other structured knowledge fields are retained but - do not yet have a curated searchable representation. + 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 @@ -212,3 +215,15 @@ No new table, migration, scheduler or source credential exists here. NativeImpor 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/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/projection.py b/backend/airweave/domains/native_ingestion/projection.py index a263162e8d..55ad69c04d 100644 --- a/backend/airweave/domains/native_ingestion/projection.py +++ b/backend/airweave/domains/native_ingestion/projection.py @@ -7,6 +7,7 @@ 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 @@ -80,6 +81,7 @@ class NativeContentEntity(BaseEntity): 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) @@ -170,6 +172,7 @@ def map_native(record: SourceRecord) -> ProjectionInputs: 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, diff --git a/backend/airweave/domains/native_ingestion/tests/test_projection.py b/backend/airweave/domains/native_ingestion/tests/test_projection.py index 01417617b3..c49cc9f3a2 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_projection.py +++ b/backend/airweave/domains/native_ingestion/tests/test_projection.py @@ -211,3 +211,134 @@ async def test_inactive_import_and_invalid_text_are_not_successful_empty_indexin 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(), 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 From 32038d64cb7eb9deee169e4523221aa21b0498f6 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 13:05:10 -0700 Subject: [PATCH 246/266] Retain Outlook calendar originals with resumable scoped capture and offline search --- .../canonical/outlook_calendar_projection.py | 225 +++++++++++ .../entities/canonical/projection_mappers.py | 7 + .../domains/entities/canonical/scan_store.py | 23 +- .../tests/test_outlook_calendar_projection.py | 221 +++++++++++ .../tests/test_outlook_calendar_recovery.py | 241 ++++++++++++ .../domains/owned_provisioning/models.py | 9 +- .../tests/test_provisioning.py | 17 +- .../source_connections/tests/test_update.py | 14 +- .../domains/source_connections/update.py | 11 +- backend/airweave/domains/syncs/service.py | 11 +- backend/airweave/platform/configs/config.py | 26 +- .../platform/entities/outlook_calendar.py | 66 ++++ .../platform/sources/outlook_calendar.py | 13 +- .../sources/outlook_calendar_capture.py | 364 ++++++++++++++++++ .../sources/outlook_calendar_models.py | 100 +++++ .../tests/test_outlook_calendar_capture.py | 281 ++++++++++++++ .../sources/tests/test_outlook_identity.py | 24 ++ 17 files changed, 1618 insertions(+), 35 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/outlook_calendar_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_recovery.py create mode 100644 backend/airweave/platform/sources/outlook_calendar_capture.py create mode 100644 backend/airweave/platform/sources/outlook_calendar_models.py create mode 100644 backend/airweave/platform/sources/tests/test_outlook_calendar_capture.py 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/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index ae9b9f77ae..fedc443d4d 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -408,6 +408,13 @@ async def map_record( # noqa: C901 -- explicit provider dispatch keeps mapper o yield await map_gmail(record, storage, directory) return + elif source_name == "outlook_calendar": + from airweave.domains.entities.canonical.outlook_calendar_projection import ( + map_outlook_calendar, + ) + + yield await map_outlook_calendar(record, directory) + return elif source_name == "outlook_mail": from airweave.domains.entities.canonical.outlook_projection import map_outlook diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 45798282d4..80e1ff0ae1 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -139,7 +139,7 @@ def effective_mode(row: CaptureScan, cycle: CaptureCycle) -> str: def publish_scope(row: CaptureScan, cycle: CaptureCycle, sequence: int, attempt_id: UUID) -> None: """Completed evidence is derived only from this row's committed terminal state.""" - if cycle.mode != "mixed" or row.phase != "complete": + if row.execution_state is None or row.phase != "complete": return execution = ScopeExecution.model_validate(row.execution_state) checkpoint = row.continuation.get(TERMINAL_CHECKPOINT_KEY) @@ -167,10 +167,15 @@ def begin_execution( ) -> None: """Previous completion survives replacement; mismatching owner/request invalidates it.""" if cycle.mode != "mixed": - if request.plan is not None: - raise ScanConflict("Scope plans require mixed execution") - row.execution_state = None - return + if request.plan is None: + row.execution_state = None + return + if ( + cycle.mode != "full" + or request.plan.mode != "full" + or request.plan.starting_checkpoint is not None + ): + raise ScanConflict("Full capture scope plans cannot request incremental execution") plan = request.plan if plan is None: raise ScanConflict("Mixed capture requires an explicit scope plan") @@ -211,10 +216,10 @@ def validate_scope_changes( raise ScanConflict("Scope changes must start at its last published checkpoint") -def attest_planned_owner(cycle: CaptureCycle, request: BeginScan, parent: Entity | None) -> None: +def attest_planned_owner(request: BeginScan, parent: Entity | None) -> None: """The post-I/O plan must still describe the captured owner used to select it.""" if request.expected_parent_epoch != (parent.visibility_epoch if parent else None) or ( - (cycle.mode == "mixed" or request.exact_parent_observation is not None) + (request.plan is not None or request.exact_parent_observation is not None) and request.expected_parent_revision != (parent.record_revision if parent else None) ): raise ScanConflict("Scope owner changed during plan selection") @@ -353,7 +358,7 @@ async def admit(self, db: AsyncSession, request: BeginScan) -> ScanAdmission: not in cycle.configuration.exact_parent_validation ): raise ScanConflict("Scope does not accept exact parent verification") - attest_planned_owner(cycle, request, parent) + attest_planned_owner(request, parent) if request.fingerprint != cycle.configuration.fingerprint: raise CycleConflict("Scan configuration differs from the active cycle") row = await self._row(db, request.fence, request.scope) @@ -427,7 +432,7 @@ async def _begin( ) -> ScanState: """Resume unchanged scans; reuse admission checks held under the same writer lock.""" cursor, cycle, parent = context or await self._begin_context(db, request) - attest_planned_owner(cycle, request, parent) + attest_planned_owner(request, parent) inventory = cycle.configuration.fresh_inventory(request.scope.record_type) if request.fingerprint != cycle.configuration.fingerprint: raise CycleConflict("Scan configuration differs from the active cycle") diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py new file mode 100644 index 0000000000..c9763e5abb --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py @@ -0,0 +1,221 @@ +"""Offline Graph event projection uses shared conversion and honest omitted-part coverage.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest + +from airweave.domains.converters.html import HtmlConverter +from airweave.domains.entities.canonical.outlook_calendar_projection import map_outlook_calendar +from airweave.domains.entities.canonical.projection_models import ProjectionWork +from airweave.domains.entities.canonical.projector import StrictProjectionTracker, _select_inputs +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.entities.canonical.tests.test_outlook_projection import source +from airweave.domains.sync_pipeline.pipeline.text_builder import TextualRepresentationBuilder + + +def event(body): + record, _ = source(b"") + parent = RecordIdentity(record_type="calendar", native_id="calendar-1") + return record.model_copy( + update={ + "identity": RecordIdentity( + record_type="event", native_id="event-1", container_id="calendar-1" + ), + "parent": parent, + "blobs": (), + "payload": { + "id": "event-1", + "changeKey": "version-1", + "type": "seriesMaster", + "start": { + "dateTime": "2026-11-01T01:30:00.0000000", + "timeZone": "Pacific Standard Time", + }, + "end": { + "dateTime": "2026-11-01T02:30:00.0000000", + "timeZone": "Pacific Standard Time", + }, + "originalStartTimeZone": "tzone://Microsoft/Custom", + "isAllDay": False, + "isCancelled": True, + "hasAttachments": False, + "subject": "Réunion बैठक", + "body": body, + "bodyPreview": "LOSSY PREVIEW MUST NOT APPEAR", + "organizer": {"emailAddress": {"name": "रोहन", "address": "r@example.com"}}, + "attendees": [ + { + "emailAddress": {"name": "Éva", "address": "e@example.com"}, + "type": "required", + "status": {"response": "accepted"}, + "unknown": "DO NOT INDEX", + } + ], + "recurrence": { + "pattern": { + "type": "weekly", + "interval": 2, + "daysOfWeek": ["monday"], + "extension": "DO NOT INDEX", + }, + "range": { + "type": "noEnd", + "startDate": "2026-10-01", + "recurrenceTimeZone": "Pacific Standard Time", + }, + }, + "location": {"displayName": "Office"}, + "categories": ["Planning"], + "nativeUnknown": {"retained": "DO NOT INDEX"}, + }, + } + ) + + +async def build(record, tmp_path): + mapped = await map_outlook_calendar(record, tmp_path) + work = ProjectionWork( + organization_id=uuid4(), record=record, pipeline_version=2, previous_generation=None + ) + selected, coverage = _select_inputs(mapped, work, "outlook_calendar", uuid4(), lambda _: True) + registry = MagicMock() + registry.for_extension.return_value = HtmlConverter() + batch = await TextualRepresentationBuilder(registry).build_with_text( + selected, + SimpleNamespace(source_short_name="outlook_calendar", logger=MagicMock()), + SimpleNamespace(entity_tracker=StrictProjectionTracker()), + native_bodies={ + item.entity.entity_id: item.native_body + for item in mapped.parts + if item.entity is not None and item.native_body is not None + }, + ) + return mapped, coverage, batch + + +@pytest.mark.parametrize( + "body_type,text", + [ + ("text", "नमस्ते café\n & exact"), + ("text", ""), + ("html", "

नमस्ते 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 " + assert mapped.entities[0].can_edit is None + assert len(mapped.parts) == 1 + + +@pytest.mark.parametrize( + "changes", + [ + {"parent": None}, + {"parent": RecordIdentity(record_type="calendar", native_id="wrong")}, + { + "parent": RecordIdentity( + record_type="calendar", native_id="calendar-1", container_id="nested" + ) + }, + { + "identity": RecordIdentity( + record_type="event", native_id="wrong", container_id="calendar-1" + ) + }, + {"completeness": "complete"}, + ], +) +async def test_identity_parent_and_unsupported_completeness_fail_before_projection( + tmp_path, changes +): + with pytest.raises(ValueError): + await map_outlook_calendar(event(None).model_copy(update=changes), tmp_path) + + +async def test_empty_html_does_not_publish_fabricated_native_body(tmp_path): + # The shared converter cannot produce retained content for empty HTML. + # Preserve its fail-closed behavior instead of declaring generated text to be the body. + with pytest.raises(ValueError, match="Required projection input could not be converted"): + await build(event({"contentType": "html", "content": ""}), tmp_path) + + +@pytest.mark.parametrize( + "body_type,web_link", + [ + ("text", "https://outlook.example/event/1"), + ("html", "https://outlook.example/event/1"), + ("text", None), + ], +) +async def test_native_web_link_survives_vespa_payload_and_search_result( + tmp_path, body_type, web_link +): + from airweave.domains.search.adapters.vector_db.vespa_client import VespaVectorDB + from airweave.platform.destinations.vespa.transformer import EntityTransformer + + record = event({"contentType": body_type, "content": "Native body"}) + record = record.model_copy(update={"payload": {**record.payload, "webLink": web_link}}) + mapped = await map_outlook_calendar(record, tmp_path) + fields = {"entity_id": "test-event", "name": "Event", "textual_representation": "body"} + EntityTransformer(logger=MagicMock())._add_payload_field(fields, mapped.entities[0]) + engine = VespaVectorDB(app=MagicMock(), logger=MagicMock(), filter_translator=MagicMock()) + result = engine._convert_hits_to_results([{"fields": fields}]).results[0] + assert result.web_url == (web_link or "") + assert mapped.entities[0].web_url == web_link diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_recovery.py new file mode 100644 index 0000000000..3a3d5f9a8b --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_recovery.py @@ -0,0 +1,241 @@ +"""Actual composed calendar source and SQL with synthetic Graph HTTP only.""" + +from unittest.mock import MagicMock +from uuid import uuid4 + +import httpx +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.coverage import capture_coverage +from airweave.domains.entities.canonical.projection_mappers import map_record +from airweave.domains.entities.canonical.store import source_record +from airweave.domains.entities.canonical.tests.test_capture_pipeline import components +from airweave.domains.entities.canonical.tests.test_mixed_scopes import next_job +from airweave.domains.sources.exceptions import SourceEntityNotFoundError +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.platform.configs.config import CalendarOccurrenceWindow, OutlookCalendarConfig +from airweave.platform.sources.outlook_calendar import OutlookCalendarSource + + +class CalendarHTTP: + def __init__(self): + self.round = 1 + self.interrupt = False + self.missing = False + self.calls = [] + + async def handle(self, request): + path = request.url.path + if path == "/v1.0/me": + return httpx.Response(200, json={"id": "mailbox"}) + assert request.headers["Prefer"] == 'IdType="ImmutableId"' + self.calls.append((path, dict(request.url.params))) + if path == "/v1.0/me/calendars": + return httpx.Response(200, json={"value": [{"id": "cal"}]}) + if path == "/v1.0/me/calendars/cal": + return httpx.Response(200, json={"id": "cal", "canViewPrivateItems": True}) + if path.endswith("/events"): + return httpx.Response(200, json={"value": [{"id": "single"}, {"id": "master"}]}) + if path.endswith("/calendarView"): + if self.interrupt: + raise ConnectionError("Synthetic phase interruption") + ids = ["single", "history" if self.round == 1 else "current"] + return httpx.Response(200, json={"value": [{"id": item} for item in ids]}) + identity = path.rsplit("/", 1)[-1] + assert identity in {"single", "master", "history", "current"} + if identity == "history" and self.missing: + return httpx.Response(404, json={"error": {"code": "ErrorItemNotFound"}}) + kind = {"single": "singleInstance", "master": "seriesMaster"}.get(identity, "occurrence") + day = "2026-03-03" if identity == "history" else "2026-03-08" + return httpx.Response( + 200, + json={ + "id": identity, + "changeKey": "stable", + "type": kind, + "start": {"dateTime": day + "T09:00:00", "timeZone": "UTC"}, + "end": {"dateTime": day + "T10:00:00", "timeZone": "UTC"}, + "isAllDay": False, + "isCancelled": False, + "hasAttachments": False, + "subject": identity, + "body": {"contentType": "text", "content": "Exact native body — नमस्ते"}, + "unknownNativeField": {"retained": True}, + }, + ) + + +async def capture(database, source, native, attempt=1): + service, fence = source + async with httpx.AsyncClient(transport=httpx.MockTransport(native.handle)) as client: + connector = await OutlookCalendarSource.create( + auth=StaticTokenProvider("synthetic"), + logger=MagicMock(), + http_client=client, + config=OutlookCalendarConfig( + capture_originals=True, + expected_principal_id="mailbox", + ), + ) + adapter = connector.capture_page_source + assert adapter is not None + ctx, _, runtime, bus = components(database, source) + pipeline = CanonicalCapturePipeline( + service, + database, + bus, + adapter.canonical_record_types, + CaptureAttempt(id=fence.attempt_id if attempt == 1 else uuid4(), number=attempt), + adapter.canonical_container_parents, + page_source=adapter, + files=MagicMock(), + ) + await pipeline.start(ctx) + + async def no_limits(): + pass + + await pipeline.run_scans(ctx, runtime, no_limits) + await pipeline.cleanup_orphaned_entities(ctx, runtime) + await pipeline.save_checkpoint(ctx, runtime) + + +async def state(database, source): + async with database() as db: + rows = ( + await db.scalars(select(Entity).where(Entity.entity_definition_short_name == "event")) + ).all() + scan = await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) + coverage = await capture_coverage(db, source[1].organization_id, (source[1].sync_id,)) + return {row.native_id: row for row in rows}, scan, coverage[source[1].sync_id] + + +@pytest.mark.parametrize("missing", [False, True]) +async def test_calendar_phase_resume_and_exact_historical_validation( + database, source, missing, monkeypatch +): + native = CalendarHTTP() + clock = {"start": "2026-03-01T00:00:00Z"} + monkeypatch.setattr( + OutlookCalendarConfig, + "resolved_window", + lambda self: CalendarOccurrenceWindow(start=clock["start"], end="2026-04-01T00:00:00Z"), + ) + await capture(database, source, native) + initial, _, coverage = await state(database, source) + assert set(initial) == {"single", "master", "history"} + assert coverage.phase == "complete" and coverage.discovery == "incomplete" + source = await next_job(database, source) + native.round, native.interrupt = 2, True + clock["start"] = "2026-03-06T00:00:00Z" + with pytest.raises(ConnectionError, match="phase interruption"): + await capture(database, source, native) + _, interrupted, coverage = await state(database, source) + assert interrupted.phase == "collecting" + assert interrupted.continuation["phase"] == "calendarView" + assert coverage.phase == "active" + sweep, window = interrupted.sweep_id, interrupted.continuation["context"]["window"] + call_boundary = len(native.calls) + native.interrupt, native.missing = False, missing + clock["start"] = "2026-03-07T00:00:00Z" + if missing: + with pytest.raises(SourceEntityNotFoundError): + await capture(database, source, native, attempt=2) + else: + await capture(database, source, native, attempt=2) + rows, scan, coverage = await state(database, source) + assert scan.sweep_id == sweep + assert scan.continuation["context"]["window"] == window + assert set(rows) == {"single", "master", "history", "current"} + assert all(row.deleted_at is None for row in rows.values()) + assert rows["single"].id == initial["single"].id + assert rows["single"].record_revision == initial["single"].record_revision + assert all(row.parent_native_id == "cal" and row.container_id == "cal" for row in rows.values()) + assert rows["history"].id == initial["history"].id + assert rows["single"].source_payload["unknownNativeField"] == {"retained": True} + resumed = native.calls[call_boundary:] + assert resumed[0][0] == "/v1.0/me/calendars" + assert not any(path.endswith("/events") for path, _ in resumed) + assert any(path.endswith("/events/history") for path, _ in resumed) + view = next(params for path, params in resumed if path.endswith("/calendarView")) + assert view["startDateTime"].startswith("2026-03-06") + assert coverage.phase == ("active" if missing else "complete") + assert scan.phase == ("reconciling" if missing else "complete") + if not missing: + async with map_record( + source_record(rows["single"]), "outlook_calendar", MagicMock() + ) as mapped: + assert any( + part.native_body and part.native_body.text == "Exact native body — नमस्ते" + for part in mapped.parts + ) + assert any(part.entity is None for part in mapped.parts) + assert scan.execution_state["published"]["request_context"]["window"] == window + assert scan.execution_state["published"]["checkpoint"] is None + assert all(row.last_seen_run_id == sweep for row in rows.values()) + + +@pytest.mark.parametrize("mode,checkpoint", [("changes", None), ("full", {"token": "unsafe"})]) +async def test_full_scope_plan_rejects_incremental_authority(database, source, mode, checkpoint): + from airweave.domains.entities.canonical.cycle_models import ProviderCheckpoint + 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_scans import begin + + with pytest.raises(ScanConflict, match="cannot request incremental"): + await begin( + database, + *source, + plan=ScopePlan( + mode=mode, + starting_checkpoint=ProviderCheckpoint(value=checkpoint) if checkpoint else None, + ), + ) + async with database() as db: + assert await db.scalar(select(CaptureScan)) is None + + +async def test_full_scope_plan_rejects_parent_changed_during_planning( + database, source, monkeypatch +): + from datetime import datetime, timezone + + from airweave.domains.entities.canonical.requests import CaptureRecord + from airweave.domains.entities.canonical.scan_store import ScanConflict + from airweave.domains.entities.canonical.tests.helpers import capture as capture_records + from airweave.platform.sources.outlook_calendar_capture import OutlookCalendarCapture + + original = OutlookCalendarCapture.prepare_scope + + async def race(self, scope, cycle, previous, *, parent, force_full): + plan = await original(self, scope, cycle, previous, parent=parent, force_full=force_full) + if parent is not None: + await capture_records( + database, + *source, + CaptureRecord( + identity=parent.identity, + payload={**parent.payload, "name": "Renamed while selecting plan"}, + observed_at=datetime.now(timezone.utc), + ), + ) + return plan + + monkeypatch.setattr(OutlookCalendarCapture, "prepare_scope", race) + native = CalendarHTTP() + with pytest.raises(ScanConflict, match="owner changed during plan selection"): + await capture(database, source, native) + async with database() as db: + assert ( + await db.scalar(select(Entity).where(Entity.entity_definition_short_name == "event")) + is None + ) + assert ( + await db.scalar(select(CaptureScan).where(CaptureScan.record_type == "event")) is None + ) + assert not any(path.endswith("/events") for path, _ in native.calls) diff --git a/backend/airweave/domains/owned_provisioning/models.py b/backend/airweave/domains/owned_provisioning/models.py index 191123dc71..05afea45a3 100644 --- a/backend/airweave/domains/owned_provisioning/models.py +++ b/backend/airweave/domains/owned_provisioning/models.py @@ -25,6 +25,7 @@ class ManagedSource(BaseModel): "google_drive", "slack", "outlook_mail", + "outlook_calendar", "stripe", "linear", "attio", @@ -119,7 +120,7 @@ def source_config(self) -> dict[str, JsonValue]: {**self.config, "expected_account_id": self.expected_identity} ) return StripeConfig(original_capture=binding).model_dump(mode="json") - if self.provider == "outlook_mail": + if self.provider in {"outlook_mail", "outlook_calendar"}: return { **self.config, "expected_principal_id": self.expected_identity, @@ -210,6 +211,7 @@ def _account_principal(provider: str, config: dict) -> tuple[str, str | None]: GmailConfig, GoogleCalendarConfig, GoogleDriveConfig, + OutlookCalendarConfig, OutlookMailConfig, SlackConfig, StripeConfig, @@ -223,8 +225,9 @@ def _account_principal(provider: str, config: dict) -> tuple[str, str | None]: identity = GoogleCalendarConfig.model_validate(config).expected_primary_calendar_id case "google_drive": identity = GoogleDriveConfig.model_validate(config).expected_permission_id - case "outlook_mail": - outlook = OutlookMailConfig.model_validate(config) + 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 diff --git a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py index aff9921c83..160e2fa6e2 100644 --- a/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py +++ b/backend/airweave/domains/owned_provisioning/tests/test_provisioning.py @@ -455,36 +455,37 @@ async def test_slack_workspace_user_survives_pause_and_rejects_different_member( 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 + database, setup, provider ): ctx, service, request, account, lifecycle, schedules, workflows = setup creator = service.store.create - creator._source_registry.get.return_value.short_name = "outlook_mail" + 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("outlook_mail", expected_config) + creator._source_validation.seed_config_result(provider, expected_config) spec = ManagedSource.model_validate( { **request.source.model_dump(), - "provider": "outlook_mail", + "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("outlook_mail", spec.source_config()) == ("native-owner", None) + assert native_principal(provider, spec.source_config()) == ("native-owner", None) with pytest.raises(ValueError, match="must capture originals"): - native_principal("outlook_mail", {"expected_principal_id": "native-owner"}) + native_principal(provider, {"expected_principal_id": "native-owner"}) with pytest.raises(ValueError): - native_principal("outlook_mail", {"capture_originals": True}) + 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 == "outlook_mail" and source.config_fields == expected_config + 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 diff --git a/backend/airweave/domains/source_connections/tests/test_update.py b/backend/airweave/domains/source_connections/tests/test_update.py index a7968652df..990840aa97 100644 --- a/backend/airweave/domains/source_connections/tests/test_update.py +++ b/backend/airweave/domains/source_connections/tests/test_update.py @@ -201,15 +201,16 @@ async def test_update_config_valid(): 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): - sc = _make_sc(short_name="outlook_mail") +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("outlook_mail", config) + 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: @@ -222,15 +223,16 @@ async def test_outlook_capture_mode_change_requires_new_source(originals: bool): 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): - sc = _make_sc(short_name="outlook_mail") +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("outlook_mail", config) + validation.seed_config_result(provider, config) svc = _build_service(sc_repo=sc_repo, source_validation=validation) if originals: diff --git a/backend/airweave/domains/source_connections/update.py b/backend/airweave/domains/source_connections/update.py index 99ee86d2dc..f1b199037a 100644 --- a/backend/airweave/domains/source_connections/update.py +++ b/backend/airweave/domains/source_connections/update.py @@ -31,7 +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 OutlookMailConfig, StripeConfig +from airweave.platform.configs.config import OutlookCalendarConfig, OutlookMailConfig, StripeConfig from airweave.schemas.source_connection import ( AuthenticationMethod, ScheduleConfig, @@ -44,10 +44,13 @@ 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 != "outlook_mail": + if source.short_name not in {"outlook_mail", "outlook_calendar"}: return - previous = OutlookMailConfig.model_validate(source.config_fields or {}) - proposed = OutlookMailConfig.model_validate(config) + 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, diff --git a/backend/airweave/domains/syncs/service.py b/backend/airweave/domains/syncs/service.py index f1c81813cb..8ed9a5bab9 100644 --- a/backend/airweave/domains/syncs/service.py +++ b/backend/airweave/domains/syncs/service.py @@ -59,7 +59,7 @@ TemporalScheduleServiceProtocol, TemporalWorkflowServiceProtocol, ) -from airweave.platform.configs.config import OutlookMailConfig, StripeConfig +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 @@ -135,9 +135,14 @@ async def create( canonical = isinstance( source_entry.source_class_ref, (CanonicalSource, CanonicalPageSource) ) - if source_entry.short_name == "outlook_mail": + if source_entry.short_name in {"outlook_mail", "outlook_calendar"}: # Outlook composes its canonical adapter only for explicit original capture. - canonical = OutlookMailConfig.model_validate(source_config or {}).capture_originals + 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 diff --git a/backend/airweave/platform/configs/config.py b/backend/airweave/platform/configs/config.py index 60a4c8dfa8..4d90d582ed 100644 --- a/backend/airweave/platform/configs/config.py +++ b/backend/airweave/platform/configs/config.py @@ -580,7 +580,31 @@ class OracleConfig(SourceConfig): class OutlookCalendarConfig(SourceConfig): - """Outlook Calendar configuration schema.""" + """Owned calendar originals with bounded occurrence expansion on a fresh source.""" + + 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, 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/sources/outlook_calendar.py b/backend/airweave/platform/sources/outlook_calendar.py index c7d1703176..f3d7fab243 100644 --- a/backend/airweave/platform/sources/outlook_calendar.py +++ b/backend/airweave/platform/sources/outlook_calendar.py @@ -67,6 +67,10 @@ 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, @@ -81,7 +85,14 @@ async def create( instance.graph = OutlookGraphClient( auth, http_client, cls.short_name, config.expected_principal_id ) - await instance.graph.verify_principal() + 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 # ------------------------------------------------------------------ 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/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 index c746b331ba..3b550f9758 100644 --- a/backend/airweave/platform/sources/tests/test_outlook_identity.py +++ b/backend/airweave/platform/sources/tests/test_outlook_identity.py @@ -317,3 +317,27 @@ def graph(request): 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}) From 45e959c080f69faec881dce61905699bfb7d61f3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 13:21:40 -0700 Subject: [PATCH 247/266] Capture verified Slack file children before history completes --- .../domains/entities/canonical/cycle_store.py | 63 +++++++- .../domains/entities/canonical/service.py | 9 +- .../canonical/tests/test_early_child_work.py | 143 ++++++++++++++++++ .../tests/test_slack_file_recovery.py | 123 +++++++++++++++ .../domains/sync_pipeline/canonical_scan.py | 106 +++++++++---- 5 files changed, 407 insertions(+), 37 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_early_child_work.py diff --git a/backend/airweave/domains/entities/canonical/cycle_store.py b/backend/airweave/domains/entities/canonical/cycle_store.py index f3e3ebf313..31255b0bdb 100644 --- a/backend/airweave/domains/entities/canonical/cycle_store.py +++ b/backend/airweave/domains/entities/canonical/cycle_store.py @@ -227,6 +227,42 @@ def membership_ready(fence: WriterFence, state: CaptureCycle): return and_(inventory_complete(Entity, fence, state), ~stale_ancestor) +def work_membership_ready(fence: WriterFence, state: CaptureCycle): + """Early work needs a seen exact-read owner; completion still requires full inventories.""" + if state.mode != "full" or not state.configuration.exact_parent_validation: + return membership_ready(fence, state) + owner = aliased(Entity) + early = exists( + select(CaptureScan.id) + .join(owner, owner.id == CaptureScan.parent_record_id) + .where( + Entity.entity_definition_short_name.in_(state.configuration.exact_parent_validation), + CaptureScan.organization_id == fence.organization_id, + CaptureScan.sync_id == fence.sync_id, + CaptureScan.cycle_id == state.version.cycle_id, + CaptureScan.record_type == Entity.entity_definition_short_name, + CaptureScan.container_id.is_not_distinct_from(Entity.container_id), + CaptureScan.phase.in_(("collecting", "reconciling")), + Entity.last_seen_run_id == CaptureScan.sweep_id, + owner.organization_id == Entity.organization_id, + owner.sync_id == Entity.sync_id, + owner.entity_definition_short_name == Entity.parent_record_type, + owner.native_id == Entity.parent_native_id, + owner.container_id.is_not_distinct_from(Entity.parent_container_id), + CaptureScan.parent_visibility_epoch == owner.visibility_epoch, + Entity.parent_visibility_epoch == owner.visibility_epoch, + ) + ) + chain = ancestor_chain() + ancestor = aliased(Entity) + stale_ancestor = exists( + select(chain.c.id) + .join(ancestor, ancestor.id == chain.c.id) + .where(~inventory_complete(ancestor, fence, state)) + ) + return and_(or_(inventory_complete(Entity, fence, state), early), ~stale_ancestor) + + async def root_ready(db: AsyncSession, fence: WriterFence, state: CaptureCycle) -> bool: """Every declared root inventory completes before descendant work.""" for kind in state.configuration.root_record_types: @@ -302,7 +338,7 @@ async def attest_scope( raise CycleConflict("Child scope no longer has a visible parent") ready = await db.scalar( select(Entity.id).where( - Entity.id == parent.id, content_is_available(), membership_ready(fence, state) + Entity.id == parent.id, content_is_available(), work_membership_ready(fence, state) ) ) if ready is None or not await root_ready(db, fence, state): @@ -353,10 +389,12 @@ def child_scope_complete(fence: WriterFence, state: CaptureCycle, record_type: s return exists(select(CaptureScan.id).where(*predicates)) -async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) -> ScopeWork | None: +async def next_scope_work( + db: AsyncSession, fence: WriterFence, cycle_id: UUID, *, within: CompletedScope | None = None +) -> ScopeWork | None: """Reevaluate SQL frontier each time; newly discovered earlier UUIDs cannot be skipped.""" _, state = await attest_cycle(db, fence, cycle_id) - for kind in state.configuration.root_record_types: + for kind in () if within is not None else state.configuration.root_record_types: predicates = [ CaptureScan.organization_id == fence.organization_id, CaptureScan.sync_id == fence.sync_id, @@ -370,6 +408,22 @@ async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) predicates.append(CaptureScan.membership_attempt_id == fence.attempt_id) if await db.scalar(select(CaptureScan.id).where(*predicates).limit(1)) is None: return ScopeWork(record_type=kind) + selection = [] + if within is not None: + await attest_scope(db, fence, state, within) + owner = await scope_owner(db, fence, state, within) + identity = source_record(owner).identity if owner is not None else None + selection = [ + Entity.entity_definition_short_name == within.record_type, + Entity.container_id.is_not_distinct_from(within.container_id), + Entity.parent_record_type.is_not_distinct_from( + identity.record_type if identity else None + ), + Entity.parent_native_id.is_not_distinct_from(identity.native_id if identity else None), + Entity.parent_container_id.is_not_distinct_from( + identity.container_id if identity else None + ), + ] for child_type, allowed in state.configuration.parents.items(): parent_types = tuple(kind for kind in allowed if kind is not None) if not parent_types: @@ -381,7 +435,8 @@ async def next_scope_work(db: AsyncSession, fence: WriterFence, cycle_id: UUID) Entity.sync_id == fence.sync_id, Entity.entity_definition_short_name.in_(parent_types), content_is_available(), - membership_ready(fence, state), + work_membership_ready(fence, state), + *selection, ~child_scope_complete(fence, state, child_type), ) .order_by(Entity.id) diff --git a/backend/airweave/domains/entities/canonical/service.py b/backend/airweave/domains/entities/canonical/service.py index cccc879f52..5cdaf16276 100644 --- a/backend/airweave/domains/entities/canonical/service.py +++ b/backend/airweave/domains/entities/canonical/service.py @@ -119,12 +119,17 @@ async def save_checkpoint( await self.store.save_checkpoint(db, fence, cursor_data) async def next_scope_work( - self, db: AsyncSession, fence: WriterFence, cycle_id: UUID + self, + db: AsyncSession, + fence: WriterFence, + cycle_id: UUID, + *, + within: CompletedScope | None = None, ) -> ScopeWork | None: """Read the next eligible incomplete scope under the existing writer fence.""" async with UnitOfWork(db): await self.store._fenced_sync(db, fence) - return await next_scope_work(db, fence, cycle_id) + return await next_scope_work(db, fence, cycle_id, within=within) async def withdraw_scan_parent( self, diff --git a/backend/airweave/domains/entities/canonical/tests/test_early_child_work.py b/backend/airweave/domains/entities/canonical/tests/test_early_child_work.py new file mode 100644 index 0000000000..62690c2f4d --- /dev/null +++ b/backend/airweave/domains/entities/canonical/tests/test_early_child_work.py @@ -0,0 +1,143 @@ +"""Early exact-owner work is not completed inventory or a replacement discovery queue.""" + +import pytest +from sqlalchemy import select + +from airweave.domains.entities.canonical.cycle_models import BeginCycle, CompleteCycle +from airweave.domains.entities.canonical.cycle_store import CycleConflict, next_scope_work +from airweave.domains.entities.canonical.requests import CompletedScope, parent_container_key +from airweave.domains.entities.canonical.scan_models import BeginScan +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.entity import Entity + + +async def setup(database, source): + 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) + message = item("message", "M", channel.identity) + scope = CompletedScope( + record_type="message", container_id=message.identity.container_id, parent=channel.identity + ) + state = await scan(database, service, fence, cycle, scope) + state = await page(database, service, fence, state, message) + async with database() as db: + work = await next_scope_work(db, fence, cycle.version.cycle_id, within=scope) + assert work.record_type == "file" and work.parent.identity == message.identity + request = BeginScan( + fence=fence, + scope=CompletedScope( + record_type="file", + container_id=parent_container_key(message.identity), + parent=message.identity, + ), + cycle_id=cycle.version.cycle_id, + fingerprint=CONFIG.fingerprint, + expected_parent_epoch=work.parent_visibility_epoch, + expected_parent_revision=work.parent.revision, + exact_parent_observation=message, + ) + return cycle, channel, message, scope, state, request + + +async def test_early_file_commits_without_completing_history(database, source): + service, fence = source + cycle, _, message, scope, history, request = await setup(database, source) + with pytest.raises(CycleConflict, match="fresh exact"): + async with database() as db: + await service.begin_scan( + db, request.model_copy(update={"exact_parent_observation": None}) + ) + async with database() as db: + admitted = await service.admit_scan(db, request) + state = await page( + database, service, fence, admitted.state, item("file", "F", message.identity), final=True + ) + await finish(database, service, fence, state) + async with database() as db: + assert await next_scope_work(db, fence, cycle.version.cycle_id, within=scope) is None + current = await service.read_scan(db, fence, scope) + assert current.phase == "collecting" and current.version == history.version + refreshed = await service.read_cycle(db, fence) + with pytest.raises(CycleConflict, match="incomplete child"): + async with database() as db: + await service.complete_cycle(db, CompleteCycle(fence=fence, expected=refreshed.version)) + + +async def test_unseen_and_other_inventory_owners_cannot_enter_early_work(database, source): + service, fence = source + cycle, channel, _, scope, _, request = await setup(database, source) + unseen = item("message", "unseen", channel.identity) + await capture(database, service, fence, unseen) + async with database() as db: + row = await db.scalar(select(Entity).where(Entity.native_id == "unseen")) + forged = request.model_copy( + update={ + "scope": CompletedScope( + record_type="file", + container_id=parent_container_key(unseen.identity), + parent=unseen.identity, + ), + "expected_parent_revision": row.record_revision, + "expected_parent_epoch": row.visibility_epoch, + "exact_parent_observation": unseen, + } + ) + with pytest.raises(CycleConflict, match="ancestor membership"): + async with database() as db: + await service.admit_scan(db, forged) + # Same kind/container with a different owner is not a broad frontier query. + wrong = scope.model_copy( + update={"parent": channel.identity.model_copy(update={"native_id": "other"})} + ) + with pytest.raises(CycleConflict, match="parent is missing"): + async with database() as db: + await next_scope_work(db, fence, cycle.version.cycle_id, within=wrong) + + +async def test_early_receipt_invalidated_by_edit_and_inventory_restart(database, source): + service, fence = source + cycle, _, message, scope, history, request = await setup(database, source) + async with database() as db: + admitted = await service.admit_scan(db, request) + history = await page( + database, + service, + fence, + history, + message.model_copy(update={"payload": {"id": "M", "edited": True}}), + ) + with pytest.raises(CycleConflict, match="fresh exact"): + await page(database, service, fence, admitted.state, final=True) + await scan(database, service, fence, cycle, scope, expected=history.version, restart=True) + async with database() as db: + assert await next_scope_work(db, fence, cycle.version.cycle_id, within=scope) is None + with pytest.raises(CycleConflict, match="ancestor membership"): + await page(database, service, fence, admitted.state, final=True) + + +async def test_new_writer_cannot_use_early_parent_before_root_refresh(database, source): + from uuid import uuid4 + + service, fence = source + cycle, _, _, scope, _, request = await setup(database, source) + 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, + ) + with pytest.raises(CycleConflict, match="ancestor membership"): + async with database() as db: + await next_scope_work(db, newer, cycle.version.cycle_id, within=scope) + with pytest.raises(CycleConflict, match="ancestor membership"): + async with database() as db: + await service.admit_scan(db, request.model_copy(update={"fence": newer})) 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 index 3a3b81c94a..ba197455f2 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py @@ -1,5 +1,6 @@ """Actual child frontier commits/retry with synthetic Slack HTTP and original bytes.""" +import asyncio from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 @@ -7,6 +8,10 @@ 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 @@ -14,7 +19,9 @@ 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 @@ -170,3 +177,119 @@ async def test_omitted_accessible_message_does_not_remove_retained_file_children assert stored.deleted_at is None and stored.content_access == "available" if row.entity_definition_short_name == "file": assert stored.blobs + + +async def test_files_progress_before_history_finishes_and_resume_both_cursors( + database, source, tmp_path +): + storage = FilesystemBackend(tmp_path) + downloads = [] + histories = [] + interrupt = True + + async def respond(url, params): + if url.endswith("conversations.list"): + return ROOT + if "oldest" in params: + 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": [], "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"] + 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"] + 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 diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index d3110492ab..f62ec216e5 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -7,7 +7,12 @@ from sqlalchemy.ext.asyncio import AsyncSession -from airweave.domains.entities.canonical.cycle_models import BeginCycle, CaptureCycle, RestartCycle +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 ( @@ -42,6 +47,14 @@ 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.""" @@ -62,6 +75,7 @@ def __init__( 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.""" @@ -159,45 +173,69 @@ async def fresh_full(self, previous: CaptureCycle) -> CaptureCycle: async def run_cycle(self, cycle: CaptureCycle) -> CaptureCycle: """Process the current frontier; every acknowledgement remains in the existing SQL store.""" - configuration = cycle.configuration 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 - 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(configuration.children_of(work.record_type)), - ) - 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 + 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, @@ -206,6 +244,7 @@ async def scan( 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) @@ -213,7 +252,8 @@ async def scan( return state = admission.state parent = admission.parent - restarts = 0 + restart_key = (cycle.version.cycle_id, scope) + pages = 0 while state.phase == "collecting": await self.check_limits() try: @@ -221,9 +261,9 @@ async def scan( scope, state.continuation, files=self.files, parent=parent ) except (InvalidScanContinuation, InvalidScopeCheckpoint) as error: - if restarts >= 1: + if restart_key in self._restarted_scopes: raise - restarts += 1 + 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. @@ -244,6 +284,10 @@ async def scan( ) 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 From c56fa88dcc97efd01e61e91655b5f0dc68f4dbde Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 13:50:30 -0700 Subject: [PATCH 248/266] Project authenticated canonical records during active capture --- .../entities/canonical/projection_models.py | 15 ++ .../entities/canonical/projection_store.py | 64 ++++++- .../domains/entities/canonical/projector.py | 15 ++ .../entities/canonical/reprojection.py | 3 +- .../entities/canonical/tests/helpers.py | 49 ++++- .../tests/test_calendar_exclusions.py | 10 +- .../canonical/tests/test_drive_projection.py | 10 +- .../tests/test_extraction_coverage.py | 24 ++- .../canonical/tests/test_gmail_observation.py | 7 +- .../canonical/tests/test_native_text.py | 11 +- .../canonical/tests/test_notion_projection.py | 11 +- .../tests/test_notion_property_projection.py | 7 +- .../tests/test_outlook_calendar_projection.py | 10 +- .../canonical/tests/test_projection.py | 13 +- .../tests/test_projection_admission.py | 176 ++++++++++++++++++ .../canonical/tests/test_projection_gc.py | 8 +- .../canonical/tests/test_projection_ocr.py | 5 +- .../tests/test_projection_recovery.py | 51 ++++- .../canonical/tests/test_reprojection.py | 8 +- .../canonical/tests/test_retained_text.py | 6 +- .../canonical/tests/test_slack_projection.py | 11 +- .../native_ingestion/tests/test_ingestion.py | 14 +- .../native_ingestion/tests/test_projection.py | 8 +- .../native_ingestion/tests/test_search.py | 40 ++-- .../activities/discover_native_projection.py | 18 +- .../activities/project_canonical_records.py | 22 +-- .../tests/test_discover_native_projection.py | 43 +++++ .../tests/test_project_canonical_records.py | 16 ++ .../temporal/worker/tests/test_wiring.py | 6 + .../domains/temporal/worker/wiring.py | 2 +- .../workflows/recover_native_projection.py | 13 +- .../workflows/run_source_connection.py | 36 +++- .../tests/test_native_projection_recovery.py | 147 ++++++++++++++- .../workflows/tests/test_reproject_command.py | 47 +++++ .../tests/test_source_projection_replay.py | 156 ++++++++++++++++ backend/scripts/reproject_canonical.py | 24 ++- 36 files changed, 976 insertions(+), 130 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection_admission.py create mode 100644 backend/airweave/domains/temporal/activities/tests/test_discover_native_projection.py create mode 100644 backend/airweave/domains/temporal/workflows/tests/test_reproject_command.py create mode 100644 backend/airweave/domains/temporal/workflows/tests/test_source_projection_replay.py diff --git a/backend/airweave/domains/entities/canonical/projection_models.py b/backend/airweave/domains/entities/canonical/projection_models.py index 11c8e51804..f832b35afe 100644 --- a/backend/airweave/domains/entities/canonical/projection_models.py +++ b/backend/airweave/domains/entities/canonical/projection_models.py @@ -8,6 +8,11 @@ from airweave.domains.entities.canonical.models import SourceRecord +def canonical_projection_workflow_id(organization_id: UUID | str, sync_id: UUID | str) -> str: + """One execution identity shared by scheduled, capture and operator projection.""" + return f"canonical-projection:{organization_id}:{sync_id}" + + class ProjectionLocator(BaseModel): """Stored in the existing original_entity_id field, never a native provider ID.""" @@ -42,11 +47,21 @@ def parse(cls, value: str) -> "ProjectionLocator | None": ) +class ProjectionBinding(BaseModel): + """Authenticated tenant source and exact destination captured with pending work.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + source_connection_id: UUID + source_name: str + collection_id: UUID + + class ProjectionWork(BaseModel): """Snapshot plus publication pointer used for compare-and-swap.""" model_config = ConfigDict(extra="forbid", frozen=True) organization_id: UUID + binding: ProjectionBinding record: SourceRecord pipeline_version: int previous_generation: UUID | None diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index e147245452..3a7ed9350d 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -12,6 +12,7 @@ from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.projection_models import ( + ProjectionBinding, ProjectionDocument, ProjectionLocator, ProjectionSourcePage, @@ -22,6 +23,7 @@ from airweave.domains.entities.canonical.store import content_is_available, source_record from airweave.domains.entities.canonical.text_artifacts import text_manifest from airweave.domains.entities.canonical.text_models import TextArtifact +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 @@ -96,6 +98,47 @@ def _pending_record(): class CanonicalProjectionStore: """Optimistic projection computation, serialized CAS against source capture.""" + async def binding( + self, db: AsyncSession, organization_id: UUID, sync_id: UUID + ) -> ProjectionBinding | None: + """Resolve one authenticated source and collection; ambiguity fails closed.""" + rows = ( + await db.execute( + select(SourceConnection.id, SourceConnection.short_name, Collection.id) + .join( + Collection, + and_( + Collection.readable_id == SourceConnection.readable_collection_id, + Collection.organization_id == SourceConnection.organization_id, + ), + ) + .join( + Sync, + and_( + Sync.id == SourceConnection.sync_id, + Sync.organization_id == SourceConnection.organization_id, + ), + ) + .where( + SourceConnection.organization_id == organization_id, + SourceConnection.sync_id == sync_id, + SourceConnection.is_authenticated.is_(True), + ) + .limit(2) + ) + ).all() + if len(rows) != 1: + return None + source, name, collection = rows[0] + return ProjectionBinding( + source_connection_id=source, source_name=name, collection_id=collection + ) + + async def admit(self, db: AsyncSession, work: ProjectionWork) -> bool: + """Fence current authorization before any converter, storage or embedding call.""" + async with UnitOfWork(db): + return await self._current(db, work) is not None + async def pending( self, db: AsyncSession, @@ -109,6 +152,9 @@ async def pending( """One bounded UUID page; deleted records publish an empty generation.""" if not 1 <= limit <= 100: raise ValueError("Projection batch size must be between 1 and 100") + binding = await self.binding(db, organization_id, sync_id) + if binding is None: + return () statement = ( select(Entity, Sync.index_pipeline_version) .join(Sync, Sync.id == Entity.sync_id) @@ -127,6 +173,7 @@ async def pending( return tuple( ProjectionWork( organization_id=organization_id, + binding=binding, record=source_record(row), pipeline_version=version, previous_generation=row.indexed_generation, @@ -150,7 +197,16 @@ async def pending_sources( if not 1 <= limit <= 100: raise ValueError("Projection source page size must be between 1 and 100") owned_source = exists( - select(SourceConnection.id).where( + select(SourceConnection.id) + .join( + Collection, + and_( + Collection.readable_id == SourceConnection.readable_collection_id, + Collection.organization_id == SourceConnection.organization_id, + ), + ) + .where( + SourceConnection.is_authenticated.is_(True), SourceConnection.sync_id == Sync.id, SourceConnection.organization_id == Sync.organization_id, SourceConnection.short_name.in_(source_names), @@ -191,6 +247,8 @@ async def _current(self, db: AsyncSession, work: ProjectionWork) -> Entity | Non ) if sync is None or sync.index_pipeline_version != work.pipeline_version: return None + if await self.binding(db, work.organization_id, work.record.sync_id) != work.binding: + return None row = ( await db.execute( select(Entity, content_is_available()) @@ -227,6 +285,8 @@ async def prepare( text_representations: tuple[TextArtifact, ...] | None = None, ) -> bool: """Commit exact immutable deletion identities before the first remote feed.""" + if collection_id != work.binding.collection_id: + raise ValueError("Projection destination does not match authenticated binding") identities = set() indexed_parts = set() for document in documents: @@ -325,6 +385,8 @@ async def publish( if ( attempt is None or attempt.retired_at is not None + or attempt.collection_id != work.binding.collection_id + or attempt.sync_id != work.record.sync_id or attempt.organization_id != work.organization_id or attempt.record_id != work.record.id or attempt.revision != work.record.revision diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 7c8b0a3cdc..2e7c577a63 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -79,6 +79,18 @@ def __init__( self._processor = processor self._storage = storage + async def _admit( + self, work: ProjectionWork, source_name: str, destination: VespaDestination + ) -> bool: + """Close authorization transaction before any external processing begins.""" + if ( + source_name != work.binding.source_name + or destination.collection_id != work.binding.collection_id + ): + return False + async with self._sessions() as db: + return await self._store.admit(db, work) + async def project_one( self, work: ProjectionWork, @@ -92,6 +104,9 @@ async def project_one( map_record, ) + if not await self._admit(work, source_name, destination): + return False + generation = uuid4() chunks = [] coverage = ExtractionCoverage(parts=()) diff --git a/backend/airweave/domains/entities/canonical/reprojection.py b/backend/airweave/domains/entities/canonical/reprojection.py index c01984d622..39b22d5450 100644 --- a/backend/airweave/domains/entities/canonical/reprojection.py +++ b/backend/airweave/domains/entities/canonical/reprojection.py @@ -6,6 +6,7 @@ from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession +from airweave.domains.entities.canonical.projection_models import canonical_projection_workflow_id from airweave.domains.entities.canonical.search_metadata import SEARCH_METADATA_PIPELINE_VERSION from airweave.models.collection import Collection from airweave.models.entity import Entity @@ -101,7 +102,7 @@ async def plan_reprojection( target_version=target_version, captured_records=count, target_pending_records=pending, - workflow_id=f"reprojection:{organization_id}:{sync_id}:{target_version}", + workflow_id=canonical_projection_workflow_id(organization_id, sync_id), applied=apply, ) if apply: diff --git a/backend/airweave/domains/entities/canonical/tests/helpers.py b/backend/airweave/domains/entities/canonical/tests/helpers.py index 940a9cc5a0..d27531b583 100644 --- a/backend/airweave/domains/entities/canonical/tests/helpers.py +++ b/backend/airweave/domains/entities/canonical/tests/helpers.py @@ -25,8 +25,6 @@ async def capture(database, service, fence, *records): async def publish_prepared(store, db, work, generation, count, collection_id=None): """Synthetic exact manifests exercise the same mandatory pre-feed transaction.""" - from uuid import uuid4 - from airweave.domains.entities.canonical.projection_models import ( ProjectionDocument, ProjectionLocator, @@ -40,7 +38,7 @@ async def publish_prepared(store, db, work, generation, count, collection_id=Non generation=generation, part_index=0, ).encode() - collection_id = collection_id or uuid4() + collection_id = collection_id or work.binding.collection_id prepared = await store.prepare( db, work, @@ -57,3 +55,48 @@ async def publish_prepared(store, db, work, generation, count, collection_id=Non ), ) return prepared and await store.publish(db, work, generation, count) + + +async def bind_projection(database, fence, source_name="gmail", collection_id=None): + """Create the real authenticated source/collection required for synthetic projection.""" + from uuid import uuid4 + + from airweave.domains.entities.canonical.projection_models import ProjectionBinding + from airweave.models.collection import Collection + from airweave.models.source_connection import SourceConnection + from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata + + collection_id = collection_id or uuid4() + readable_id = str(collection_id) + source_id = uuid4() + 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( + id=collection_id, + name="Synthetic projection", + readable_id=readable_id, + organization_id=fence.organization_id, + vector_db_deployment_metadata_id=deployment.id, + ) + ) + await db.flush() + db.add( + SourceConnection( + id=source_id, + name="Synthetic projection", + short_name=source_name, + organization_id=fence.organization_id, + sync_id=fence.sync_id, + readable_collection_id=readable_id, + is_authenticated=True, + ) + ) + await db.commit() + return ProjectionBinding( + source_connection_id=source_id, source_name=source_name, collection_id=collection_id + ) diff --git a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py index 7ae56e5ce9..14aa947c73 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py +++ b/backend/airweave/domains/entities/canonical/tests/test_calendar_exclusions.py @@ -14,12 +14,14 @@ 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 +from airweave.domains.entities.canonical.tests.helpers import bind_projection from airweave.models import Entity, ProjectionGeneration from airweave.platform.sources.records.google_calendar import record async def test_cancelled_exception_read_project_reinstate_and_revoke(database, source): capture, fence = source + binding = await bind_projection(database, fence, "google_calendar") parent = record("calendar", {"id": "cal"}) native = { "id": "instance", @@ -54,7 +56,7 @@ async def test_cancelled_exception_read_project_reinstate_and_revoke(database, s processor = MagicMock() processor.build_text = AsyncMock() destination = MagicMock() - destination.collection_id = uuid4() + destination.collection_id = binding.collection_id destination.feed_prepared = AsyncMock() projector = CanonicalProjector(projections, database, processor, AsyncMock()) assert await projector.project_one(work, "google_calendar", destination, MagicMock()) @@ -122,6 +124,7 @@ async def test_cancelled_exception_read_project_reinstate_and_revoke(database, s async def test_live_event_cannot_publish_empty_or_exclusion_nonempty(database, source): capture, fence = source + binding = await bind_projection(database, fence, "google_calendar") fact = record( "event", { @@ -146,13 +149,14 @@ async def test_live_event_cannot_publish_empty_or_exclusion_nonempty(database, s } ) generation = uuid4() - assert await store.prepare(db, active, generation, uuid4(), ()) + assert await store.prepare(db, active, generation, binding.collection_id, ()) with pytest.raises(ValueError, match="indexed content"): await store.publish(db, active, generation, 0) async def test_operational_occurrences_publish_zero_documents_without_retry(database, source): capture, fence = source + binding = await bind_projection(database, fence, "google_calendar") parent = record("calendar", {"id": "cal"}) occurrence = record( "event_occurrence", @@ -177,7 +181,7 @@ async def test_operational_occurrences_publish_zero_documents_without_retry(data processor = MagicMock() processor.build_text = AsyncMock() destination = MagicMock() - destination.collection_id = uuid4() + destination.collection_id = binding.collection_id destination.feed_prepared = AsyncMock() projector = CanonicalProjector(projections, database, processor, AsyncMock()) assert await projector.project_one(work, "google_calendar", destination, MagicMock()) diff --git a/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py b/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py index 7f9b8575a7..b52e048767 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_drive_projection.py @@ -10,7 +10,7 @@ from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.projection_mappers import map_record -from airweave.domains.entities.canonical.projection_models import ProjectionWork +from airweave.domains.entities.canonical.projection_models import ProjectionBinding, ProjectionWork from airweave.domains.entities.canonical.projector import _select_inputs from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity from airweave.platform.sources.records.workspace_manifest import ( @@ -126,7 +126,13 @@ def docs_record(*, native_available): def coverage(result, source): work = ProjectionWork( - organization_id=uuid4(), record=source, pipeline_version=2, previous_generation=None + binding=ProjectionBinding( + source_connection_id=uuid4(), source_name="google_drive", collection_id=uuid4() + ), + organization_id=uuid4(), + record=source, + pipeline_version=2, + previous_generation=None, ) return _select_inputs(result, work, "google_drive", uuid4(), lambda _: True)[1] diff --git a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py index 181c418b71..81c4d6ea0f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py +++ b/backend/airweave/domains/entities/canonical/tests/test_extraction_coverage.py @@ -23,7 +23,7 @@ ) from airweave.domains.entities.canonical.projector import CanonicalProjector from airweave.domains.entities.canonical.requests import RecordIdentity -from airweave.domains.entities.canonical.tests.helpers import capture, observation +from airweave.domains.entities.canonical.tests.helpers import bind_projection, capture, observation from airweave.domains.entities.canonical.tests.test_gmail_projection import part from airweave.domains.search.owned import OwnedSearchService from airweave.domains.search.owned_models import OwnedSearchRequest @@ -52,8 +52,8 @@ def original(*attachments, completeness="complete"): ) -def destination(): - result = MagicMock(collection_id=uuid4(), feed_prepared=AsyncMock()) +def destination(collection_id): + result = MagicMock(collection_id=collection_id, feed_prepared=AsyncMock()) result.prepare_documents = lambda chunks: { "base_entity": [ EntityTransformer(collection_id=result.collection_id).transform(c) for c in chunks @@ -75,6 +75,7 @@ async def test_message_pdf_video_and_missing_part_publish_truthful_coverage( database, source, tmp_path ): service, fence = source + binding = await bind_projection(database, fence, "gmail") pdf = fitz.open() page = pdf.new_page() page.insert_text( @@ -98,7 +99,7 @@ async def test_message_pdf_video_and_missing_part_publish_truthful_coverage( ) await capture(database, service, fence, item) store = CanonicalProjectionStore() - target = destination() + target = destination(binding.collection_id) result = await projector(database, FilesystemBackend(tmp_path)).batch( fence.organization_id, fence.sync_id, "gmail", target, logger ) @@ -182,13 +183,14 @@ async def test_supported_attachment_conversion_failure_and_feed_failure_stay_pen database, source, tmp_path ): service, fence = source + binding = await bind_projection(database, fence, "gmail") await capture( database, service, fence, original(part(b"broken binary\x00", mime="application/pdf", filename="broken.pdf")), ) - target = destination() + target = destination(binding.collection_id) project = projector(database, FilesystemBackend(tmp_path)) assert ( await project.batch(fence.organization_id, fence.sync_id, "gmail", target, logger) @@ -224,6 +226,7 @@ async def test_unsupported_only_publication_is_explicit_and_retries_on_pipeline_ database, source, tmp_path ): service, fence = source + binding = await bind_projection(database, fence, "google_drive") from airweave.domains.entities.canonical.tests.test_gmail_projection import blob, record captured = record({}, sync_id=fence.sync_id) @@ -240,7 +243,7 @@ async def test_unsupported_only_publication_is_explicit_and_retries_on_pipeline_ blobs=(ref,), ), ) - target = destination() + target = destination(binding.collection_id) assert ( await projector(database, storage).batch( fence.organization_id, fence.sync_id, "google_drive", target, logger @@ -288,11 +291,12 @@ async def test_coverage_manifest_is_exact_immutable_and_capture_fences_publicati ) service, fence = source + binding = await bind_projection(database, fence, "gmail") await capture(database, service, fence, original()) store = CanonicalProjectionStore() async with database() as db: work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] - generation, collection = uuid4(), uuid4() + generation, collection = uuid4(), binding.collection_id locator = ProjectionLocator( record_id=work.record.id, revision=work.record.revision, @@ -356,11 +360,12 @@ async def test_inline_image_without_ocr_keeps_body_partial_but_converter_failure database, source, tmp_path ): service, fence = source + binding = await bind_projection(database, fence, "gmail") item = original(part(b"retained image bytes", mime="image/png", filename="inline.png")) item.payload["payload"]["mimeType"] = "multipart/related" await capture(database, service, fence, item) storage = FilesystemBackend(tmp_path) - target = destination() + target = destination(binding.collection_id) result = await projector(database, storage).batch( fence.organization_id, fence.sync_id, "gmail", target, logger ) @@ -408,12 +413,13 @@ async def test_drive_metadata_only_reports_missing_original_then_new_bytes_are_i from airweave.domains.entities.canonical.tests.test_gmail_projection import blob, record service, fence = source + binding = await bind_projection(database, fence, "google_drive") identity = RecordIdentity(record_type="file", native_id="retained-file") payload = {"id": identity.native_id, "name": "Notes.txt", "mimeType": "text/plain"} metadata = observation(identity=identity, payload=payload, completeness="metadata_only") await capture(database, service, fence, metadata) storage = FilesystemBackend(tmp_path) - target = destination() + target = destination(binding.collection_id) project = projector(database, storage) assert ( await project.batch(fence.organization_id, fence.sync_id, "google_drive", target, logger) diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py index f12ff32f40..3d3382b28d 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_observation.py @@ -14,7 +14,11 @@ RecordIdentity, ) from airweave.domains.entities.canonical.store import capture_fingerprint -from airweave.domains.entities.canonical.tests.helpers import capture, publish_prepared +from airweave.domains.entities.canonical.tests.helpers import ( + bind_projection, + capture, + publish_prepared, +) pytestmark = pytest.mark.integration @@ -23,6 +27,7 @@ async def test_locator_rotation_preserves_originals_and_requires_current_publica database, source ): service, fence = source + await bind_projection(database, fence, "gmail") digest = sha256(b"identical message body").hexdigest() payload = { "id": "synthetic-message", diff --git a/backend/airweave/domains/entities/canonical/tests/test_native_text.py b/backend/airweave/domains/entities/canonical/tests/test_native_text.py index 9e9caf9ac2..f35b9e239f 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_native_text.py +++ b/backend/airweave/domains/entities/canonical/tests/test_native_text.py @@ -88,7 +88,11 @@ async def test_slack_native_retained_read_matches_text_used_for_index( 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 capture, observation + from airweave.domains.entities.canonical.tests.helpers import ( + bind_projection, + capture, + observation, + ) from airweave.domains.entities.canonical.tests.test_extraction_coverage import ( destination, projector, @@ -97,6 +101,7 @@ async def test_slack_native_retained_read_matches_text_used_for_index( from airweave.models.sync import Sync service, fence = source + binding = await bind_projection(database, fence, "slack") original = observation( identity=RecordIdentity(record_type="message", native_id="1.000001", container_id="C1"), payload={"ts": "1.000001", **({"text": body} if body is not None else {})}, @@ -120,7 +125,9 @@ async def fixed_chunks(entities, context, runtime): work = (await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id))[ 0 ] - assert await projection.project_one(work, "slack", destination(), MagicMock()) + assert await projection.project_one( + work, "slack", destination(binding.collection_id), MagicMock() + ) reader = CanonicalTextReader( CanonicalQueryService(CanonicalRecordStore(), CanonicalQueryStore(), "key"), storage, diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py index 37dea6ddf7..d303af71db 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_projection.py @@ -19,7 +19,12 @@ ) from airweave.domains.entities.canonical.projector import CanonicalProjector from airweave.domains.entities.canonical.requests import RecordIdentity -from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared +from airweave.domains.entities.canonical.tests.helpers import ( + bind_projection, + capture, + observation, + publish_prepared, +) from airweave.models import Entity, Sync PAGE, BLOCK = str(UUID(int=1)), str(UUID(int=2)) @@ -232,6 +237,7 @@ def test_original_search_view_is_in_existing_entity_registry(): @pytest.mark.integration async def test_notion_publication_uses_current_revision_and_ancestor_authority(database, source): service, fence = source + await bind_projection(database, fence, "notion") root = original(properties={"title": {"type": "title", "title": [rich("Page")]}}) child = block("paragraph", {"rich_text": [rich("First version")]}) root_observation = observation( @@ -281,6 +287,7 @@ async def test_notion_publication_uses_current_revision_and_ancestor_authority(d @pytest.mark.integration async def test_unknown_notion_block_stays_pending_in_actual_projector(database, source): service, fence = source + binding = await bind_projection(database, fence, "notion") root = original(properties={"title": {"type": "title", "title": []}}) child = block("future_type", {"body": "not raw JSON prose"}) await capture( @@ -300,7 +307,7 @@ async def test_unknown_notion_block_stays_pending_in_actual_projector(database, pending = await store.pending(db, fence.organization_id, fence.sync_id) root_work = next(item for item in pending if item.record.identity.record_type == "page") assert await publish_prepared(store, db, root_work, uuid4(), 1) - processor, destination = MagicMock(), MagicMock() + processor, destination = MagicMock(), MagicMock(collection_id=binding.collection_id) projector = CanonicalProjector(store, database, processor, AsyncMock()) result = await projector.batch( fence.organization_id, fence.sync_id, "notion", destination, MagicMock() diff --git a/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py b/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py index 81a00e4e43..2ce6953fe2 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_notion_property_projection.py @@ -11,6 +11,7 @@ map_record, ) from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity +from airweave.domains.entities.canonical.tests.helpers import bind_projection from airweave.domains.entities.canonical.tests.test_notion_projection import PAGE, STAMP, original @@ -204,6 +205,7 @@ async def test_property_publication_obeys_current_parent(database, source): from airweave.models import Entity, Sync service, fence = source + await bind_projection(database, fence, "notion") record, content = fixture("number", [item("number", 1)]) blob = record.blobs[0].model_copy( update={"key": f"canonical/{fence.sync_id}/blobs/sha256/{record.blobs[0].sha256}"} @@ -273,6 +275,7 @@ async def test_unsupported_property_publishes_explicit_extraction(database, sour from airweave.models import Entity service, fence = source + binding = await bind_projection(database, fence, "notion") record, content = fixture("formula", [item("formula", {"type": "unsupported"})], "unsupported") blob = record.blobs[0].model_copy( update={"key": f"canonical/{fence.sync_id}/blobs/sha256/{record.blobs[0].sha256}"} @@ -301,7 +304,9 @@ async def test_unsupported_property_publishes_explicit_extraction(database, sour ) if w.record.identity.record_type == "page_property" ) - assert await projector(database, storage).project_one(work, "notion", destination(), logger) + assert await projector(database, storage).project_one( + work, "notion", destination(binding.collection_id), logger + ) async with database() as db: row = await db.scalar(select(Entity).where(Entity.id == work.record.id)) coverage = await current_extraction( diff --git a/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py index c9763e5abb..fb85ed93df 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_outlook_calendar_projection.py @@ -8,7 +8,7 @@ from airweave.domains.converters.html import HtmlConverter from airweave.domains.entities.canonical.outlook_calendar_projection import map_outlook_calendar -from airweave.domains.entities.canonical.projection_models import ProjectionWork +from airweave.domains.entities.canonical.projection_models import ProjectionBinding, ProjectionWork from airweave.domains.entities.canonical.projector import StrictProjectionTracker, _select_inputs from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.entities.canonical.tests.test_outlook_projection import source @@ -77,7 +77,13 @@ def event(body): async def build(record, tmp_path): mapped = await map_outlook_calendar(record, tmp_path) work = ProjectionWork( - organization_id=uuid4(), record=record, pipeline_version=2, previous_generation=None + binding=ProjectionBinding( + source_connection_id=uuid4(), source_name="outlook_calendar", collection_id=uuid4() + ), + organization_id=uuid4(), + record=record, + pipeline_version=2, + previous_generation=None, ) selected, coverage = _select_inputs(mapped, work, "outlook_calendar", uuid4(), lambda _: True) registry = MagicMock() diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection.py b/backend/airweave/domains/entities/canonical/tests/test_projection.py index 29d24a8859..8f54f6faeb 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection.py @@ -11,7 +11,12 @@ publication_matches, ) from airweave.domains.entities.canonical.requests import RecordIdentity -from airweave.domains.entities.canonical.tests.helpers import capture, observation, publish_prepared +from airweave.domains.entities.canonical.tests.helpers import ( + bind_projection, + capture, + observation, + publish_prepared, +) from airweave.models import Entity, Sync pytestmark = pytest.mark.integration @@ -24,6 +29,7 @@ async def pending(database, fence): 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() @@ -64,6 +70,7 @@ async def test_publication_compare_and_swap_revision_generation_and_pipeline(dat 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) @@ -88,6 +95,7 @@ async def test_parent_access_loss_fences_publication_and_tombstone_is_empty(data 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() @@ -116,6 +124,7 @@ async def test_owned_payload_mapper_to_publication_pipeline(database, source): from airweave.platform.destinations.vespa.transformer import EntityTransformer service, fence = source + binding = await bind_projection(database, fence, "slack") await capture( database, service, @@ -152,7 +161,7 @@ async def build_text(entities, context, runtime, *, native_bodies=None): processor.build_text = build_text processor.process_built_text = process destination = MagicMock() - destination.collection_id = uuid4() + destination.collection_id = binding.collection_id destination.prepare_documents = lambda chunks: { "base_entity": [ EntityTransformer(collection_id=destination.collection_id).transform(chunk) 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_gc.py b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py index 634b897572..b9f0b0343d 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_gc.py @@ -12,13 +12,14 @@ scope_projection_document_id, ) from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore -from airweave.domains.entities.canonical.tests.helpers import capture, observation +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: @@ -31,7 +32,7 @@ async def prepare(database, source, count=3): generation=generation, part_index=0, ).encode() - collection_id = uuid4() + collection_id = binding.collection_id docs = tuple( ProjectionDocument( schema_name="base_entity", @@ -259,6 +260,7 @@ async def test_manifest_rejects_foreign_scope_and_unanchored_locator(database, s 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, @@ -270,7 +272,7 @@ async def test_empty_tombstone_publication_does_not_become_pending_each_gc_pass( 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, uuid4(), ()) + 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: diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py b/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py index e0f8647595..94b8897837 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_ocr.py @@ -15,7 +15,7 @@ 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.helpers import bind_projection, capture, observation from airweave.domains.sync_pipeline.processors.chunk_embed import ChunkEmbedProcessor from airweave.models import Entity @@ -30,6 +30,7 @@ async def test_scanned_pdf_without_ocr_stays_pending_with_owned_bytes(database, 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}", @@ -49,7 +50,7 @@ async def test_scanned_pdf_without_ocr_stays_pending_with_owned_bytes(database, blobs=(blob,), ), ) - destination = MagicMock(feed_prepared=AsyncMock()) + destination = MagicMock(collection_id=binding.collection_id, feed_prepared=AsyncMock()) store = CanonicalProjectionStore() projector = CanonicalProjector( store, diff --git a/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py index bf656a88d5..0302c91b6b 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_projection_recovery.py @@ -8,8 +8,9 @@ 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 capture, observation -from airweave.models import Entity, Organization, SourceConnection, Sync, SyncJob +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): @@ -17,11 +18,29 @@ async def test_source_pages_skip_failed_indexed_unsupported_and_foreign_rows(dat 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", - "fresh-two", + "provider", + "unauthenticated", + "missing-collection", + "foreign-collection", "tombstone", "failed", "indexed", @@ -48,8 +67,21 @@ async def test_source_pages_skip_failed_indexed_unsupported_and_foreign_rows(dat sync_id=sync_id, organization_id=foreign if case == "foreign-source" else seed.organization_id, name=case, - short_name="legacy" if case == "unsupported" else "almanac", - is_authenticated=False, + 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() @@ -83,14 +115,16 @@ async def test_source_pages_skip_failed_indexed_unsupported_and_foreign_rows(dat elif case == "withdrawn": record.removal_reason = "access_revoked" await db.commit() - if case in ("fresh", "fresh-two", "tombstone"): + if case in ("fresh", "provider", "tombstone"): expected.add(sync_id) store = CanonicalProjectionStore() async with database() as db: - first = await store.pending_sources(db, ("almanac",), limit=2) + 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",), after_id=first.next_cursor, limit=2) + 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( @@ -102,6 +136,7 @@ async def test_source_pages_skip_failed_indexed_unsupported_and_foreign_rows(dat 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: diff --git a/backend/airweave/domains/entities/canonical/tests/test_reprojection.py b/backend/airweave/domains/entities/canonical/tests/test_reprojection.py index 94a8571bfd..11e9aff798 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_reprojection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_reprojection.py @@ -60,7 +60,9 @@ async def test_wrong_scope_and_changed_version_fail(database, search_indexed): async def test_launch_failure_keeps_queue_and_identical_retry_starts_once( - database, search_indexed, monkeypatch # noqa: F811 + database, + search_indexed, # noqa: F811 + monkeypatch, ): from argparse import Namespace from unittest.mock import AsyncMock @@ -91,6 +93,8 @@ async def test_launch_failure_keeps_queue_and_identical_retry_starts_once( 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"].endswith(":3") + 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_retained_text.py b/backend/airweave/domains/entities/canonical/tests/test_retained_text.py index 4442e432bc..75d3f944fe 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_retained_text.py +++ b/backend/airweave/domains/entities/canonical/tests/test_retained_text.py @@ -4,7 +4,6 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock -from uuid import uuid4 import fitz import httpx @@ -25,7 +24,7 @@ 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 capture, observation +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 @@ -37,6 +36,7 @@ @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): @@ -89,7 +89,7 @@ async def fixed_chunks(entities, context, runtime): ) processor._chunk_entities = fixed_chunks - destination = MagicMock(collection_id=uuid4(), feed_prepared=AsyncMock()) + 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) diff --git a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py index 2192bff390..b6eec492eb 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_projection.py @@ -121,7 +121,11 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source 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 capture, observation + from airweave.domains.entities.canonical.tests.helpers import ( + bind_projection, + capture, + observation, + ) from airweave.domains.entities.canonical.tests.test_extraction_coverage import ( destination, projector, @@ -130,6 +134,7 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source 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() @@ -155,7 +160,7 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source ) await capture(database, service, fence, item) result = await projector(database, storage).batch( - fence.organization_id, fence.sync_id, "slack", destination(), logger + 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: @@ -229,7 +234,7 @@ async def test_slack_pdf_publication_records_partial_extraction(database, source 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(), logger + 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: diff --git a/backend/airweave/domains/native_ingestion/tests/test_ingestion.py b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py index 34aed39522..cdf9ef88a2 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_ingestion.py +++ b/backend/airweave/domains/native_ingestion/tests/test_ingestion.py @@ -18,6 +18,7 @@ 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, @@ -47,17 +48,10 @@ def snapshot(native_id="one", revision=1, **changes): async def bind(database, fence, *, dataset="knowledge", provider="almanac"): + binding = await bind_projection(database, fence, provider) async with database() as db: - db.add( - SourceConnection( - id=uuid4(), - organization_id=fence.organization_id, - sync_id=fence.sync_id, - name="Native test", - short_name=provider, - config_fields={"owner_id": "owner-one", "dataset": dataset}, - ) - ) + source = await db.get(SourceConnection, binding.source_connection_id) + source.config_fields = {"owner_id": "owner-one", "dataset": dataset} await db.commit() diff --git a/backend/airweave/domains/native_ingestion/tests/test_projection.py b/backend/airweave/domains/native_ingestion/tests/test_projection.py index c49cc9f3a2..fcdab49cf3 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_projection.py +++ b/backend/airweave/domains/native_ingestion/tests/test_projection.py @@ -167,7 +167,9 @@ async def fixed_chunks(entities, context, runtime): 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(), MagicMock()) + 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: @@ -246,7 +248,9 @@ async def fixed_chunks(entities, context, runtime): work = (await CanonicalProjectionStore().pending(db, fence.organization_id, fence.sync_id))[ 0 ] - assert await projection.project_one(work, "almanac", destination(), MagicMock()) + 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] diff --git a/backend/airweave/domains/native_ingestion/tests/test_search.py b/backend/airweave/domains/native_ingestion/tests/test_search.py index 12c45e41a7..cbbd7d28f4 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_search.py +++ b/backend/airweave/domains/native_ingestion/tests/test_search.py @@ -6,7 +6,7 @@ import pytest from fastapi import HTTPException -from sqlalchemy import update +from sqlalchemy import select, update from airweave.domains.entities.canonical.projection_models import ProjectionLocator from airweave.domains.entities.canonical.projection_store import CanonicalProjectionStore @@ -19,7 +19,6 @@ from airweave.domains.search.types import SearchResults from airweave.models.collection import Collection from airweave.models.source_connection import SourceConnection -from airweave.models.vector_db_deployment_metadata import VectorDbDeploymentMetadata async def test_native_search_requires_current_publication_and_source_access(database, source): @@ -28,25 +27,17 @@ async def test_native_search_requires_current_publication_and_source_access(data await bind(database, fence) await ingest(database, fence, snapshot()) async with database() as db: - deployment = VectorDbDeploymentMetadata( - dense_embedder="fake", embedding_dimensions=3, sparse_embedder="fake" - ) - db.add(deployment) - await db.flush() - collection = Collection( - name="Native synthetic", - readable_id="native-search", - organization_id=fence.organization_id, - vector_db_deployment_metadata_id=deployment.id, - ) - db.add(collection) - await db.flush() - await db.execute( - update(SourceConnection) - .where(SourceConnection.sync_id == fence.sync_id) - .values(readable_collection_id=collection.readable_id, is_authenticated=True) + 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, + ) ) - await db.commit() + assert collection is not None store = CanonicalProjectionStore() async with database() as db: work = (await store.pending(db, fence.organization_id, fence.sync_id))[0] @@ -80,9 +71,12 @@ async def test_native_search_requires_current_publication_and_source_access(data ) == [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 - ) == [] + assert ( + await visible_results( + db, fence.organization_id, collection.readable_id, [legacy], registry + ) + == [] + ) async def revoke_during_retrieval(**kwargs): async with database() as db: diff --git a/backend/airweave/domains/temporal/activities/discover_native_projection.py b/backend/airweave/domains/temporal/activities/discover_native_projection.py index 96d20a3570..33ed3c7007 100644 --- a/backend/airweave/domains/temporal/activities/discover_native_projection.py +++ b/backend/airweave/domains/temporal/activities/discover_native_projection.py @@ -1,23 +1,35 @@ -"""Discover committed native work without a second dispatch journal.""" +"""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: - """Native HTTP imports lack the provider workflow's durable projection child.""" + """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=("almanac",), + source_names=source_names, after_id=UUID(after_id) if after_id else None, limit=20, ) diff --git a/backend/airweave/domains/temporal/activities/project_canonical_records.py b/backend/airweave/domains/temporal/activities/project_canonical_records.py index 1b3dd5ad7f..5dab931d0d 100644 --- a/backend/airweave/domains/temporal/activities/project_canonical_records.py +++ b/backend/airweave/domains/temporal/activities/project_canonical_records.py @@ -3,7 +3,6 @@ from dataclasses import dataclass from uuid import UUID -from sqlalchemy import select from temporalio import activity from airweave.core.logging import logger @@ -13,8 +12,6 @@ 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.models.collection import Collection -from airweave.models.source_connection import SourceConnection from airweave.platform.destinations.vespa.destination import VespaDestination @@ -37,23 +34,10 @@ async def run( organization = UUID(organization_id) sync = UUID(sync_id) async with get_db_context() as db: - row = ( - await db.execute( - select(SourceConnection.short_name, Collection.id) - .join( - Collection, - Collection.readable_id == SourceConnection.readable_collection_id, - ) - .where( - SourceConnection.organization_id == organization, - SourceConnection.sync_id == sync, - Collection.organization_id == organization, - ) - ) - ).one_or_none() - if row is None: + binding = await CanonicalProjectionStore().binding(db, organization, sync) + if binding is None: return ProjectionBatchResult().model_dump(mode="json") - source_name, collection_id = row + 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. 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 index c2d67a2ac5..83d507354d 100644 --- a/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py +++ b/backend/airweave/domains/temporal/activities/tests/test_project_canonical_records.py @@ -41,6 +41,7 @@ async def activity_setup(database, source, monkeypatch): organization_id=fence.organization_id, sync_id=fence.sync_id, name="Native", + is_authenticated=True, short_name="almanac", readable_collection_id="native", ) @@ -129,3 +130,18 @@ async def test_projector_failure_propagates_and_closes_destination(activity_setu 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/worker/tests/test_wiring.py b/backend/airweave/domains/temporal/worker/tests/test_wiring.py index cb5f87774e..5948f2e0a9 100644 --- a/backend/airweave/domains/temporal/worker/tests/test_wiring.py +++ b/backend/airweave/domains/temporal/worker/tests/test_wiring.py @@ -37,6 +37,12 @@ def test_create_activities_returns_list(): 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 diff --git a/backend/airweave/domains/temporal/worker/wiring.py b/backend/airweave/domains/temporal/worker/wiring.py index 6e76f83c54..f6da596670 100644 --- a/backend/airweave/domains/temporal/worker/wiring.py +++ b/backend/airweave/domains/temporal/worker/wiring.py @@ -64,7 +64,7 @@ def create_activities() -> list: logger.debug("Wiring activities with container dependencies") return [ - DiscoverNativeProjectionActivity().run, + DiscoverNativeProjectionActivity(container.source_registry).run, CleanupProjectionGenerationsActivity(container.storage_backend).run, ProjectCanonicalRecordsActivity( projector=CanonicalProjector( diff --git a/backend/airweave/domains/temporal/workflows/recover_native_projection.py b/backend/airweave/domains/temporal/workflows/recover_native_projection.py index bc8f00bfa7..f637ce431a 100644 --- a/backend/airweave/domains/temporal/workflows/recover_native_projection.py +++ b/backend/airweave/domains/temporal/workflows/recover_native_projection.py @@ -1,4 +1,4 @@ -"""Recover native publications independently of unrelated cleanup outcomes.""" +"""Recover canonical publications independently of capture and cleanup outcomes.""" from datetime import timedelta @@ -7,7 +7,10 @@ from temporalio.exceptions import WorkflowAlreadyStartedError with workflow.unsafe.imports_passed_through(): - from airweave.domains.entities.canonical.projection_models import ProjectionSourcePage + from airweave.domains.entities.canonical.projection_models import ( + ProjectionSourcePage, + canonical_projection_workflow_id, + ) from airweave.domains.temporal.activities.discover_native_projection import ( DiscoverNativeProjectionActivity, ) @@ -39,7 +42,11 @@ async def run(self, after_sync_id: str | None = None) -> None: await workflow.start_child_workflow( ProjectCanonicalRecordsWorkflow.run, args=[str(source.organization_id), str(source.sync_id), None, 0, 0, True], - id=f"native-projection:{source.organization_id}:{source.sync_id}", + 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, ) diff --git a/backend/airweave/domains/temporal/workflows/run_source_connection.py b/backend/airweave/domains/temporal/workflows/run_source_connection.py index 3b10645b0f..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, @@ -212,14 +216,30 @@ async def _execute_sync( finally: # Starting the child is durable workflow history. Capture may be partial; # projection consumes only committed rows and owns its own retries/status. - 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, + 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_native_projection_recovery.py b/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py index 2ceeb29753..f621d831b8 100644 --- a/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py +++ b/backend/airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py @@ -117,7 +117,7 @@ async def generations(): active = await env.client.start_workflow( ProjectCanonicalRecordsWorkflow.run, args=[organization, sources[0], None, 0, 0, True], - id=f"native-projection:{organization}:{sources[0]}", + id=f"canonical-projection:{organization}:{sources[0]}", task_queue=queue, ) await asyncio.wait_for(started.wait(), timeout=20) @@ -131,7 +131,7 @@ async def generations(): await active.result() for sync in sources[1:]: await env.client.get_workflow_handle( - f"native-projection:{organization}:{sync}" + f"canonical-projection:{organization}:{sync}" ).result() finally: release.set() @@ -163,3 +163,146 @@ async def project(org: str, sync: str, after: str | None = None, skip_failed: bo 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/scripts/reproject_canonical.py b/backend/scripts/reproject_canonical.py index b6199a1df0..6fc056f37f 100644 --- a/backend/scripts/reproject_canonical.py +++ b/backend/scripts/reproject_canonical.py @@ -4,7 +4,8 @@ import asyncio from uuid import UUID -from temporalio.common import WorkflowIDConflictPolicy, WorkflowIDReusePolicy +from temporalio.common import WorkflowIDReusePolicy +from temporalio.exceptions import WorkflowAlreadyStartedError from airweave.core.config import settings from airweave.db.session import AsyncSessionLocal @@ -36,14 +37,19 @@ async def run(args: argparse.Namespace) -> None: # 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() - await client.start_workflow( - ProjectCanonicalRecordsWorkflow.run, - args=[str(args.organization), str(args.sync)], - id=plan.workflow_id, - task_queue=settings.TEMPORAL_TASK_QUEUE, - id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, - id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, - ) + 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.") From ab03a34d6a1b033eea5a57004501bd626872f009 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 13:54:55 -0700 Subject: [PATCH 249/266] Include projection admission workflow gates in worker checks --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 31c7d5d18e..51f408277a 100644 --- a/Makefile +++ b/Makefile @@ -54,7 +54,7 @@ test-health: test-worker: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_discover_native_projection.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_source_projection_replay.py airweave/domains/temporal/workflows/tests/test_reproject_command.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py test-ocr: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/ocr/tests airweave/core/container/tests/test_ocr_wiring.py From 46ca83d3c41df8416860d6c58674c3bedc12f4a9 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 14:15:28 -0700 Subject: [PATCH 250/266] Attest child inventories atomically with fresh source pages --- .../domains/entities/canonical/page_source.py | 11 +- .../domains/entities/canonical/scan_models.py | 14 + .../domains/entities/canonical/scan_store.py | 81 ++++- .../tests/test_page_child_receipts.py | 320 ++++++++++++++++++ .../tests/test_slack_file_recovery.py | 61 +++- .../canonical/tests/test_slack_recovery.py | 23 +- .../domains/sync_pipeline/canonical_scan.py | 1 + backend/airweave/platform/sources/slack.py | 21 +- .../platform/sources/test_slack_capture.py | 42 +++ 9 files changed, 539 insertions(+), 35 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/tests/test_page_child_receipts.py diff --git a/backend/airweave/domains/entities/canonical/page_source.py b/backend/airweave/domains/entities/canonical/page_source.py index 8181134055..efe1d3cd68 100644 --- a/backend/airweave/domains/entities/canonical/page_source.py +++ b/backend/airweave/domains/entities/canonical/page_source.py @@ -17,7 +17,11 @@ CompletedScope, ScopeRemovalReason, ) -from airweave.domains.entities.canonical.scan_models import ScanContinuation, ScanState +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 @@ -32,6 +36,11 @@ class CapturePage(BaseModel): 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 diff --git a/backend/airweave/domains/entities/canonical/scan_models.py b/backend/airweave/domains/entities/canonical/scan_models.py index cb1e9a63ee..2828e0e0d6 100644 --- a/backend/airweave/domains/entities/canonical/scan_models.py +++ b/backend/airweave/domains/entities/canonical/scan_models.py @@ -38,6 +38,19 @@ def bounded(cls, value: dict[str, JsonValue]) -> dict[str, JsonValue]: return value +class ChildScopeObservation(BaseModel): + """Current native page certifies this exact owner's complete child inventory. + + The provider validates inventory completeness, independently of unrelated body + coverage. SQL derives the receipt from the captured owner, never caller versions. + This declaration is page input, not a persisted queue or a resumed-page receipt. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + scope: CompletedScope + continuation: ScanContinuation = Field(default_factory=ScanContinuation) + + class BeginScan(BaseModel): """Resume a cycle, or explicitly CAS-restart its sweep after invalid continuation.""" @@ -100,6 +113,7 @@ class CommitScanPage(BaseModel): max_length=500, description="Verified independent originals from the bound source; never placeholders", ) + child_scope_observations: tuple[ChildScopeObservation, ...] = Field(default=(), max_length=500) continuation: ScanContinuation final: bool = False provider_checkpoint: ProviderCheckpoint | None = None diff --git a/backend/airweave/domains/entities/canonical/scan_store.py b/backend/airweave/domains/entities/canonical/scan_store.py index 80e1ff0ae1..ba6d8ef5f6 100644 --- a/backend/airweave/domains/entities/canonical/scan_store.py +++ b/backend/airweave/domains/entities/canonical/scan_store.py @@ -387,6 +387,19 @@ async def admit(self, db: AsyncSession, request: BeginScan) -> ScanAdmission: available = await db.scalar(select(content_is_available()).where(Entity.id == parent.id)) if parent.deleted_at is not None or not available: return ScanAdmission(state=None, parent=source_record(parent), capture=captured) + state = await self._begin_verified(db, request, context=(cursor, cycle, parent)) + return ScanAdmission(state=state, parent=source_record(parent), capture=captured) + + async def _begin_verified( + self, + db: AsyncSession, + request: BeginScan, + *, + context: tuple[SyncCursor, CaptureCycle, Entity], + ) -> ScanState: + """Share child admission after a fenced native owner observation is captured.""" + cursor, cycle, parent = context + row = await self._row(db, request.fence, request.scope) updated = request.model_copy( update={ "expected_parent_epoch": parent.visibility_epoch, @@ -401,9 +414,7 @@ async def admit(self, db: AsyncSession, request: BeginScan) -> ScanAdmission: row.parent_verified_revision = parent.record_revision row.revision += 1 await db.flush() - return ScanAdmission( - state=await self._state(db, row), parent=source_record(parent), capture=captured - ) + return await self._state(db, row) async def begin(self, db: AsyncSession, request: BeginScan) -> ScanState: """Ordinary admission cannot accept uncommitted owner observations.""" @@ -503,6 +514,68 @@ async def _begin( await db.flush() return await self._state(db, row) + @staticmethod + def _validate_child_observations(request: CommitScanPage, cycle: CaptureCycle) -> None: + """Only exact active owners in this native page can attest child inventories.""" + if not request.child_scope_observations: + return + if cycle.mode != "full": + raise ScanConflict("Page child observations require a full capture cycle") + records = {record.identity: record for record in request.records} + seen = set() + for observation in request.child_scope_observations: + scope = observation.scope + parent = records.get(scope.parent) + if ( + parent is None + or parent.kind != "upsert" + or parent.allow_reparent + or parent.removal_reason is not None + or parent.identity.record_type not in cycle.configuration.exact_parent_validation + or parent.identity.record_type + not in cycle.configuration.parents.get(scope.record_type, ()) + ): + raise ScanConflict( + "Child observation requires an allowed current-page upsert owner" + ) + key = (scope.parent, scope.record_type) + if key in seen: + raise ScanConflict("Duplicate child observation for the same owner and kind") + seen.add(key) + + async def _observe_child_scopes(self, db: AsyncSession, request: CommitScanPage) -> None: + """Derive existing child receipts inside the page's atomic writer transaction.""" + for observation in request.child_scope_observations: + # _begin may advance the cycle version; never reuse an earlier snapshot. + cursor, cycle = await attest_cycle(db, request.fence, request.cycle_id) + scope = observation.scope + await attest_scope(db, request.fence, cycle, scope, require_owner_receipt=False) + parent = await scope_owner(db, request.fence, cycle, scope) + if parent is None or parent.deleted_at is not None: + raise ScanConflict("Child observation owner is unavailable") + previous = await self._row(db, request.fence, scope) + request_begin = BeginScan( + fence=request.fence, + scope=scope, + cycle_id=request.cycle_id, + fingerprint=cycle.configuration.fingerprint, + expected=( + ScanVersion(sweep_id=previous.sweep_id, revision=previous.revision) + if previous is not None + else None + ), + restart=bool( + previous + and previous.cycle_id == request.cycle_id + and cycle.configuration.fresh_inventory(scope.record_type) + and previous.membership_attempt_id != request.fence.attempt_id + ), + expected_parent_epoch=parent.visibility_epoch, + expected_parent_revision=parent.record_revision, + continuation=observation.continuation, + ) + await self._begin_verified(db, request_begin, context=(cursor, cycle, parent)) + async def page( self, db: AsyncSession, @@ -555,6 +628,7 @@ async def page( for record in request.discovered_records ): raise ScanConflict("Discovered originals must be verified independent declared roots") + self._validate_child_observations(request, cycle) capture = capture_page or self.records._capture_locked captured = await capture( db, @@ -573,6 +647,7 @@ async def page( sequence=discovered.sequence, unchanged=captured.unchanged + discovered.unchanged, ) + await self._observe_child_scopes(db, request) row.continuation = continuation row.revision += 1 if request.final: 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_slack_file_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py index ba197455f2..7afbc97f97 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_file_recovery.py @@ -53,9 +53,9 @@ async def attempt( expected_team_id="T1", expected_user_id="U1", capture_files=True ) connector._verified_principal = SlackPrincipal(ok=True, team_id="T1", user_id="U1") - # Fresh attempts enumerate history and then revalidate the attachment owner. - # A retry reuses completed history and only refreshes the owner. - defaults = [ROOT, *([{"messages": [MESSAGE]}] if number == 1 else []), {"messages": [MESSAGE]}] + # 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) @@ -179,18 +179,21 @@ async def test_omitted_accessible_message_does_not_remove_retained_file_children 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 + 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) @@ -199,13 +202,17 @@ async def respond(url, params): if cursor == "H2": if interrupt: raise asyncio.CancelledError() - return {"messages": [], "response_metadata": {"next_cursor": "H3"}} + 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") @@ -219,6 +226,7 @@ async def respond(url, params): 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") @@ -293,3 +301,46 @@ async def capture(connector, scope, continuation, **kwargs): 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_recovery.py b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py index e3f82f53c2..081d306af0 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py +++ b/backend/airweave/domains/entities/canonical/tests/test_slack_recovery.py @@ -87,29 +87,12 @@ async def test_cancel_after_history_resumes_pending_reply_without_repeating_hist second, connector, _ = await runner( database, source, [ROOT, REPLIES], attempt=2, capture_files=capture_files ) - if capture_files: - initial = iter([ROOT, REPLIES]) - - async def respond(url, params): - if "oldest" not in params: - return next(initial) - assert params["oldest"] == params["latest"] - return { - "messages": [ - message for message in REPLIES["messages"] if message["ts"] == params["oldest"] - ] - } - - connector._get = AsyncMock(side_effect=respond) await run(second) operations = [call.args[0].split("/")[-1] for call in connector._get.call_args_list] assert operations[:2] == ["conversations.list", "conversations.replies"] - assert len(operations) == (4 if capture_files else 2) - if capture_files: - # Exact owner reads must not become another unbounded history traversal. - for call in connector._get.call_args_list[2:]: - assert call.args[1]["oldest"] == call.args[1]["latest"] - assert call.args[1]["limit"] == 1 + # 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 diff --git a/backend/airweave/domains/sync_pipeline/canonical_scan.py b/backend/airweave/domains/sync_pipeline/canonical_scan.py index f62ec216e5..42573091d6 100644 --- a/backend/airweave/domains/sync_pipeline/canonical_scan.py +++ b/backend/airweave/domains/sync_pipeline/canonical_scan.py @@ -277,6 +277,7 @@ async def scan( 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, diff --git a/backend/airweave/platform/sources/slack.py b/backend/airweave/platform/sources/slack.py index 511309d587..064a8e8c5d 100644 --- a/backend/airweave/platform/sources/slack.py +++ b/backend/airweave/platform/sources/slack.py @@ -39,7 +39,7 @@ RecordIdentity, parent_container_key, ) -from airweave.domains.entities.canonical.scan_models import ScanContinuation +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, @@ -338,14 +338,18 @@ def _capture_message(message: dict, channel_id: str) -> CaptureRecord: def child_scope(self, parent: SourceRecord, record_type: str) -> CompletedScope: """Keep the established native container identity for this flat provider.""" - if self.canonical_container_parents.get(record_type) != parent.identity.record_type: + 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.identity) + container_id=parent_container_key(parent) if record_type == "file" - else parent.identity.native_id, - parent=parent.identity, + else parent.native_id, + parent=parent, ) async def capture_page( @@ -363,9 +367,14 @@ async def capture_page( 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": tuple(self._attachment_owner(record) for record in page.records) + "records": records, + "child_scope_observations": tuple( + ChildScopeObservation(scope=self._child_scope(record.identity, "file")) + for record in records + ), } ) return page diff --git a/backend/tests/unit/platform/sources/test_slack_capture.py b/backend/tests/unit/platform/sources/test_slack_capture.py index cf5bc0c1c8..cbd541481f 100644 --- a/backend/tests/unit/platform/sources/test_slack_capture.py +++ b/backend/tests/unit/platform/sources/test_slack_capture.py @@ -292,3 +292,45 @@ async def test_exact_attachment_owner_refresh_preserves_native_inventory_and_par connector._get = AsyncMock(return_value={"messages": messages}) with pytest.raises(ValueError, match=error): await connector.refresh_known(record, files=MagicMock()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_files", [None, [], [{"id": "F1"}]]) +async def test_fresh_message_page_declares_exact_file_inventory_scope(native_files): + connector = await source() + connector.slack_config = SlackConfig( + expected_team_id="T1", expected_user_id="U1", capture_files=True + ) + native = {"ts": "1", "text": "Retained unchanged"} + if native_files is not None: + native["files"] = native_files + connector._get = AsyncMock(return_value={"messages": [native]}) + page = await connector.capture_page( + CompletedScope(record_type="message", container_id="C1"), + ScanContinuation(), + files=MagicMock(), + ) + assert page.records[0].payload == native + (declaration,) = page.child_scope_observations + owner = SourceRecord.model_construct(identity=page.records[0].identity) + assert declaration.scope == connector.child_scope(owner, "file") + assert declaration.continuation == ScanContinuation() + assert page.records[0].payload_schema_version == 2 + connector._get.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_invalid_attachment_inventory_cannot_emit_fresh_scope_receipt(): + connector = await source() + connector.slack_config = SlackConfig( + expected_team_id="T1", expected_user_id="U1", capture_files=True + ) + connector._get = AsyncMock( + return_value={"messages": [{"ts": "1", "files": [{"id": "F1"}, {"id": "F1"}]}]} + ) + with pytest.raises(ValueError, match="duplicate file identities"): + await connector.capture_page( + CompletedScope(record_type="message", container_id="C1"), + ScanContinuation(), + files=MagicMock(), + ) From f10e456e5e633219842b07d0c56e980206ed984c Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 14:34:20 -0700 Subject: [PATCH 251/266] Wake and serialize worker shutdown across startup and signal handlers --- Makefile | 2 +- .../domains/temporal/worker/__init__.py | 122 +++++++++------ .../temporal/worker/tests/test_worker.py | 146 +++++++++++++++++- .../worker/tests/test_worker_ocr_guard.py | 5 +- 4 files changed, 217 insertions(+), 58 deletions(-) diff --git a/Makefile b/Makefile index 51f408277a..034e83be98 100644 --- a/Makefile +++ b/Makefile @@ -54,7 +54,7 @@ test-health: test-worker: @test -n "$$CANONICAL_TEST_DATABASE_URL" || (echo 'Set CANONICAL_TEST_DATABASE_URL to a disposable PostgreSQL database.' >&2; exit 1) - cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_discover_native_projection.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_source_projection_replay.py airweave/domains/temporal/workflows/tests/test_reproject_command.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py + cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/temporal/activities/tests/test_project_canonical_records.py airweave/domains/temporal/activities/tests/test_discover_native_projection.py airweave/domains/temporal/activities/tests/test_cleanup_stuck_sync_jobs.py airweave/domains/temporal/workflows/tests/test_native_projection_recovery.py airweave/domains/temporal/workflows/tests/test_source_projection_replay.py airweave/domains/temporal/workflows/tests/test_reproject_command.py airweave/domains/temporal/workflows/tests/test_cleanup_workflows.py airweave/domains/temporal/workflows/tests/test_canonical_projection.py airweave/domains/temporal/worker/tests/test_wiring.py airweave/domains/temporal/worker/tests/test_config.py airweave/domains/temporal/worker/tests/test_worker.py airweave/domains/temporal/worker/tests/test_worker_ocr_guard.py test-ocr: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/ocr/tests airweave/core/container/tests/test_ocr_wiring.py diff --git a/backend/airweave/domains/temporal/worker/__init__.py b/backend/airweave/domains/temporal/worker/__init__.py index 077aafdef2..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 @@ -69,6 +68,8 @@ def __init__(self, config: WorkerConfig) -> None: ) 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}") @@ -186,12 +201,13 @@ async def main() -> None: 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/tests/test_worker.py b/backend/airweave/domains/temporal/worker/tests/test_worker.py index 4bbf121279..c690a3ec9e 100644 --- a/backend/airweave/domains/temporal/worker/tests/test_worker.py +++ b/backend/airweave/domains/temporal/worker/tests/test_worker.py @@ -5,6 +5,7 @@ WorkerMetricsRegistry) is constructed for real so tests exercise actual wiring. """ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -17,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) @@ -206,3 +207,138 @@ async def test_control_server_binds_only_loopback(mock_runtime_cls): 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 2ecc85cbaf..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,5 +1,6 @@ """OCR capability does not prevent the real worker composition from starting.""" +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -61,7 +62,9 @@ def health_client(*args, **kwargs): monkeypatch.setattr(container_mod, "container", None) lifecycle = MagicMock(start=AsyncMock(), stop=AsyncMock()) monkeypatch.setattr(worker, "TemporalWorker", lambda _: lifecycle) - monkeypatch.setattr(worker.signal, "signal", lambda *_: None) + 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() From 380b4dbce9504e1f61c5b112037f5958b0a37aaa Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 16:05:54 -0700 Subject: [PATCH 252/266] Evaluate delivered search cards against frozen source identities --- Makefile | 1 + .../unit/search/test_owned_evaluation.py | 153 ++++++++++++++++++ evaluation/README.md | 31 ++++ evaluation/owned_retrieval.py | 130 +++++++++++++++ evaluation/retrieval.py | 35 +++- evaluation/tests/test_retrieval.py | 28 ++++ 6 files changed, 375 insertions(+), 3 deletions(-) create mode 100644 backend/tests/unit/search/test_owned_evaluation.py create mode 100644 evaluation/owned_retrieval.py diff --git a/Makefile b/Makefile index 034e83be98..57a5939357 100644 --- a/Makefile +++ b/Makefile @@ -45,6 +45,7 @@ test-provisioning: test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search airweave/domains/embedders + PYTHONPATH=$(BACKEND):. $(BACKEND)/$(PYTEST) -q -o log_cli=false $(BACKEND)/tests/unit/search/test_owned_evaluation.py test-auth: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/api/tests/test_service_auth.py airweave/api/tests/test_sync_authorization.py airweave/api/tests/test_context_resolver_auth.py diff --git a/backend/tests/unit/search/test_owned_evaluation.py b/backend/tests/unit/search/test_owned_evaluation.py new file mode 100644 index 0000000000..10d96782f2 --- /dev/null +++ b/backend/tests/unit/search/test_owned_evaluation.py @@ -0,0 +1,153 @@ +"""Synthetic delivery boundaries, independent of corpus indexing and credentials.""" + +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest +from evaluation.owned_retrieval import ConversationIdentity, CorpusRecord, delivered_result + +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.search.owned_models import ( + OwnedSearchGroup, + OwnedSearchHit, + OwnedSearchMatch, + OwnedSearchResponse, +) + + +def match(source_id="account-one", sync_id=None, native_id="message", provider="almanac"): + """Build a match plus its separately retained census identity.""" + hit = OwnedSearchHit( + record_id=uuid4(), + revision=1, + sync_id=sync_id or uuid4(), + source_connection_id=uuid4(), + provider=provider, + identity=RecordIdentity(record_type="message", native_id=native_id, container_id="session"), + title="Synthetic", + excerpts=("Synthetic passage",), + observed_at=datetime.now(timezone.utc), + source_created_at=None, + source_updated_at=None, + completeness="complete", + ) + record = CorpusRecord( + record_id=hit.record_id, + sync_id=hit.sync_id, + revision=hit.revision, + source_id=source_id, + provider=provider, + identity=hit.identity, + ) + return hit, record + + +def response(*hits, incomplete=False): + """A delivered bounded response, with no hidden members in items.""" + return OwnedSearchResponse( + items=hits, + sources=(), + candidate_window_full=False, + engine_partial=False, + excluded_candidates=0, + postfilter_excluded=0, + retrieval_incomplete=incomplete, + ) + + +def test_cards_and_displayed_originals_are_different_ranked_units(): + first, first_record = match() + second, second_record = match(sync_id=first.sync_id, native_id="second") + third, third_record = match(native_id="third") + first.group = OwnedSearchGroup( + kind="session", + native_id="session", + matched_records=17, + additional_matches=( + OwnedSearchMatch.model_validate(second.model_dump(exclude={"group"}, round_trip=True)), + ), + ) + membership = ConversationIdentity(kind="session", native_id="session") + first_record = first_record.model_copy(update={"conversation": membership}) + second_record = second_record.model_copy(update={"conversation": membership}) + records = (first_record, second_record, third_record) + delivered = response(first, third, incomplete=True) + cards = delivered_result("q", delivered, records, unit="card") + originals = delivered_result("q", delivered, records, unit="displayed_original") + assert cards.record_ids == ( + first_record.card_id, + third_record.evaluation_id, + ) + assert originals.record_ids == tuple(record.evaluation_id for record in records) + assert cards.status == originals.status == "partial" + assert len(originals.record_ids) == 3 # matched_records is not a delivered member list. + + +def test_card_identity_is_source_scoped_and_stable_across_destination_rebuilds(): + first, first_record = match() + other, other_record = match(source_id="account-two") + rebuilt, rebuilt_record = match() + for hit in (first, other, rebuilt): + hit.group = OwnedSearchGroup(kind="session", native_id="session", matched_records=1) + membership = ConversationIdentity(kind="session", native_id="session") + first_record = first_record.model_copy(update={"conversation": membership}) + other_record = other_record.model_copy(update={"conversation": membership}) + rebuilt_record = rebuilt_record.model_copy(update={"conversation": membership}) + result = delivered_result( + "q", response(first, other), (first_record, other_record), unit="card" + ) + rerun = delivered_result("q", response(rebuilt), (rebuilt_record,), unit="card") + assert len(set(result.record_ids)) == 2 + assert result.record_ids[0] == rerun.record_ids[0] + assert first_record.evaluation_id == rebuilt_record.evaluation_id + assert first_record.evaluation_id != other_record.evaluation_id + + +def test_stale_or_wrong_source_matches_cannot_be_scored_as_frozen_records(): + first, record = match() + with pytest.raises(ValueError, match="frozen source identity"): + delivered_result( + "q", response(first.model_copy(update={"revision": 2})), (record,), unit="card" + ) + with pytest.raises(ValueError, match="Duplicate destination"): + delivered_result("q", response(first), (record, record), unit="card") + second, second_record = match(source_id="account-two") + first.group = OwnedSearchGroup( + kind="session", + native_id="session", + matched_records=2, + additional_matches=(OwnedSearchMatch.model_validate(second.model_dump(exclude={"group"})),), + ) + record = record.model_copy( + update={"conversation": ConversationIdentity(kind="session", native_id="session")} + ) + with pytest.raises(ValueError, match="different source identity"): + delivered_result("q", response(first), (record, second_record), unit="card") + + +@pytest.mark.parametrize("violation", ("wrong_card", "missing_attestation", "foreign_member")) +def test_group_membership_must_match_frozen_census(violation): + first, first_record = match() + second, second_record = match(sync_id=first.sync_id, native_id="second") + membership = ConversationIdentity(kind="session", native_id="session") + first_record = first_record.model_copy(update={"conversation": membership}) + second_record = second_record.model_copy(update={"conversation": membership}) + first.group = OwnedSearchGroup( + kind="session", + native_id="session", + matched_records=2, + additional_matches=(OwnedSearchMatch.model_validate(second.model_dump(exclude={"group"})),), + ) + if violation == "wrong_card": + first.group = first.group.model_copy(update={"native_id": "different-session"}) + elif violation == "missing_attestation": + first_record = first_record.model_copy(update={"conversation": None}) + else: + second_record = second_record.model_copy( + update={ + "conversation": ConversationIdentity(kind="session", native_id="different-session") + } + ) + for unit in ("card", "displayed_original"): + with pytest.raises(ValueError, match="frozen conversation membership"): + delivered_result("q", response(first), (first_record, second_record), unit=unit) diff --git a/evaluation/README.md b/evaluation/README.md index 85a74f5768..f5c8eef8f6 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -213,3 +213,34 @@ revisions, model settings and assessment provenance fixed for both experiments. The one-session trial in the product worklog is diagnostic, not a judged benchmark: it found the session in every mode, but the expected passage only in keyword and hybrid mode. More diverse sessions and independently judged queries are required. + +`evaluation.owned_retrieval.delivered_result` exports these two units from the +existing typed `OwnedSearchResponse`. Supply a frozen tuple of `CorpusRecord` +locators with destination record/sync/revision, preserved provider identity, and +a stable original source ID (account or attested owner/dataset). The adapter +checks each returned original against the census and hashes source plus native +identity for rebuild-stable evaluation IDs. Conversation IDs additionally bind +group kind and native ID; accounts with the same native thread ID stay separate. +Attest optional `ConversationIdentity` from canonical parent identity or a +validated retained email thread before retrieval. A returned group or additional +member contradicting that frozen membership is rejected. Use `CorpusRecord.card_id` +for card judgments, independent of which member becomes the representative. +Use `unit="card"` for primary UI/API ranking and `unit="displayed_original"` +for the explicitly flattened diagnostic. Returned match counts never invent +hidden original members. An incomplete response retains its ranking and is +marked partial. This helper requires the backend environment; it does not query +the service or verify publication by itself. + +`evaluation.retrieval.pool_unjudged(dataset, runs)` returns each query's sorted +union of delivered IDs without existing judgments. It checks the exact dataset +and corpus binding for every run. Assess that private pool before comparing +metrics; membership in the pool is neither a positive nor a negative label. +If labels change, deliberately regenerate the dataset fingerprint and rebind +the preserved observations; never rerun queries selectively to improve scores. + +The delivery adapter's synthetic tests use backend fixtures: + +```sh +PYTHONPATH=backend:. backend/.venv/bin/python -m pytest \ + backend/tests/unit/search/test_owned_evaluation.py -q +``` diff --git a/evaluation/owned_retrieval.py b/evaluation/owned_retrieval.py new file mode 100644 index 0000000000..a6ab89a53e --- /dev/null +++ b/evaluation/owned_retrieval.py @@ -0,0 +1,130 @@ +"""Export existing owned-search responses for the offline relevance evaluator. + +Requires the backend environment. Source identities come from the frozen corpus +census; destination UUIDs are locators, never cross-rebuild evaluation identities. +""" + +import hashlib +import json +from typing import Literal +from uuid import UUID + +from pydantic import Field + +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.search.owned_models import OwnedSearchMatch, OwnedSearchResponse +from evaluation.retrieval import Identifier, Result, Value + + +def _identifier(kind: str, values: object) -> str: + encoded = json.dumps( + values, sort_keys=True, ensure_ascii=False, separators=(",", ":") + ) + return kind + ":" + hashlib.sha256(encoded.encode()).hexdigest() + + +class ConversationIdentity(Value): + """Frozen grouping attestation from retained canonical identity, before retrieval.""" + + kind: Literal["session", "email_thread"] + native_id: str = Field(min_length=1) + + +class CorpusRecord(Value): + """One exact destination locator and its preserved stable source identity.""" + + record_id: UUID + sync_id: UUID + revision: int = Field(ge=1) + source_id: Identifier + provider: Identifier + identity: RecordIdentity + conversation: ConversationIdentity | None = None + + @property + def evaluation_id(self) -> str: + """Keep native IDs account-scoped, including their original container.""" + return _identifier( + "original", + [self.source_id, self.provider, self.identity.model_dump(mode="json")], + ) + + @property + def card_id(self) -> str: + """Use frozen membership even when a different original wins the ranking.""" + if self.conversation is None: + return self.evaluation_id + return _identifier( + "card", + [ + self.source_id, + self.provider, + self.conversation.kind, + self.conversation.native_id, + ], + ) + + +def delivered_result( + query_id: str, + response: OwnedSearchResponse, + records: tuple[CorpusRecord, ...], + *, + unit: Literal["card", "displayed_original"], +) -> Result: + """Preserve card order; flatten only explicitly delivered originals when requested.""" + census = {record.record_id: record for record in records} + if len(census) != len(records): + raise ValueError("Duplicate destination record in corpus census") + + def reference(match: OwnedSearchMatch) -> CorpusRecord: + record = census[match.record_id] + if ( + record.sync_id != match.sync_id + or record.revision != match.revision + or record.provider != match.provider + or record.identity != match.identity + ): + raise ValueError("Delivered match conflicts with frozen source identity") + return record + + ranking = [] + for hit in response.items: + original = reference(hit) + conversation = ( + None + if hit.group is None + else ConversationIdentity( + kind=hit.group.kind, native_id=hit.group.native_id + ) + ) + if conversation != original.conversation: + raise ValueError( + "Delivered group conflicts with frozen conversation membership" + ) + additional = () if hit.group is None else hit.group.additional_matches + for match in additional: + member = reference(match) + if (member.source_id, member.provider, member.sync_id) != ( + original.source_id, + original.provider, + original.sync_id, + ): + raise ValueError("Conversation contains a different source identity") + if member.conversation != conversation: + raise ValueError( + "Additional match conflicts with frozen conversation membership" + ) + if unit == "card" and hit.group is not None: + ranking.append(original.card_id) + else: + ranking.append(original.evaluation_id) + if unit == "displayed_original": + ranking.extend(reference(match).evaluation_id for match in additional) + return Result( + query_id=query_id, + status="partial" + if response.retrieval_incomplete or response.engine_partial + else "success", + record_ids=tuple(ranking), + ) diff --git a/evaluation/retrieval.py b/evaluation/retrieval.py index c2db8f5726..f7b3c37fcd 100644 --- a/evaluation/retrieval.py +++ b/evaluation/retrieval.py @@ -14,7 +14,6 @@ from statistics import mean from typing import Annotated, Literal -import ir_measures from pydantic import BaseModel, ConfigDict, Field, model_validator Identifier = Annotated[str, Field(min_length=1, pattern=r"^\S+$")] @@ -209,8 +208,8 @@ class Report(Value): warnings: tuple[str, ...] -def summarize(dataset: Dataset, run: Run) -> Report: - """Use established IR measures, preserving API rank order with unique synthetic scores.""" +def validate_run(dataset: Dataset, run: Run) -> None: + """Require every observation to belong to the exact frozen evaluation dataset.""" if ( run.corpus_id != dataset.corpus_id or run.dataset_sha256 != dataset.fingerprint() @@ -223,6 +222,36 @@ def summarize(dataset: Dataset, run: Run) -> Report: ) if any(t.query_id not in query_ids for t in run.timings): raise ValueError("Timing references an unknown query") + + +def pool_unjudged( + dataset: Dataset, runs: tuple[Run, ...] +) -> dict[str, tuple[str, ...]]: + """Union delivered candidates needing assessment; never invent negative labels.""" + for run in runs: + validate_run(dataset, run) + return { + query.id: tuple( + sorted( + { + record + for run in runs + for result in run.results + if result.query_id == query.id + for record in result.record_ids + } + - {judgment.record_id for judgment in query.judgments} + ) + ) + for query in dataset.queries + } + + +def summarize(dataset: Dataset, run: Run) -> Report: + """Use established IR measures, preserving API rank order with unique synthetic scores.""" + import ir_measures + + validate_run(dataset, run) measures = [ ir_measures.nDCG @ 10, ir_measures.RR @ 10, diff --git a/evaluation/tests/test_retrieval.py b/evaluation/tests/test_retrieval.py index b85bc35a0c..be54f222fd 100644 --- a/evaluation/tests/test_retrieval.py +++ b/evaluation/tests/test_retrieval.py @@ -11,6 +11,7 @@ Result, Run, Timing, + pool_unjudged, summarize, ) @@ -145,6 +146,33 @@ def test_missing_query_or_changed_labels_cannot_improve_score_silently(): ) +def test_pool_unions_only_unjudged_delivered_records_and_preserves_dataset_binding(): + dataset = corpus() + original = run(dataset) + additional = run( + dataset, + results=( + Result( + query_id="known", + status="success", + record_ids=("meeting", "other", "unknown"), + ), + Result(query_id="failed", status="timeout"), + Result(query_id="absent", status="success", record_ids=("candidate",)), + ), + ) + assert pool_unjudged(dataset, (original, additional)) == { + "known": ("other", "unknown"), + "failed": (), + "absent": ("candidate",), + } + assert not any( + j.record_id == "unknown" for q in dataset.queries for j in q.judgments + ) + with pytest.raises(ValueError, match="exact dataset"): + pool_unjudged(dataset.model_copy(update={"version": "changed"}), (original,)) + + def test_no_answer_success_is_separate_from_answerable_metrics(): dataset = Dataset(corpus_id="empty", version="v1", queries=(corpus().queries[2],)) observation = Run( From 7c7d06a58e6dec4d77ae56498ee4dfa59cefd10f Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 16:51:09 -0700 Subject: [PATCH 253/266] Add frozen-corpus HTTP search replay and pure wire models --- Makefile | 2 +- .../domains/entities/canonical/coverage.py | 48 +- .../entities/canonical/coverage_models.py | 47 ++ .../entities/canonical/cycle_models.py | 3 +- .../entities/canonical/extraction_models.py | 14 +- .../domains/entities/canonical/models.py | 13 + .../entities/canonical/query_models.py | 14 +- .../airweave/domains/search/owned_models.py | 4 +- .../domains/search/retrieval_strategy.py | 11 + .../airweave/domains/search/types/__init__.py | 2 +- backend/airweave/domains/search/types/plan.py | 10 +- .../unit/search/test_owned_evaluation.py | 43 ++ .../tests/unit/search/test_owned_replay.py | 520 ++++++++++++++++++ evaluation/README.md | 69 ++- evaluation/replay.py | 350 ++++++++++++ evaluation/replay_cli.py | 156 ++++++ 16 files changed, 1234 insertions(+), 72 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/coverage_models.py create mode 100644 backend/airweave/domains/search/retrieval_strategy.py create mode 100644 backend/tests/unit/search/test_owned_replay.py create mode 100644 evaluation/replay.py create mode 100644 evaluation/replay_cli.py diff --git a/Makefile b/Makefile index 57a5939357..edf47e04fc 100644 --- a/Makefile +++ b/Makefile @@ -45,7 +45,7 @@ test-provisioning: test-search: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/domains/search airweave/domains/embedders - PYTHONPATH=$(BACKEND):. $(BACKEND)/$(PYTEST) -q -o log_cli=false $(BACKEND)/tests/unit/search/test_owned_evaluation.py + PYTHONPATH=$(BACKEND):. $(BACKEND)/$(PYTEST) -q -o log_cli=false $(BACKEND)/tests/unit/search/test_owned_evaluation.py $(BACKEND)/tests/unit/search/test_owned_replay.py test-auth: cd $(BACKEND) && $(PYTEST) -q -o log_cli=false airweave/api/tests/test_service_auth.py airweave/api/tests/test_sync_authorization.py airweave/api/tests/test_context_resolver_auth.py diff --git a/backend/airweave/domains/entities/canonical/coverage.py b/backend/airweave/domains/entities/canonical/coverage.py index ce9967c6a0..1e2655b0f8 100644 --- a/backend/airweave/domains/entities/canonical/coverage.py +++ b/backend/airweave/domains/entities/canonical/coverage.py @@ -1,16 +1,19 @@ """Describe persisted capture promises independently of indexing counts or job success.""" -from typing import Literal from uuid import UUID -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError, model_validator +from pydantic import ValidationError from sqlalchemy import and_, case, func, or_, select from sqlalchemy.ext.asyncio import AsyncSession +from airweave.domains.entities.canonical.coverage_models import ( + CaptureCoverage, + CaptureScopeSummary, + FullCaptureCoverage, +) from airweave.domains.entities.canonical.cycle_models import ( CYCLE_KEY, CaptureCycle, - CompletionPolicy, ) from airweave.domains.entities.canonical.requests import WriterFence from airweave.domains.entities.canonical.store import content_is_available @@ -20,45 +23,6 @@ from airweave.models.sync_cursor import SyncCursor -class FullCaptureCoverage(BaseModel): - """Public coverage evidence excludes private provider continuation values.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - cycle_id: UUID - completed_at: AwareDatetime - discovery: Literal["incomplete", "scope_enumeration_complete"] - - -class CaptureScopeSummary(BaseModel): - """Counts describe eligible exact scopes, never original/indexed record totals.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - eligible: int = Field(ge=0, strict=True) - completed_full: int = Field(ge=0, strict=True) - completed_changes: int = Field(ge=0, strict=True) - unfinished: int = Field(ge=0, strict=True) - - @model_validator(mode="after") - def partition(self): - """Every eligible scope belongs to exactly one progress category.""" - if self.eligible != self.completed_full + self.completed_changes + self.unfinished: - raise ValueError("Scope progress must partition eligible scopes") - return self - - -class CaptureCoverage(BaseModel): - """No cycle means unknown; completion only certifies the declared scope policies.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - phase: Literal["active", "complete"] - mode: Literal["full", "changes", "mixed"] = "full" - scope_summary: CaptureScopeSummary | None = None - last_full_capture: FullCaptureCoverage | None = None - provider_checkpoint_promoted_at: AwareDatetime | None = None - policies: dict[str, CompletionPolicy] - discovery: Literal["incomplete", "pending", "scope_enumeration_complete"] - - async def mixed_scope_summary( db: AsyncSession, organization_id: UUID, sync_id: UUID, cycle: CaptureCycle ) -> CaptureScopeSummary | None: diff --git a/backend/airweave/domains/entities/canonical/coverage_models.py b/backend/airweave/domains/entities/canonical/coverage_models.py new file mode 100644 index 0000000000..5f54350506 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/coverage_models.py @@ -0,0 +1,47 @@ +"""Public capture coverage facts, independent of SQL and runtime settings.""" + +from typing import Literal +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator + +CompletionPolicy = Literal["exhaustive", "discovery_only", "discovery_with_validation"] + + +class FullCaptureCoverage(BaseModel): + """Public coverage evidence excludes private provider continuation values.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + cycle_id: UUID + completed_at: AwareDatetime + discovery: Literal["incomplete", "scope_enumeration_complete"] + + +class CaptureScopeSummary(BaseModel): + """Counts describe eligible exact scopes, never original/indexed record totals.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + eligible: int = Field(ge=0, strict=True) + completed_full: int = Field(ge=0, strict=True) + completed_changes: int = Field(ge=0, strict=True) + unfinished: int = Field(ge=0, strict=True) + + @model_validator(mode="after") + def partition(self): + """Every eligible scope belongs to exactly one progress category.""" + if self.eligible != self.completed_full + self.completed_changes + self.unfinished: + raise ValueError("Scope progress must partition eligible scopes") + return self + + +class CaptureCoverage(BaseModel): + """No cycle means unknown; completion only certifies the declared scope policies.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + phase: Literal["active", "complete"] + mode: Literal["full", "changes", "mixed"] = "full" + scope_summary: CaptureScopeSummary | None = None + last_full_capture: FullCaptureCoverage | None = None + provider_checkpoint_promoted_at: AwareDatetime | None = None + policies: dict[str, CompletionPolicy] + discovery: Literal["incomplete", "pending", "scope_enumeration_complete"] diff --git a/backend/airweave/domains/entities/canonical/cycle_models.py b/backend/airweave/domains/entities/canonical/cycle_models.py index 3715f42c71..26b0912ed7 100644 --- a/backend/airweave/domains/entities/canonical/cycle_models.py +++ b/backend/airweave/domains/entities/canonical/cycle_models.py @@ -16,11 +16,10 @@ model_validator, ) +from airweave.domains.entities.canonical.coverage_models import CompletionPolicy from airweave.domains.entities.canonical.models import SourceRecord from airweave.domains.entities.canonical.requests import RecordKind, ScanVersion, WriterFence -CompletionPolicy = Literal["exhaustive", "discovery_only", "discovery_with_validation"] - CaptureMode = Literal["full", "changes", "mixed"] CompletedDiscovery = Literal["incomplete", "scope_enumeration_complete"] diff --git a/backend/airweave/domains/entities/canonical/extraction_models.py b/backend/airweave/domains/entities/canonical/extraction_models.py index f51778c655..cda9502fe2 100644 --- a/backend/airweave/domains/entities/canonical/extraction_models.py +++ b/backend/airweave/domains/entities/canonical/extraction_models.py @@ -1,6 +1,6 @@ """Generation-owned extraction evidence, separate from captured original completeness.""" -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator @@ -51,6 +51,18 @@ class ExtractionCoverage(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) parts: tuple[ExtractionOutcome, ...] + @model_validator(mode="before") + @classmethod + def validate_serialized_status(cls, value: Any) -> Any: + """Accept wire status only as a checked attestation of the underlying parts.""" + if not isinstance(value, dict) or "status" not in value: + return value + facts = {key: item for key, item in value.items() if key != "status"} + coverage = cls.model_validate(facts) + if value["status"] != coverage.status: + raise ValueError("Serialized extraction status conflicts with its parts") + return facts + @model_validator(mode="after") def unique_parts(self): """Each stable expected part appears once, with no ordinal gaps.""" diff --git a/backend/airweave/domains/entities/canonical/models.py b/backend/airweave/domains/entities/canonical/models.py index 9560c22cfe..04b9dae568 100644 --- a/backend/airweave/domains/entities/canonical/models.py +++ b/backend/airweave/domains/entities/canonical/models.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue +from airweave.domains.entities.canonical.coverage_models import CaptureCoverage from airweave.domains.entities.canonical.requests import BlobReference, RecordIdentity @@ -68,3 +69,15 @@ class ChangePage(BaseModel): 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/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index e9360a7d55..a100eaa730 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -6,9 +6,9 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue -from airweave.domains.entities.canonical.coverage import CaptureCoverage from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord +from airweave.domains.entities.canonical.models import RecordPage as RecordPage from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.platform.sources.records.sheets_manifest import GridBounds, GridGap from airweave.platform.sources.records.sheets_models import SpreadsheetCell @@ -73,18 +73,6 @@ class RecordListQuery(BaseModel): cursor: str | None = None -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" - - class RecordChangePage(BaseModel): """A bounded journal window plus a cursor for future polling.""" diff --git a/backend/airweave/domains/search/owned_models.py b/backend/airweave/domains/search/owned_models.py index 817b060a65..5800a99c4f 100644 --- a/backend/airweave/domains/search/owned_models.py +++ b/backend/airweave/domains/search/owned_models.py @@ -5,10 +5,10 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator -from airweave.domains.entities.canonical.coverage import CaptureCoverage +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.types import RetrievalStrategy +from airweave.domains.search.retrieval_strategy import RetrievalStrategy class OwnedSearchRequest(BaseModel): 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/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/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/tests/unit/search/test_owned_evaluation.py b/backend/tests/unit/search/test_owned_evaluation.py index 10d96782f2..82e4673e49 100644 --- a/backend/tests/unit/search/test_owned_evaluation.py +++ b/backend/tests/unit/search/test_owned_evaluation.py @@ -6,6 +6,10 @@ import pytest from evaluation.owned_retrieval import ConversationIdentity, CorpusRecord, delivered_result +from airweave.domains.entities.canonical.extraction_models import ( + ExtractionCoverage, + ExtractionOutcome, +) from airweave.domains.entities.canonical.requests import RecordIdentity from airweave.domains.search.owned_models import ( OwnedSearchGroup, @@ -151,3 +155,42 @@ def test_group_membership_must_match_frozen_census(violation): for unit in ("card", "displayed_original"): with pytest.raises(ValueError, match="frozen conversation membership"): delivered_result("q", response(first), (first_record, second_record), unit=unit) + + +@pytest.mark.parametrize("location", ("representative", "additional")) +@pytest.mark.parametrize("violation", (None, "wrong_status", "unknown_status", "unknown_field")) +def test_actual_search_wire_extraction_roundtrip_stays_strict(location, violation): + """The normal API dump includes computed status, including additional matches.""" + coverage = ExtractionCoverage( + parts=(ExtractionOutcome(part_index=0, key="body", kind="body", outcome="indexed"),) + ) + first, _ = match() + second, _ = match(sync_id=first.sync_id, native_id="second") + first.extraction = second.extraction = coverage + first.group = OwnedSearchGroup( + kind="session", + native_id="session", + matched_records=2, + additional_matches=(OwnedSearchMatch.model_validate(second.model_dump(exclude={"group"})),), + ) + wire = response(first).model_dump(mode="json") + extraction = ( + wire["items"][0]["extraction"] + if location == "representative" + else wire["items"][0]["group"]["additional_matches"][0]["extraction"] + ) + assert extraction["status"] == "complete" + if violation == "wrong_status": + extraction["status"] = "partial" + elif violation == "unknown_status": + extraction["status"] = "invented" + elif violation == "unknown_field": + extraction["future_field"] = True + if violation is None: + parsed = OwnedSearchResponse.model_validate(wire) + assert parsed.model_dump(mode="json") == wire + assert ExtractionCoverage.model_validate(coverage.persisted()).status == "complete" + assert "status" not in coverage.persisted() + else: + with pytest.raises(ValueError): + OwnedSearchResponse.model_validate(wire) diff --git a/backend/tests/unit/search/test_owned_replay.py b/backend/tests/unit/search/test_owned_replay.py new file mode 100644 index 0000000000..e2248e8c7a --- /dev/null +++ b/backend/tests/unit/search/test_owned_replay.py @@ -0,0 +1,520 @@ +"""HTTP replay boundaries; synthetic fixtures do not qualify a real corpus.""" + +import asyncio +import json +import os +import subprocess +import sys +from pathlib import Path +from uuid import UUID + +import httpx +import pytest +from evaluation.native_import_cli import Settings +from evaluation.owned_retrieval import ConversationIdentity +from evaluation.replay import ( + CorpusMismatch, + FrozenCorpus, + FrozenRecord, + ReplayConfiguration, + replay, +) +from evaluation.replay_cli import ReplayArguments, execute, main +from evaluation.retrieval import Dataset, Judgment, Query + +from airweave.domains.entities.canonical.extraction_models import ( + ExtractionCoverage, + ExtractionOutcome, +) +from airweave.domains.entities.canonical.models import RecordPage, SourceRecord +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.search.owned_models import ( + OwnedSearchGroup, + OwnedSearchHit, + OwnedSearchMatch, + OwnedSearchResponse, +) + + +def fixture(): + """Two delivered members with frozen capture, group and extraction facts.""" + sync_id = UUID(int=100) + parent = RecordIdentity(record_type="session", native_id="session") + extraction = ExtractionCoverage( + parts=(ExtractionOutcome(part_index=0, key="body", kind="body", outcome="indexed"),) + ) + records = tuple( + FrozenRecord( + record_id=UUID(int=i), + sync_id=sync_id, + revision=3, + source_id="synthetic-source", + provider="almanac", + identity=RecordIdentity( + record_type="message", native_id=str(i), container_id="session" + ), + conversation=ConversationIdentity(kind="session", native_id="session"), + capture_hash="a" * 64, + indexed_pipeline_version=2, + extraction=extraction, + parent=parent, + ) + for i in (1, 2) + ) + corpus = FrozenCorpus( + corpus_id="synthetic-corpus", + organization_id=UUID(int=200), + sync_ids=(sync_id,), + records=records, + ) + dataset = Dataset( + corpus_id=corpus.corpus_id, + version="v1", + queries=tuple( + Query( + id=f"q{i}", + text=f"Synthetic query {i}", + tags=("unit:card",), + expectation="relevant_records", + expectation_assessor="synthetic", + judgments=( + Judgment(record_id=records[0].card_id, relevance=3, assessor="synthetic"), + ), + ) + for i in range(4) + ), + ) + live = tuple( + SourceRecord( + id=r.record_id, + sync_id=r.sync_id, + identity=r.identity, + parent=r.parent, + revision=r.revision, + payload={"synthetic": True}, + payload_schema_version=1, + capture_hash=r.capture_hash, + content_hash=None, + completeness="complete", + observed_at="2026-10-01T00:00:00Z", + source_created_at=None, + source_updated_at=None, + deleted_at=None, + removal_reason=None, + blobs=(), + indexed_revision=r.revision, + indexed_pipeline_version=r.indexed_pipeline_version, + ) + for r in records + ) + hits = tuple( + OwnedSearchHit( + record_id=r.record_id, + sync_id=r.sync_id, + revision=r.revision, + source_connection_id=UUID(int=300), + provider=r.provider, + identity=r.identity, + extraction=r.extraction, + title="Synthetic", + excerpts=("Synthetic passage",), + observed_at="2026-10-01T00:00:00Z", + source_created_at=None, + source_updated_at=None, + completeness="complete", + ) + for r in records + ) + hits[0].group = OwnedSearchGroup( + kind="session", + native_id="session", + matched_records=2, + additional_matches=( + OwnedSearchMatch.model_validate(hits[1].model_dump(exclude={"group"})), + ), + ) + wire = OwnedSearchResponse( + items=(hits[0],), + sources=(), + candidate_window_full=True, + engine_partial=False, + excluded_candidates=0, + postfilter_excluded=0, + retrieval_incomplete=True, + ).model_dump(mode="json") + return corpus, dataset, live, wire + + +def run_with_handler(tmp_path, handler): + """Use the same asynchronous HTTP boundary as the real runner.""" + corpus, dataset, _, _ = fixture() + + async def execute(): + async with httpx.AsyncClient( + base_url="http://localhost/api/v1/", + transport=httpx.MockTransport(handler), + follow_redirects=False, + ) as client: + return await replay( + client, + dataset, + corpus, + ReplayConfiguration(system="Synthetic MiniLM", mode="hybrid"), + tmp_path / "run", + ) + + return asyncio.run(execute()) + + +def page(records, *, next_cursor=None): + """Actual public wire serialization, including the live traversal contract.""" + return RecordPage( + records=records, next_cursor=next_cursor, has_more=next_cursor is not None + ).model_dump(mode="json") + + +def test_wire_replay_retains_failed_outcomes_raw_bodies_and_private_permissions(tmp_path): + _, dataset, live, wire = fixture() + calls = [] + + def handler(request): + calls.append(request) + if request.method == "GET": + assert request.url.params["state"] == "active" + assert request.url.params["limit"] == "500" + return httpx.Response( + 200, + json=page(live[:1], next_cursor="second") + if "cursor" not in request.url.params + else page(live[1:]), + ) + query = json.loads(request.content)["query"] + assert json.loads(request.content)["mode"] == "hybrid" + if query.endswith("0"): + return httpx.Response(200, json=wire) + if query.endswith("1"): + return httpx.Response(429, content=b"private provider failure") + if query.endswith("2"): + raise httpx.ReadTimeout("private timeout", request=request) + return httpx.Response(200, content=b"private malformed body") + + run = run_with_handler(tmp_path, handler) + assert [r.status for r in run.results] == ["partial", "error", "timeout", "error"] + assert run.dataset_sha256 == dataset.fingerprint() + assert len(run.results[0].record_ids) == 1 + assert len(run.timings) == 4 + assert [request.method for request in calls] == [ + "GET", + "GET", + "POST", + "POST", + "POST", + "POST", + "GET", + "GET", + ] + output = tmp_path / "run" + assert (output / "request-00001-response.json").read_bytes() == b"private provider failure" + assert not (output / "request-00002-response.json").exists() + evidence = json.loads((output / "evidence.json").read_bytes()) + assert evidence["outcome"] == "verified" + assert evidence["observations"][0]["ranking"]["fallback_reason"] == "unconfigured" + assert evidence["census_before"] and evidence["census_after"] + assert output.stat().st_mode & 0o777 == 0o700 + assert all(p.stat().st_mode & 0o777 == 0o600 for p in output.rglob("*") if p.is_file()) + + +@pytest.mark.parametrize( + "violation", + ("revision", "capture", "pipeline", "missing", "extra", "unavailable", "cursor_loop"), +) +def test_before_census_rejects_changed_or_incomplete_qualified_set(tmp_path, violation): + _, _, live, _ = fixture() + + def handler(request): + assert request.method == "GET" # A rejected census must not issue search. + records = live + if violation == "revision": + records = (live[0].model_copy(update={"revision": 4}), live[1]) + elif violation == "capture": + records = (live[0].model_copy(update={"capture_hash": "b" * 64}), live[1]) + elif violation == "pipeline": + records = (live[0].model_copy(update={"indexed_pipeline_version": 3}), live[1]) + elif violation == "missing": + records = live[:1] + elif violation == "extra": + records = (*live, live[1].model_copy(update={"id": UUID(int=3)})) + elif violation == "unavailable": + records = (live[0].model_copy(update={"content_access": "unavailable"}), live[1]) + return httpx.Response( + 200, json=page(records, next_cursor="loop" if violation == "cursor_loop" else None) + ) + + with pytest.raises(CorpusMismatch): + run_with_handler(tmp_path, handler) + assert not (tmp_path / "run/run.json").exists() + assert json.loads((tmp_path / "run/evidence.json").read_bytes())["outcome"] == "invalid" + + +def test_after_census_rejects_revision_change_and_retains_valid_raw_responses(tmp_path): + _, _, live, wire = fixture() + before = True + + def handler(request): + nonlocal before + if request.method == "POST": + before = False + return httpx.Response(200, json=wire) + records = live if before else (live[0].model_copy(update={"indexed_revision": 4}), live[1]) + return httpx.Response(200, json=page(records)) + + with pytest.raises(CorpusMismatch): + run_with_handler(tmp_path, handler) + output = tmp_path / "run" + assert len(list(output.glob("*-response.json"))) == 4 + assert not (output / "run.json").exists() + assert json.loads((output / "evidence.json").read_bytes())["census_before"] + + +@pytest.mark.parametrize("violation", ("extraction", "revision", "foreign_group")) +def test_delivered_proof_mismatch_aborts_comparison(tmp_path, violation): + _, _, live, wire = fixture() + + def handler(request): + if request.method == "GET": + return httpx.Response(200, json=page(live)) + if violation == "extraction": + wire["items"][0]["group"]["additional_matches"][0]["extraction"] = None + elif violation == "revision": + wire["items"][0]["revision"] = 4 + else: + wire["items"][0]["group"]["native_id"] = "foreign" + return httpx.Response(200, json=wire) + + with pytest.raises(CorpusMismatch): + run_with_handler(tmp_path, handler) + assert (tmp_path / "run/request-00000-response.json").exists() + outcome = json.loads((tmp_path / "run/request-00000-outcome.json").read_bytes()) + assert outcome["status"] == "error" + assert not (tmp_path / "run/run.json").exists() + + +def test_cli_auth_scope_and_redirect_rules_are_applied_to_actual_http(tmp_path, monkeypatch): + corpus, dataset, live, _ = fixture() + dataset_file, census_file = tmp_path / "dataset.json", tmp_path / "corpus.json" + dataset_file.write_text(dataset.model_dump_json()) + census_file.write_text(corpus.model_dump_json()) + constructor = httpx.AsyncClient + calls = [] + + def handler(request): + calls.append(request) + assert request.headers["X-API-Key"] == "synthetic-key" + assert request.headers["X-Organization-ID"] == str(corpus.organization_id) + assert request.url.path.startswith("/api/v1/sync/") + if request.method == "GET": + return httpx.Response(200, json=page(live)) + assert json.loads(request.content)["sync_ids"] == [str(corpus.sync_ids[0])] + return httpx.Response(302, headers={"Location": "https://foreign.invalid/private"}) + + def client_factory(**kwargs): + assert kwargs["follow_redirects"] is False + assert kwargs["trust_env"] is False + return constructor(**kwargs, transport=httpx.MockTransport(handler)) + + monkeypatch.setattr("evaluation.replay_cli.httpx.AsyncClient", client_factory) + result = asyncio.run( + execute( + ReplayArguments( + url="http://localhost/api/v1/", + dataset=dataset_file, + census=census_file, + output=tmp_path / "run", + system="synthetic", + mode="keyword", + ), + Settings(api_key="synthetic-key"), + ) + ) + assert "4 failed" in result + assert len(calls) == 6 # Two census pages and four requests, no redirect destinations. + run = json.loads((tmp_path / "run/run.json").read_bytes()) + assert all(r["status"] == "error" and not r["record_ids"] for r in run["results"]) + + +@pytest.mark.parametrize("violation", ("corpus", "unit", "unknown_label")) +def test_wrong_dataset_binding_rejects_before_network_or_output(tmp_path, violation): + corpus, dataset, _, _ = fixture() + wire = dataset.model_dump(mode="json") + if violation == "corpus": + wire["corpus_id"] = "other" + elif violation == "unit": + wire["queries"][0]["tags"] = ["unit:displayed_original"] + else: + wire["queries"][0]["judgments"][0]["record_id"] = "unknown" + dataset = Dataset.model_validate(wire) + + def handler(request): + pytest.fail("Invalid frozen binding must not make HTTP requests") + + async def execute_invalid(): + async with httpx.AsyncClient( + base_url="http://localhost/", transport=httpx.MockTransport(handler) + ) as client: + return await replay( + client, + dataset, + corpus, + ReplayConfiguration(system="synthetic", mode="keyword"), + tmp_path / "run", + ) + + with pytest.raises(CorpusMismatch): + asyncio.run(execute_invalid()) + assert not (tmp_path / "run").exists() + + +@pytest.mark.parametrize("unit", ("card", "displayed_original")) +def test_oversized_card_response_receives_no_scoring_credit(tmp_path, unit): + corpus, dataset, live, wire = fixture() + additional = wire["items"][0]["group"]["additional_matches"][0] + wire["items"][0]["group"] = None + wire["items"].append({**additional, "group": None}) + corpus = corpus.model_copy( + update={ + "records": tuple( + r.model_copy(update={"conversation": None, "parent": None}) for r in corpus.records + ) + } + ) + live = tuple(r.model_copy(update={"parent": None}) for r in live) + dataset = dataset.model_copy( + update={ + "queries": tuple( + q.model_copy( + update={ + "tags": (f"unit:{unit}",), + "judgments": ( + Judgment( + record_id=corpus.records[0].evaluation_id, + relevance=3, + assessor="synthetic", + ), + ), + } + ) + for q in dataset.queries + ) + } + ) + + def handler(request): + return httpx.Response(200, json=page(live) if request.method == "GET" else wire) + + async def execute_oversized(): + async with httpx.AsyncClient( + base_url="http://localhost/", transport=httpx.MockTransport(handler) + ) as client: + return await replay( + client, + dataset, + corpus, + ReplayConfiguration(system="synthetic", mode="keyword", limit=1, unit=unit), + tmp_path / "run", + ) + + run = asyncio.run(execute_oversized()) + assert all(r.status == "error" and not r.record_ids for r in run.results) + assert ( + len(json.loads((tmp_path / "run/request-00000-response.json").read_bytes())["items"]) == 2 + ) + + +def test_displayed_original_expansion_can_exceed_card_limit(tmp_path): + corpus, dataset, live, wire = fixture() + dataset = dataset.model_copy( + update={ + "queries": tuple( + q.model_copy( + update={ + "tags": ("unit:displayed_original",), + "judgments": ( + Judgment( + record_id=corpus.records[0].evaluation_id, + relevance=3, + assessor="synthetic", + ), + ), + } + ) + for q in dataset.queries + ) + } + ) + + def handler(request): + return httpx.Response(200, json=page(live) if request.method == "GET" else wire) + + async def execute_grouped(): + async with httpx.AsyncClient( + base_url="http://localhost/", transport=httpx.MockTransport(handler) + ) as client: + return await replay( + client, + dataset, + corpus, + ReplayConfiguration( + system="synthetic", mode="keyword", limit=1, unit="displayed_original" + ), + tmp_path / "run", + ) + + run = asyncio.run(execute_grouped()) + assert all(r.status == "partial" and len(r.record_ids) == 2 for r in run.results) + + +def test_http_dtos_and_help_import_without_server_settings_or_database(tmp_path): + root = Path(__file__).resolve().parents[4] + env = {"PATH": os.environ["PATH"], "PYTHONPATH": f"{root / 'backend'}:{root}"} + code = ( + "from evaluation.replay_cli import main; import sys; " + "assert 'airweave.core.config' not in sys.modules; " + "assert 'airweave.platform.sources' not in sys.modules; " + "assert 'sqlalchemy' not in sys.modules" + ) + imported = subprocess.run( + [sys.executable, "-c", code], env=env, cwd=tmp_path, capture_output=True, text=True + ) + assert imported.returncode == 0, imported.stderr + helped = subprocess.run( + [sys.executable, "-m", "evaluation.replay_cli", "--help"], + env=env, + cwd=tmp_path, + capture_output=True, + text=True, + ) + assert helped.returncode == 0, helped.stderr + assert "--census" in helped.stdout + + +@pytest.mark.parametrize( + "url", + ( + "http://remote.example/api/v1/", + "https://user:private@host/api/v1/", + "https://host/api/v1/?private", + ), +) +def test_cli_reuses_destination_transport_rules(url): + with pytest.raises(ValueError): + ReplayArguments( + url=url, dataset="dataset", census="census", output="run", system="test", mode="keyword" + ) + + +def test_argument_errors_do_not_echo_private_input(capsys): + with pytest.raises(SystemExit) as failure: + main(["--url", "https://private.invalid/secret"]) + assert failure.value.code == 2 + captured = capsys.readouterr() + assert "secret" not in captured.err diff --git a/evaluation/README.md b/evaluation/README.md index f5c8eef8f6..ed61fad63a 100644 --- a/evaluation/README.md +++ b/evaluation/README.md @@ -202,7 +202,7 @@ when its displayed message is not the answer. Keep these two experiments separat whole. Use a separately fingerprinted dataset tagged `unit:card`. - **Displayed-original ranking:** representative original followed by the explicitly returned additional matches, preserving card order and within-card order. Judge - those exact original identities in a separate dataset tagged `unit:original`. + those exact original identities in a separate dataset tagged `unit:displayed_original`. Do not add unreturned members merely because they belong to a returned group. This is a ranking of inspectable originals, not the number of visible cards. @@ -244,3 +244,70 @@ The delivery adapter's synthetic tests use backend fixtures: PYTHONPATH=backend:. backend/.venv/bin/python -m pytest \ backend/tests/unit/search/test_owned_evaluation.py -q ``` + +## HTTP replay of a retained corpus + +The replay command runs one mode against an **already qualified, retained** corpus. +It does not copy records, publish, fetch providers or rebuild an index. The earlier +mixed-corpus trial was cleaned; its census cannot be replayed against missing records. +Prepare and qualify a retained fixture separately before using this command. + +```sh +PYTHONPATH=backend:. backend/.venv/bin/python -m evaluation.replay_cli \ + --url http://127.0.0.1:18086/api/v1/ \ + --dataset /private/path/card-dataset.json \ + --census /private/path/qualified-corpus.json \ + --output /private/path/new-hybrid-run \ + --system 'reviewed-commit+embedding+retrieval+reranker+configuration' \ + --mode hybrid --unit card --limit 20 +``` + +Authentication uses `AIRWEAVE_API_KEY` and the census organization via +`X-Organization-ID`. The native publisher's URL rules also apply here: external +HTTPS, loopback HTTP, no URL credentials, no redirects or environment proxies. +No server database credentials or dotenv files are needed; `--help` needs no key. +Each HTTP operation has a 10-second connect and 60-second read/write/pool timeout. +Requests run once, sequentially, without retries or a whole-run deadline. + +Inputs are strict `Dataset` and `evaluation.replay.FrozenCorpus` JSON, each at most +32 MiB. The census contains `corpus_id`, `organization_id`, explicit `sync_ids` +(at most 20) and `records`. Each `FrozenRecord` extends `CorpusRecord` with the +qualified `capture_hash`, `indexed_pipeline_version`, exact `extraction` parts, +and canonical `parent` identity. Preserve source revisions/native versions in the +qualified corpus specification identified by `corpus_id`; destination revisions +are separately checked here. Freeze conversation membership before retrieval. +The census must enumerate **every active, content-available record in each selected +sync**, including intentionally excluded records; it need not describe other +sources in the provider account. Empty selected syncs are allowed. Every query +must carry exactly the matching `unit:card` or `unit:displayed_original` tag and +every judgment must refer to an identity in this corpus. + +Before and after querying, the runner paginates the existing record-list endpoint +and checks exact eligible membership, native/parent identities, destination +revisions, capture hashes and current index/pipeline revisions. Missing, extra, +withdrawn or changed records reject the comparison. Every explicitly returned +representative and additional match must also have exactly the qualified +extraction coverage. No per-record fetches are added. Availability is API +metadata/permission evidence, **not physical blob integrity**. These live census +traversals are non-atomic; changes that occur and revert between observations can +escape detection. The API does not attest the supplied model/system descriptor. + +The fresh output directory is mode 0700; files are 0600 and never overwritten. +It retains validated inputs, paginated census responses, numbered search request +and raw response files, every query outcome, mode/fallback metadata and HTTP +timings with unknown cache state. Response bytes are saved before decoding. +HTTP failures, timeouts and malformed responses remain failed query outcomes; +corpus or delivered-proof contradictions abort scoring while preserving evidence. +Only successful before/after checks produce `run.json` and an unjudged candidate +pool. Query failures still stay in the metric denominator. No labels are invented. + +Use `--report` to invoke the existing offline scorer; install its pinned +`ir-measures==0.4.3` dependency in that environment first. Alternatively, score +the retained `dataset.json` and `run.json` with the standalone command above. +Exit codes: 0 replay retained (inspect failed/partial counts); 2 invalid inputs or +output; 3 failed census HTTP verification; 4 corpus/proof mismatch; 5 unavailable +offline scoring dependency; 130 interrupted; 1 unexpected failure. + +`make test-search` runs both the delivery adapter and replay boundary tests. +These synthetic HTTP tests establish runner behavior, not live ranking quality +or qualification of any existing local corpus. diff --git a/evaluation/replay.py b/evaluation/replay.py new file mode 100644 index 0000000000..a9c8bdaf9c --- /dev/null +++ b/evaluation/replay.py @@ -0,0 +1,350 @@ +"""Replay frozen queries over HTTP against an already qualified, retained corpus. + +No capture, indexing, SQL, provider calls or relevance judgments happen here. +""" + +import hashlib +import json +import os +from pathlib import Path +from time import perf_counter +from typing import Literal +from uuid import UUID + +import httpx +from pydantic import BaseModel, Field, model_validator + +from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage +from airweave.domains.entities.canonical.models import RecordPage +from airweave.domains.entities.canonical.requests import RecordIdentity +from airweave.domains.search.owned_models import ( + OwnedRanking, + OwnedSearchRequest, + OwnedSearchResponse, +) +from airweave.domains.search.retrieval_strategy import RetrievalStrategy +from evaluation.owned_retrieval import CorpusRecord, delivered_result +from evaluation.retrieval import ( + Dataset, + Identifier, + Result, + Run, + Timing, + Value, + pool_unjudged, +) + + +class FrozenRecord(CorpusRecord): + """Replay-only proof of one qualified locator's capture and extraction revision.""" + + capture_hash: str = Field(pattern=r"^[a-f0-9]{64}$") + indexed_pipeline_version: int = Field(ge=1) + extraction: ExtractionCoverage + parent: RecordIdentity | None = None + + +class FrozenCorpus(Value): + """Complete eligible record set in each explicitly selected source sync.""" + + corpus_id: Identifier + organization_id: UUID + sync_ids: tuple[UUID, ...] = Field(min_length=1, max_length=20) + records: tuple[FrozenRecord, ...] = Field(min_length=1) + + @model_validator(mode="after") + def validate_scope(self): + """Require unique locators and stable source scopes without silently filtering.""" + if len(set(self.sync_ids)) != len(self.sync_ids): + raise ValueError("Duplicate source scope") + if len({r.record_id for r in self.records}) != len(self.records): + raise ValueError("Duplicate corpus locator") + if len({r.evaluation_id for r in self.records}) != len(self.records): + raise ValueError("Duplicate native source identity") + scopes: dict[UUID, tuple[str, str]] = {} + for record in self.records: + if record.sync_id not in self.sync_ids: + raise ValueError("Corpus record is outside selected sources") + identity = (record.source_id, record.provider) + if scopes.setdefault(record.sync_id, identity) != identity: + raise ValueError("Source scope has conflicting identity") + if len(set(scopes.values())) != len(scopes): + raise ValueError("Native source identity spans multiple destination syncs") + return self + + +class ReplayConfiguration(Value): + """Explicit request settings; system/model description is operator supplied.""" + + system: str = Field(min_length=1) + mode: RetrievalStrategy + unit: Literal["card", "displayed_original"] = "card" + limit: int = Field(default=20, ge=1, le=200) + + +class Observation(Value): + """One outcome, including failures, with raw HTTP bytes retained separately.""" + + query_id: Identifier + request_id: Identifier + status: Literal["success", "partial", "error", "timeout"] + http_status: int | None = None + ranking: OwnedRanking | None = None + candidate_window_full: bool | None = None + engine_partial: bool | None = None + retrieval_incomplete: bool | None = None + + +class ReplayEvidence(Value): + """Evidence and practical limits; invalid corpus checks never produce a scored run.""" + + outcome: Literal["verified", "invalid"] + configuration: ReplayConfiguration + corpus_sha256: str + dataset_sha256: str + census_before: bool + census_after: bool + observations: tuple[Observation, ...] + limitations: tuple[str, ...] = ( + "Before/after census traversals are live and non-atomic; transient changes can escape detection.", + "Record availability is metadata/permission evidence, not physical blob integrity.", + "Extraction is checked for delivered matches; other records retain qualification-time proof.", + "System/model description is supplied by the operator, not independently verified by the API.", + "Timings cover one HTTP request each, with unknown cache state and no retries.", + ) + + +class CorpusMismatch(ValueError): + """The supplied qualified corpus no longer describes the destination.""" + + +def save_bytes(path: Path, data: bytes) -> None: + """Never overwrite evidence or rely on a permissive caller umask.""" + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "wb") as output: + output.write(data) + + +def save_model(path: Path, model: BaseModel) -> None: + """Retain validated inputs and results privately.""" + save_bytes(path, (model.model_dump_json(indent=2) + "\n").encode()) + + +async def verify_census( + client: httpx.AsyncClient, corpus: FrozenCorpus, output: Path +) -> None: + """Compare paginated active, content-available records with the complete frozen set.""" + expected = {record.record_id: record for record in corpus.records} + eligible: set[UUID] = set() + for scope_index, sync_id in enumerate(corpus.sync_ids): + seen: set[UUID] = set() + cursors: set[str] = set() + cursor = None + previous_id: UUID | None = None + page_index = 0 + while True: + params: dict[str, str | int] = {"state": "active", "limit": 500} + if cursor is not None: + params["cursor"] = cursor + response = await client.get(f"sync/{sync_id}/records", params=params) + save_bytes( + output / f"scope-{scope_index:02d}-page-{page_index:05d}.json", + response.content, + ) + response.raise_for_status() + page = RecordPage.model_validate_json(response.content) + for live in page.records: + if ( + live.sync_id != sync_id + or live.id in seen + or (previous_id is not None and live.id <= previous_id) + ): + raise CorpusMismatch( + "Census traversal has invalid scope or ordering" + ) + seen.add(live.id) + previous_id = live.id + if live.deleted_at is not None: + raise CorpusMismatch("Active traversal returned a deleted record") + if live.content_access != "available": + continue + frozen = expected.get(live.id) + if frozen is None or ( + frozen.sync_id != live.sync_id + or frozen.identity != live.identity + or frozen.parent != live.parent + or frozen.revision != live.revision + or frozen.capture_hash != live.capture_hash + or live.indexed_revision != frozen.revision + or live.indexed_pipeline_version != frozen.indexed_pipeline_version + ): + raise CorpusMismatch("Destination differs from qualified corpus") + eligible.add(live.id) + if page.has_more != (page.next_cursor is not None): + raise CorpusMismatch("Census continuation is contradictory") + if not page.has_more: + break + if not page.records or page.next_cursor in cursors: + raise CorpusMismatch("Census continuation did not advance") + cursor = page.next_cursor + cursors.add(cursor) + page_index += 1 + if eligible != set(expected): + raise CorpusMismatch("Qualified corpus records are missing or unavailable") + + +def validate_extraction(response: OwnedSearchResponse, corpus: FrozenCorpus) -> None: + """Check the exact extraction proof of every explicitly delivered original.""" + expected = {record.record_id: record for record in corpus.records} + for hit in response.items: + additional = () if hit.group is None else hit.group.additional_matches + for match in (hit, *additional): + record = expected.get(match.record_id) + if record is None or match.extraction != record.extraction: + raise CorpusMismatch( + "Delivered extraction differs from qualified proof" + ) + + +async def replay( + client: httpx.AsyncClient, + dataset: Dataset, + corpus: FrozenCorpus, + configuration: ReplayConfiguration, + output: Path, +) -> Run: + """One sequential trial, with raw failures retained and no corpus rebuilds.""" + if dataset.corpus_id != corpus.corpus_id: + raise CorpusMismatch("Dataset and corpus identities differ") + ids = { + r.card_id if configuration.unit == "card" else r.evaluation_id + for r in corpus.records + } + for query in dataset.queries: + units = {tag for tag in query.tags if tag.startswith("unit:")} + if units != {f"unit:{configuration.unit}"} or any( + judgment.record_id not in ids for judgment in query.judgments + ): + raise CorpusMismatch( + "Dataset labels do not describe the selected delivery unit" + ) + requests = tuple( + OwnedSearchRequest( + query=q.text, + sync_ids=corpus.sync_ids, + mode=configuration.mode, + limit=configuration.limit, + ) + for q in dataset.queries + ) + output.mkdir(mode=0o700, parents=False, exist_ok=False) + save_model(output / "dataset.json", dataset) + save_model(output / "corpus.json", corpus) + before = after = False + observations: list[Observation] = [] + results: list[Result] = [] + timings: list[Timing] = [] + try: + (output / "before").mkdir(mode=0o700) + await verify_census(client, corpus, output / "before") + before = True + for index, (query, request) in enumerate( + zip(dataset.queries, requests, strict=True) + ): + request_id = f"request-{index:05d}" + save_model(output / f"{request_id}-input.json", request) + started = perf_counter() + status_code = None + parsed = None + abort: CorpusMismatch | None = None + try: + response = await client.post( + "sync/search", json=request.model_dump(mode="json") + ) + duration = (perf_counter() - started) * 1000 + status_code = response.status_code + save_bytes(output / f"{request_id}-response.json", response.content) + response.raise_for_status() + parsed = OwnedSearchResponse.model_validate_json(response.content) + if len(parsed.items) > request.limit: + raise ValueError("Delivered card count exceeds requested limit") + validate_extraction(parsed, corpus) + try: + result = delivered_result( + query.id, parsed, corpus.records, unit=configuration.unit + ) + except (ValueError, KeyError) as exc: + raise CorpusMismatch( + "Delivered identity differs from qualified corpus" + ) from exc + except httpx.TimeoutException: + duration = (perf_counter() - started) * 1000 + result = Result(query_id=query.id, status="timeout") + except httpx.HTTPError: + duration = (perf_counter() - started) * 1000 + result = Result(query_id=query.id, status="error") + except CorpusMismatch as exc: + abort = exc + result = Result(query_id=query.id, status="error") + except ValueError: + result = Result(query_id=query.id, status="error") + results.append(result) + timings.append( + Timing( + request_id=request_id, + query_id=query.id, + phase="end_to_end", + condition="unknown", + status=result.status, + duration_ms=duration, + ) + ) + observation = Observation( + query_id=query.id, + request_id=request_id, + status=result.status, + http_status=status_code, + ranking=None if parsed is None else parsed.ranking, + candidate_window_full=None + if parsed is None + else parsed.candidate_window_full, + engine_partial=None if parsed is None else parsed.engine_partial, + retrieval_incomplete=None + if parsed is None + else parsed.retrieval_incomplete, + ) + observations.append(observation) + save_model(output / f"{request_id}-outcome.json", observation) + if abort is not None: + raise abort + (output / "after").mkdir(mode=0o700) + await verify_census(client, corpus, output / "after") + after = True + finally: + save_model( + output / "evidence.json", + ReplayEvidence( + outcome="verified" if after else "invalid", + configuration=configuration, + corpus_sha256=hashlib.sha256( + corpus.model_dump_json().encode() + ).hexdigest(), + dataset_sha256=dataset.fingerprint(), + census_before=before, + census_after=after, + observations=tuple(observations), + ), + ) + run = Run( + system=configuration.system, + corpus_id=corpus.corpus_id, + dataset_sha256=dataset.fingerprint(), + results=tuple(results), + timings=tuple(timings), + ) + save_model(output / "run.json", run) + save_bytes( + output / "unjudged.json", + (json.dumps(pool_unjudged(dataset, (run,)), indent=2) + "\n").encode(), + ) + return run diff --git a/evaluation/replay_cli.py b/evaluation/replay_cli.py new file mode 100644 index 0000000000..875a164cb4 --- /dev/null +++ b/evaluation/replay_cli.py @@ -0,0 +1,156 @@ +"""Replay queries against an already-qualified retained owned-search corpus.""" + +import asyncio +import sys +from pathlib import Path + +import httpx +from pydantic import ValidationError, field_validator + +from evaluation.native_import_cli import ( + MAX_INPUT_BYTES, + Parser, + PublishArguments, + Settings, +) +from evaluation.replay import ( + CorpusMismatch, + FrozenCorpus, + ReplayConfiguration, + replay, + save_model, +) +from evaluation.retrieval import Dataset, summarize + + +class ReplayArguments(ReplayConfiguration): + """Explicit HTTP destination and frozen input files; credentials stay in env.""" + + url: str + dataset: Path + census: Path + output: Path + report: bool = False + + @field_validator("url") + @classmethod + def destination_url(cls, value: str) -> str: + """Reuse the native publisher's destination and credential transport boundary.""" + return PublishArguments.destination_url(value) + + +def parse_args(argv: list[str] | None) -> ReplayArguments: + """Require explicit corpus, request mode and system descriptor.""" + parser = Parser(description=__doc__) + parser.add_argument( + "--url", required=True, help="API root, e.g. https://host/api/v1/" + ) + parser.add_argument( + "--dataset", required=True, help="Frozen Dataset JSON (max 32 MiB)" + ) + parser.add_argument( + "--census", required=True, help="Qualified FrozenCorpus JSON (max 32 MiB)" + ) + parser.add_argument("--output", required=True, help="New private run directory") + parser.add_argument( + "--system", required=True, help="Operator-supplied system/model descriptor" + ) + parser.add_argument( + "--mode", required=True, choices=("keyword", "semantic", "hybrid") + ) + parser.add_argument( + "--unit", default="card", choices=("card", "displayed_original") + ) + parser.add_argument("--limit", type=int, default=20) + parser.add_argument( + "--report", + action="store_true", + help="Also score using the existing offline evaluator", + ) + return ReplayArguments.model_validate(vars(parser.parse_args(argv))) + + +def read_input(path: Path) -> bytes: + """Read a bounded caller-owned file without printing its private content.""" + with path.open("rb") as staged: + data = staged.read(MAX_INPUT_BYTES + 1) + if len(data) > MAX_INPUT_BYTES: + raise ValueError("Input exceeds the size limit") + return data + + +async def execute(arguments: ReplayArguments, settings: Settings) -> str: + """Use the existing API models and offline scorer, without server configuration.""" + dataset = Dataset.model_validate_json(read_input(arguments.dataset)) + corpus = FrozenCorpus.model_validate_json(read_input(arguments.census)) + configuration = ReplayConfiguration( + system=arguments.system, + mode=arguments.mode, + unit=arguments.unit, + limit=arguments.limit, + ) + async with httpx.AsyncClient( + base_url=arguments.url, + headers={ + "X-API-Key": settings.api_key.get_secret_value(), + "X-Organization-ID": str(corpus.organization_id), + }, + follow_redirects=False, + trust_env=False, + timeout=httpx.Timeout(60, connect=10), + ) as client: + run = await replay(client, dataset, corpus, configuration, arguments.output) + if arguments.report: + save_model(arguments.output / "report.json", summarize(dataset, run)) + failed = sum(result.status in {"error", "timeout"} for result in run.results) + return f"Corpus checks passed; retained {len(run.results)} query outcomes ({failed} failed)." + + +def main(argv: list[str] | None = None) -> int: + """Emit static errors; private query text, response bodies and secrets stay in files.""" + try: + arguments = parse_args(argv) + summary = asyncio.run(execute(arguments, Settings())) + except CorpusMismatch: + print( + "Corpus or delivered proof differs from frozen input; comparison rejected. Retain evidence.", + file=sys.stderr, + ) + return 4 + except httpx.HTTPError: + print( + "Census HTTP verification failed; comparison rejected. Retain evidence.", + file=sys.stderr, + ) + return 3 + except (ValidationError, ValueError, OSError): + print( + "Invalid configuration, frozen input or output. Check --help; inputs are limited to 32 MiB.", + file=sys.stderr, + ) + return 2 + except ModuleNotFoundError: + print( + "Offline scoring dependency unavailable; a verified run may already be retained. " + "Install the pinned evaluation dependencies before scoring.", + file=sys.stderr, + ) + return 5 + except KeyboardInterrupt: + print( + "Interrupted; retain private evidence. No corpus mutations were requested.", + file=sys.stderr, + ) + return 130 + except Exception: + print( + "Unexpected replay failure; retain private evidence. No comparison is established.", + file=sys.stderr, + ) + return 1 + print(summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 3fdf1c4ae0b7221e0e12aee22414c9e5f78f1a50 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 17:22:15 -0700 Subject: [PATCH 254/266] Use production HTTP app in retained evaluation server --- backend/tests/live/evaluation.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/backend/tests/live/evaluation.py b/backend/tests/live/evaluation.py index 9e62ffa863..1f7f257afc 100644 --- a/backend/tests/live/evaluation.py +++ b/backend/tests/live/evaluation.py @@ -329,21 +329,25 @@ async def project(root: Path, manifest: Manifest, sessions): async def serve(root: Path, manifest: Manifest, sessions): + """Serve the production HTTP app without running schedules or startup services.""" import uvicorn - from fastapi import FastAPI - from airweave.api.v1.api import api_router - from airweave.api.v1.endpoints.records import record_error_response from airweave.db import session - from airweave.domains.entities.canonical.store import CanonicalStoreError - configure(root, manifest) + runtime = configure(root, manifest) session.AsyncSessionLocal = sessions - app = FastAPI() - app.include_router(api_router) - app.add_exception_handler(CanonicalStoreError, record_error_response) + from airweave.main import app + + app.state.http_metrics = runtime.metrics.http await uvicorn.Server( - uvicorn.Config(app, host="127.0.0.1", port=18081, access_log=False, log_config=None) + uvicorn.Config( + app, + host="127.0.0.1", + port=18081, + access_log=False, + log_config=None, + lifespan="off", + ) ).serve() From 95318ccf96337db2c918e0e104ad951313e42db7 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 17:57:15 -0700 Subject: [PATCH 255/266] Keep indexed record read contract independent of provider setup --- backend/airweave/domains/entities/canonical/models.py | 7 +++++++ .../airweave/domains/entities/canonical/query_models.py | 8 +------- backend/tests/unit/search/test_owned_replay.py | 2 ++ 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/models.py b/backend/airweave/domains/entities/canonical/models.py index 04b9dae568..1c0bb36f68 100644 --- a/backend/airweave/domains/entities/canonical/models.py +++ b/backend/airweave/domains/entities/canonical/models.py @@ -7,6 +7,7 @@ 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 @@ -35,6 +36,12 @@ class SourceRecord(BaseModel): 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.""" diff --git a/backend/airweave/domains/entities/canonical/query_models.py b/backend/airweave/domains/entities/canonical/query_models.py index a100eaa730..df37eb203d 100644 --- a/backend/airweave/domains/entities/canonical/query_models.py +++ b/backend/airweave/domains/entities/canonical/query_models.py @@ -6,7 +6,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue -from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage +from airweave.domains.entities.canonical.models import IndexedRecordRead as IndexedRecordRead from airweave.domains.entities.canonical.models import ObservedChange, SourceRecord from airweave.domains.entities.canonical.models import RecordPage as RecordPage from airweave.domains.entities.canonical.requests import RecordIdentity @@ -15,12 +15,6 @@ from airweave.platform.sources.records.workspace_manifest import WorkspaceManifestV1 -class IndexedRecordRead(SourceRecord): - """Current original plus independently attested current extraction coverage.""" - - extraction: ExtractionCoverage | None = None - - class DocumentRead(BaseModel): """One authorized stored revision's native document and representation coverage.""" diff --git a/backend/tests/unit/search/test_owned_replay.py b/backend/tests/unit/search/test_owned_replay.py index e2248e8c7a..6b157455c8 100644 --- a/backend/tests/unit/search/test_owned_replay.py +++ b/backend/tests/unit/search/test_owned_replay.py @@ -478,6 +478,8 @@ def test_http_dtos_and_help_import_without_server_settings_or_database(tmp_path) env = {"PATH": os.environ["PATH"], "PYTHONPATH": f"{root / 'backend'}:{root}"} code = ( "from evaluation.replay_cli import main; import sys; " + "from airweave.domains.entities.canonical.models import IndexedRecordRead; " + "assert 'extraction' in IndexedRecordRead.model_json_schema()['properties']; " "assert 'airweave.core.config' not in sys.modules; " "assert 'airweave.platform.sources' not in sys.modules; " "assert 'sqlalchemy' not in sys.modules" From cdc5b943296489781163438fd61dc7c2881871f9 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 19:07:46 -0700 Subject: [PATCH 256/266] feat(search): preserve full chunks while selecting lexical snippets --- .../canonical/tests/test_owned_search.py | 41 ++++++++++++++++++- .../tests/test_vespa_client_name_substring.py | 20 +++++++++ .../search/adapters/vector_db/vespa_client.py | 5 +++ backend/airweave/domains/search/owned.py | 18 +++++++- .../airweave/domains/search/types/results.py | 3 ++ vespa/README.md | 27 ++++++++++++ vespa/app/schemas/base_entity.sd | 12 ++++-- 7 files changed, 121 insertions(+), 5 deletions(-) create mode 100644 vespa/README.md diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index bcfc86300c..87e9cf92e7 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -12,7 +12,7 @@ from airweave.api import deps from airweave.api.v1.endpoints.records import router from airweave.db.session import get_db -from airweave.domains.embedders.fakes.embedder import FakeSparseEmbedder +from airweave.domains.embedders.fakes.embedder import FakeDenseEmbedder, FakeSparseEmbedder 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 @@ -629,3 +629,42 @@ async def test_visibility_batches_exact_parts_generations_and_duplicate_chunks(d assert await visible_results( db, fence.organization_id, connection.readable_collection_id, inputs, registry ) == [valid, valid] + + +@pytest.mark.parametrize( + "mode,snippet,use_fragment", + [ + ("keyword", "… नमस्ते 原文 ", True), + ("hybrid", "… नमस्ते 原文", True), + ("semantic", "… नमस्ते 原文", False), + ("hybrid", "An unmatched lead without term markers", False), + ("keyword", None, False), + ], +) +async def test_owned_lexical_fragment_keeps_original_chunk( + database, indexed, http_search, mode, snippet, use_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 = snippet.replace("", "").replace("", "") if use_fragment 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/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 ed58eda12a..f9b5eefc3c 100644 --- a/backend/airweave/domains/search/adapters/vector_db/vespa_client.py +++ b/backend/airweave/domains/search/adapters/vector_db/vespa_client.py @@ -463,6 +463,11 @@ 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), # Native snapshots can be opened by canonical record ID without diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 882b51de1c..25290843b0 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -575,7 +575,23 @@ async def _enrich( hits[row.id].group = self._conversation(row, hits[row.id]) scores[row.id] = (1 / (60 + rank), locator) matched_text[row.id] = result.textual_representation - excerpt = result.textual_representation[:2000] + # 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("", "").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)}) diff --git a/backend/airweave/domains/search/types/results.py b/backend/airweave/domains/search/types/results.py index 8db44554cb..c1fe9459f2 100644 --- a/backend/airweave/domains/search/types/results.py +++ b/backend/airweave/domains/search/types/results.py @@ -100,6 +100,9 @@ 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") diff --git a/vespa/README.md b/vespa/README.md new file mode 100644 index 0000000000..43ce536dff --- /dev/null +++ b/vespa/README.md @@ -0,0 +1,27 @@ +# Owned lexical snippets + +The base schema adds a query-time `query_snippet` dynamic summary sourced from +`textual_representation`. It does not replace stored text, embeddings, canonical +originals or reranker input. Owned search consumes marked fragments only for +keyword/hybrid requests; missing, unmarked or semantic results retain the bounded +chunk-prefix fallback. The strings are untrusted display text, not HTML. + +Vespa's default summary includes every summary-indexed field and fields declared +in any document summary. Keeping the default selection preserves child-schema +fields and canonical identity metadata alongside the new alias: +https://docs.vespa.ai/en/querying/document-summaries.html + +The source is an index/summary field, not an attribute. Vespa documents adding a +new-name non-attribute summary alias as requiring reindexing: +https://docs.vespa.ai/en/reference/schemas/schemas.html#modifying-schemas +Inspect prepare actions before activation and follow Vespa's reindex procedure. +This does not require provider recapture or changing embedding models. No schema +has been deployed by this code change, and existing indexes safely omit the field. + +Qualification gate: validate/deploy only an isolated application first; verify +inherited email/file schemas return both full text and the dynamic alias, with +unchanged canonical metadata. Exercise a keyword near the end of a long chunk, +hybrid lexical and vector-only matches, semantic-only fallback, and multilingual +text. Confirm full reranker input and stale-publication exclusion remain intact. +Current tests simulate Vespa summary output; they do not prove live snippet +generation, highlighting defaults, or deployment compatibility. diff --git a/vespa/app/schemas/base_entity.sd b/vespa/app/schemas/base_entity.sd index 5d3b3083d6..32e452ed37 100644 --- a/vespa/app/schemas/base_entity.sd +++ b/vespa/app/schemas/base_entity.sd @@ -219,9 +219,15 @@ schema base_entity { fields: name, textual_representation, payload } - # Note: We use Vespa's default summary (no explicit document-summary) - # which automatically includes ALL fields with "summary" in their indexing, - # including fields from child schemas (e.g., url from file_entity) + # The default summary includes this alias as well as all stored summary fields, + # including child-schema fields. Keep full text unchanged for reranking/readers. + # Adding this non-attribute summary alias requires Vespa reindexing; see README. + document-summary owned_snippets { + summary query_snippet type string { + source: textual_representation + dynamic + } + } # ========================================================================= # BASE RANK PROFILE - shared inputs and functions From bd2739dace5b0aef3c070b71704d88c0f36c6e8a Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 19:18:12 -0700 Subject: [PATCH 257/266] test(search): qualify lexical snippets against real Vespa in CI --- backend/tests/integration/test_owned_vespa.py | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/backend/tests/integration/test_owned_vespa.py b/backend/tests/integration/test_owned_vespa.py index c2491e6960..4097054ec8 100644 --- a/backend/tests/integration/test_owned_vespa.py +++ b/backend/tests/integration/test_owned_vespa.py @@ -362,3 +362,83 @@ async def test_real_minilm_paraphrase_retrieval_and_exact_keyword(): 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 From c0632c202cab4fa4ecbabc510201647647439438 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 19:37:46 -0700 Subject: [PATCH 258/266] Render Vespa fragment separators as plain text --- .../canonical/tests/test_owned_search.py | 20 +++++++++++-------- .../tests/test_real_vespa_projection.py | 8 +++++++- backend/airweave/domains/search/owned.py | 5 ++++- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py index 87e9cf92e7..f8e18acf56 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_search.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_search.py @@ -632,17 +632,21 @@ async def test_visibility_batches_exact_parts_generations_and_duplicate_chunks(d @pytest.mark.parametrize( - "mode,snippet,use_fragment", + "mode,snippet,expected_fragment", [ - ("keyword", "… नमस्ते 原文 ", True), - ("hybrid", "… नमस्ते 原文", True), - ("semantic", "… नमस्ते 原文", False), - ("hybrid", "An unmatched lead without term markers", False), - ("keyword", None, False), + ( + "keyword", + "नमस्ते 原文 ", + "… नमस्ते 原文 …", + ), + ("hybrid", "… नमस्ते 原文", "… नमस्ते 原文"), + ("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, use_fragment + database, indexed, http_search, mode, snippet, expected_fragment ): fence, locator, _ = indexed client, vector, _, _, _ = http_search @@ -664,7 +668,7 @@ async def test_owned_lexical_fragment_keeps_original_chunk( }, ) assert response.status_code == 200, response.text - expected = snippet.replace("", "").replace("", "") if use_fragment else original[:2000] + 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_real_vespa_projection.py b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py index 4cb42c317f..9d17b5b346 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_real_vespa_projection.py @@ -164,7 +164,13 @@ async def session(): 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] - assert "fundraising meeting" in " ".join(result.json()["items"][0]["excerpts"]) + 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"}} diff --git a/backend/airweave/domains/search/owned.py b/backend/airweave/domains/search/owned.py index 25290843b0..64944ae396 100644 --- a/backend/airweave/domains/search/owned.py +++ b/backend/airweave/domains/search/owned.py @@ -587,7 +587,10 @@ async def _enrich( # Remove only Vespa's presentation delimiters. All other markup is # ordinary untrusted text, never parsed/rendered as HTML here. excerpt = ( - snippet.replace("", "").replace("", "").strip() + snippet.replace("", "") + .replace("", "") + .replace("", " … ") + .strip() or result.textual_representation )[:2000] else: From 4c938cce0433005e512d84bdc087c926139013f3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 19:54:14 -0700 Subject: [PATCH 259/266] Document real Vespa snippet qualification and rollout limits --- vespa/README.md | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vespa/README.md b/vespa/README.md index 43ce536dff..46850b2d19 100644 --- a/vespa/README.md +++ b/vespa/README.md @@ -4,7 +4,8 @@ The base schema adds a query-time `query_snippet` dynamic summary sourced from `textual_representation`. It does not replace stored text, embeddings, canonical originals or reranker input. Owned search consumes marked fragments only for keyword/hybrid requests; missing, unmarked or semantic results retain the bounded -chunk-prefix fallback. The strings are untrusted display text, not HTML. +chunk-prefix fallback. Highlight delimiters are removed and Vespa fragment +separators become plain ellipses. The strings are untrusted display text, not HTML. Vespa's default summary includes every summary-indexed field and fields declared in any document summary. Keeping the default selection preserves child-schema @@ -23,5 +24,19 @@ inherited email/file schemas return both full text and the dynamic alias, with unchanged canonical metadata. Exercise a keyword near the end of a long chunk, hybrid lexical and vector-only matches, semantic-only fallback, and multilingual text. Confirm full reranker input and stale-publication exclusion remain intact. -Current tests simulate Vespa summary output; they do not prove live snippet -generation, highlighting defaults, or deployment compatibility. +Qualification evidence (2026-10-02 UTC): the disposable CI application at +[c0632c2, run 36956517573](https://github.com/AlmanacCode/airweave/actions/runs/36956517573) +passed all six real-engine tests in 74.14 seconds. The new file/email case proves +default-summary alias availability, unchanged full chunks/identity/child URL, +near-end lexical highlighting, and adjacent Hindi/Chinese text. The canonical +projection test also verifies published search, semantic fallback, stale revision +exclusion, source access withdrawal, and deletion. These are synthetic correctness +checks, not corpus relevance measurements or an existing-index rollout. + +The preceding run at bd2739d passed the new summary case but failed an older +assertion requiring a particular body sentence for a filename-matching query. +Actual Vespa selected the matching filename and returned `` delimiters. +The follow-up normalizes those delimiters and checks lexical query evidence while +retaining the full-body assertion for semantic fallback. No retained local index +was deployed, reindexed, or otherwise changed. Existing-index prepare/reindex +compatibility remains a rollout gate. From 807505b6539c4baf8124a555b81c0a843636e7ce Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 20:06:56 -0700 Subject: [PATCH 260/266] Share provider-scoped policy for intentional search exclusions --- .../entities/canonical/projection_mappers.py | 16 +- .../entities/canonical/projection_policy.py | 18 +++ .../entities/canonical/projection_store.py | 13 +- .../domains/entities/canonical/projector.py | 6 +- .../tests/test_projection_exclusions.py | 150 ++++++++++++++++++ .../native_ingestion/tests/test_projection.py | 3 +- 6 files changed, 177 insertions(+), 29 deletions(-) create mode 100644 backend/airweave/domains/entities/canonical/projection_policy.py create mode 100644 backend/airweave/domains/entities/canonical/tests/test_projection_exclusions.py diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index fedc443d4d..508e3ea763 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -12,10 +12,10 @@ from pydantic import JsonValue, TypeAdapter -from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event 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 @@ -335,8 +335,6 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: - if record.identity.record_type in {"meeting_listing", "scratchpad_listing"}: - return () if record.identity.record_type not in {"meeting", "scratchpad_note"}: raise ProjectionMappingError("Unsupported Wispr record type") responses = record.payload.get("responses") @@ -372,18 +370,6 @@ def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: ) -def excluded_from_search(record: SourceRecord, source_name: str) -> bool: - """Retained calendar exclusions intentionally publish no searchable meeting.""" - if source_name == "almanac": - from airweave.domains.native_ingestion.projection import excluded_native - - return excluded_native(record) - return source_name == "google_calendar" and ( - record.identity.record_type == "event_occurrence" - or (record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload)) - ) - - @asynccontextmanager async def map_record( # noqa: C901 -- explicit provider dispatch keeps mapper ownership visible record: SourceRecord, diff --git a/backend/airweave/domains/entities/canonical/projection_policy.py b/backend/airweave/domains/entities/canonical/projection_policy.py new file mode 100644 index 0000000000..20d0a0ed77 --- /dev/null +++ b/backend/airweave/domains/entities/canonical/projection_policy.py @@ -0,0 +1,18 @@ +"""Pure, provider-scoped publication policy shared by mapping and store validation.""" + +from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event +from airweave.domains.entities.canonical.models import SourceRecord + + +def excluded_from_search(record: SourceRecord, source_name: str) -> bool: + """Records retained for authority/navigation that intentionally have no search text.""" + if source_name == "wispr": + return record.identity.record_type in {"meeting_listing", "scratchpad_listing"} + if source_name == "almanac": + from airweave.domains.native_ingestion.projection import excluded_native + + return excluded_native(record) + return source_name == "google_calendar" and ( + record.identity.record_type == "event_occurrence" + or (record.identity.record_type == "event" and is_cancelled_recurring_event(record.payload)) + ) diff --git a/backend/airweave/domains/entities/canonical/projection_store.py b/backend/airweave/domains/entities/canonical/projection_store.py index 3a7ed9350d..1c033aefe6 100644 --- a/backend/airweave/domains/entities/canonical/projection_store.py +++ b/backend/airweave/domains/entities/canonical/projection_store.py @@ -9,7 +9,6 @@ from sqlalchemy.ext.asyncio import AsyncSession from airweave.db.unit_of_work import UnitOfWork -from airweave.domains.entities.canonical.calendar import is_cancelled_recurring_event from airweave.domains.entities.canonical.extraction_models import ExtractionCoverage from airweave.domains.entities.canonical.projection_models import ( ProjectionBinding, @@ -20,6 +19,7 @@ ProjectionWork, projection_document_locator, ) +from airweave.domains.entities.canonical.projection_policy import excluded_from_search from airweave.domains.entities.canonical.store import content_is_available, source_record from airweave.domains.entities.canonical.text_artifacts import text_manifest from airweave.domains.entities.canonical.text_models import TextArtifact @@ -358,18 +358,13 @@ async def publish( chunk_count: int, ) -> bool: """Publish only a complete current generation; late writers cannot replace it.""" - exclusion = work.record.identity.record_type == "event" and is_cancelled_recurring_event( - work.record.payload - ) - no_content = ( - work.record.deleted_at is not None - or exclusion - or work.record.identity.record_type == "event_occurrence" + no_content = work.record.deleted_at is not None or excluded_from_search( + work.record, work.binding.source_name ) if chunk_count < 0: raise ValueError("Active records require a nonempty complete projection") if no_content and chunk_count != 0: - raise ValueError("Deleted records and recurrence exclusions must publish no content") + raise ValueError("Deleted records and search exclusions must publish no content") async with UnitOfWork(db): entity = await self._current(db, work) if entity is None: diff --git a/backend/airweave/domains/entities/canonical/projector.py b/backend/airweave/domains/entities/canonical/projector.py index 2e7c577a63..1c7e432c98 100644 --- a/backend/airweave/domains/entities/canonical/projector.py +++ b/backend/airweave/domains/entities/canonical/projector.py @@ -99,10 +99,8 @@ async def project_one( logger: ContextualLogger, ) -> bool: """Never feed native capture JSON; mapper emits explicit safe projection entities.""" - from airweave.domains.entities.canonical.projection_mappers import ( - excluded_from_search, - map_record, - ) + from airweave.domains.entities.canonical.projection_mappers import map_record + from airweave.domains.entities.canonical.projection_policy import excluded_from_search if not await self._admit(work, source_name, destination): return False 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/native_ingestion/tests/test_projection.py b/backend/airweave/domains/native_ingestion/tests/test_projection.py index fcdab49cf3..36e6849547 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_projection.py +++ b/backend/airweave/domains/native_ingestion/tests/test_projection.py @@ -5,7 +5,8 @@ import pytest from airweave.adapters.storage.filesystem import FilesystemBackend -from airweave.domains.entities.canonical.projection_mappers import excluded_from_search, map_record +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 94fb371fde3e3d0f43dc7cfb0860d5dce81794ce Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 20:11:15 -0700 Subject: [PATCH 261/266] Attest complete selected transcripts without reconciling other sessions --- .../domains/native_ingestion/import_models.py | 15 +- .../domains/native_ingestion/import_store.py | 5 +- .../domains/native_ingestion/scope_store.py | 4 +- .../native_ingestion/tests/test_imports.py | 61 ++++++ .../tests/test_transcript_reconciliation.py | 203 ++++++++++++++++++ 5 files changed, 282 insertions(+), 6 deletions(-) create mode 100644 backend/airweave/domains/native_ingestion/tests/test_transcript_reconciliation.py diff --git a/backend/airweave/domains/native_ingestion/import_models.py b/backend/airweave/domains/native_ingestion/import_models.py index 13216a13ac..22575a7650 100644 --- a/backend/airweave/domains/native_ingestion/import_models.py +++ b/backend/airweave/domains/native_ingestion/import_models.py @@ -6,15 +6,28 @@ 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 imports never authorize absence deletion.""" + """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): diff --git a/backend/airweave/domains/native_ingestion/import_store.py b/backend/airweave/domains/native_ingestion/import_store.py index fb50cdfcd0..54c656248c 100644 --- a/backend/airweave/domains/native_ingestion/import_store.py +++ b/backend/airweave/domains/native_ingestion/import_store.py @@ -102,6 +102,8 @@ async def start( ) -> 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: @@ -140,11 +142,10 @@ async def start( if bound.binding.dataset == "knowledge" else {"session": (None,), "message": ("session",)} ) - policy = "exhaustive" if request.coverage == "complete" else "discovery_only" configuration = CycleConfiguration( fingerprint=fingerprint, parents=parents, - completion_policies={kind: policy for kind in 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)) diff --git a/backend/airweave/domains/native_ingestion/scope_store.py b/backend/airweave/domains/native_ingestion/scope_store.py index 401ed0b12f..95b713b447 100644 --- a/backend/airweave/domains/native_ingestion/scope_store.py +++ b/backend/airweave/domains/native_ingestion/scope_store.py @@ -92,9 +92,7 @@ async def read( raise NativeAdmissionError("Scope is unavailable or superseded for this import") state = (await self.scans._state(db, row)).model_copy( update={ - "completion_policy": "exhaustive" - if imported.request.coverage == "complete" - else "discovery_only" + "completion_policy": imported.request.completion_policy(request.scope.record_type) } ) return scope_state(state) diff --git a/backend/airweave/domains/native_ingestion/tests/test_imports.py b/backend/airweave/domains/native_ingestion/tests/test_imports.py index 14d09f7cba..d7d73c7b7e 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_imports.py +++ b/backend/airweave/domains/native_ingestion/tests/test_imports.py @@ -1,6 +1,8 @@ """Real PostgreSQL import request recovery, concurrent creation and rollback.""" import asyncio +import hashlib +import json from uuid import uuid4 import pytest @@ -17,6 +19,7 @@ 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 @@ -141,6 +144,64 @@ async def fail(*args, **kwargs): 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 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 From 2bf2c55bb5e17d2361439fd0e04843c363f40f5d Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 20:26:56 -0700 Subject: [PATCH 262/266] Use valid native originals in grouping and visibility fixtures --- .../canonical/tests/test_owned_grouping.py | 33 +++++++++++++------ .../native_ingestion/tests/test_search.py | 5 +-- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py b/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py index 65b1d5f562..3c0aca4c50 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py +++ b/backend/airweave/domains/entities/canonical/tests/test_owned_grouping.py @@ -17,6 +17,13 @@ indexed, ) from airweave.domains.entities.canonical.tests.test_search_visibility import hit +from airweave.domains.native_ingestion.tests.test_ingestion import ingest +from airweave.domains.native_ingestion.tests.test_projection import ( + message as native_message, +) +from airweave.domains.native_ingestion.tests.test_projection import ( + session as native_session, +) from airweave.domains.search.types import SearchResults from airweave.models.collection import Collection from airweave.models.source_connection import SourceConnection @@ -41,27 +48,33 @@ async def test_group_after_final_gate_promotes_survivor_and_limits_cards( for i in range(4) ] if withdraw == "parent": - parent = RecordIdentity(record_type="session", native_id="session") - await capture(database, capture_service, fence, observation(identity=parent)) - messages = [ - message.model_copy( + session = native_session() + parent = session.identity + original = native_message("Original message") + native_messages = [ + original.model_copy( update={ - "identity": message.identity.model_copy( - update={"container_id": parent.native_id} + "identity": RecordIdentity( + record_type="message", native_id=f"m{i}", container_id=parent.native_id ), - "parent": parent, + "original": original.original | {"id": f"m{i}", "ordinal": i}, } ) - for message in messages + for i in range(4) ] async with database() as db: await db.execute( update(SourceConnection) .where(SourceConnection.id == connection.id) - .values(short_name="almanac") + .values( + short_name="almanac", + config_fields={"owner_id": "owner-one", "dataset": "sessions"}, + ) ) await db.commit() - await capture(database, capture_service, fence, *messages) + await ingest(database, fence, session, *native_messages) + else: + await capture(database, capture_service, fence, *messages) store = CanonicalProjectionStore() async with database() as db: collection_id = await db.scalar( diff --git a/backend/airweave/domains/native_ingestion/tests/test_search.py b/backend/airweave/domains/native_ingestion/tests/test_search.py index cbbd7d28f4..9e4ad00b8e 100644 --- a/backend/airweave/domains/native_ingestion/tests/test_search.py +++ b/backend/airweave/domains/native_ingestion/tests/test_search.py @@ -12,7 +12,8 @@ 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, snapshot +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 @@ -25,7 +26,7 @@ async def test_native_search_requires_current_publication_and_source_access(data """Index transport is synthetic; capture, scope checks and publication are real SQL.""" _, fence = source await bind(database, fence) - await ingest(database, fence, snapshot()) + await ingest(database, fence, knowledge()) async with database() as db: collection = await db.scalar( select(Collection) From 48fbd5d9dc81f6375794a8632a5b978f0e2c60aa Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 20:43:09 -0700 Subject: [PATCH 263/266] test: remove superseded private Wispr listing mapper check --- .../tests/unit/platform/sources/test_wispr_capture.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/backend/tests/unit/platform/sources/test_wispr_capture.py b/backend/tests/unit/platform/sources/test_wispr_capture.py index 399aef0154..1ed0d25a50 100644 --- a/backend/tests/unit/platform/sources/test_wispr_capture.py +++ b/backend/tests/unit/platform/sources/test_wispr_capture.py @@ -141,15 +141,6 @@ async def test_listing_preserves_native_row_and_continuation(monkeypatch): ) -@pytest.mark.asyncio -async def test_listing_projection_is_explicitly_empty(): - from types import SimpleNamespace - - from airweave.domains.entities.canonical.projection_mappers import _wispr - - assert _wispr(SimpleNamespace(identity=SimpleNamespace(record_type="meeting_listing"))) == () - - @pytest.mark.asyncio async def test_partitioned_listing_resumes_both_half_open_windows(monkeypatch): from airweave.domains.entities.canonical.page_source import CanonicalPageSource From d1e3685e0c4fdc3d062ec58cc0ff3b9d51e542c0 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 20:58:09 -0700 Subject: [PATCH 264/266] fix: preserve Wispr transcript range text boundaries --- .../entities/canonical/projection_mappers.py | 43 ++++++- .../entities/test_projection_mappers.py | 116 +++++++++++++++++- 2 files changed, 148 insertions(+), 11 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/projection_mappers.py b/backend/airweave/domains/entities/canonical/projection_mappers.py index 508e3ea763..ce774293ae 100644 --- a/backend/airweave/domains/entities/canonical/projection_mappers.py +++ b/backend/airweave/domains/entities/canonical/projection_mappers.py @@ -292,11 +292,22 @@ def _wispr_range_text(response: dict[str, JsonValue], field: str, offset: int) - return native -def _wispr_text(responses: list[JsonValue], field: str) -> str: - """Follow exactly the requested ranges; strip server continuation guidance.""" +_WISPR_TRANSCRIPT_HEADER = ( + "<<>>\n" +) +_WISPR_TRANSCRIPT_FOOTER = "\n<<>>" + + +def _wispr_text( # noqa: C901 -- keep native framing, order and offset-unit guards together + responses: list[JsonValue], field: str +) -> str: + """Reassemble retained ranges without changing native whitespace or word boundaries.""" expected = 0 complete = False pieces: list[str] = [] + wrapped: bool | None = None + offset_units = {"codepoints", "utf16"} for page in responses: if not isinstance(page, dict): raise ProjectionMappingError("Invalid Wispr response range") @@ -310,10 +321,19 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: if not isinstance(window, dict) or window.get("start_char") != expected or complete: raise ProjectionMappingError("Wispr text ranges overlap or have a gap") text = _wispr_range_text(response, field, expected) + if field == "transcript": + page_wrapped = text.startswith(_WISPR_TRANSCRIPT_HEADER) + if page_wrapped != text.endswith(_WISPR_TRANSCRIPT_FOOTER): + raise ProjectionMappingError("Wispr transcript framing is incomplete") + if wrapped is not None and page_wrapped != wrapped: + raise ProjectionMappingError("Wispr transcript framing changed between ranges") + wrapped = page_wrapped + if page_wrapped: + text = text[len(_WISPR_TRANSCRIPT_HEADER) : -len(_WISPR_TRANSCRIPT_FOOTER)] markers = list( re.finditer( rf"(?m)^\(\.\.\.truncated, \d+ chars remaining; continue with " - rf"view_{field}\.start_char=(\d+)\.\.\.\)\s*$", + rf"view_{field}\.start_char=(\d+)\.\.\.\)$", text, ) ) @@ -324,14 +344,27 @@ def _wispr_text(responses: list[JsonValue], field: str) -> str: following = int(marker.group(1)) if following <= expected: raise ProjectionMappingError("Wispr continuation does not advance") - pieces.append(text[: marker.start()].rstrip()) + prefix = text[: marker.start()] + if marker.end() != len(text) or not prefix.endswith("\n\n"): + raise ProjectionMappingError("Wispr continuation framing is invalid") + fragment = prefix[:-2] # Remove only the observed marker separator. + length = following - expected + candidates: set[str] = set() + if len(fragment) == length: + candidates.add("codepoints") + if len(fragment.encode("utf-16-le")) // 2 == length: + candidates.add("utf16") + offset_units &= candidates + if not offset_units: + raise ProjectionMappingError("Wispr range length disagrees with continuation") + pieces.append(fragment) expected = following else: pieces.append(text) complete = True if not complete: raise ProjectionMappingError("Wispr text has an unfinished continuation") - return "\n".join(pieces) + return "".join(pieces) def _wispr(record: SourceRecord) -> tuple[BaseEntity, ...]: diff --git a/backend/tests/unit/domains/entities/test_projection_mappers.py b/backend/tests/unit/domains/entities/test_projection_mappers.py index 6097a4d308..0177ea4d77 100644 --- a/backend/tests/unit/domains/entities/test_projection_mappers.py +++ b/backend/tests/unit/domains/entities/test_projection_mappers.py @@ -91,8 +91,8 @@ async def test_wispr_ranges_do_not_duplicate_default_content(): "summary": "summary", "content": "notes", "transcript": ( - "abc\n(...truncated, 3 chars remaining; " - "continue with view_transcript.start_char=3...)\n\nProvider guidance." + "abc\n\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=3...)" ), } second = {"id": "meeting", "content": "notes", "transcript": "def"} @@ -111,7 +111,7 @@ async def test_wispr_ranges_do_not_duplicate_default_content(): async with map_record(record("meeting", payload), "wispr", AsyncMock()) as entities: entities = entities.entities assert entities[0].notes == "notes" - assert entities[0].transcript == "abc\ndef" + assert entities[0].transcript == "abcdef" assert "guidance" not in entities[0].transcript @@ -127,7 +127,7 @@ async def test_wispr_unfinished_ranges_fail_instead_of_indexing_truncated_body() "response": { "content": "", "transcript": ( - "(...truncated, 3 chars remaining; " + "abc\n\n(...truncated, 3 chars remaining; " "continue with view_transcript.start_char=3...)" ), }, @@ -239,7 +239,7 @@ async def test_wispr_scratchpad_retains_all_text_ranges_without_meeting_fields() "id": "note", "title": "Ideas", "content": ( - "abc\n(...truncated, 3 chars remaining; " + "abc\n\n(...truncated, 3 chars remaining; " "continue with view_content.start_char=3...)" ), "modified_at": "2026-09-30T00:00:00Z", @@ -262,7 +262,7 @@ async def test_wispr_scratchpad_retains_all_text_ranges_without_meeting_fields() async with map_record(original, "wispr", storage) as entities: entities = entities.entities assert entities[0].note_id == "note" - assert entities[0].content == "abc\ndef" + assert entities[0].content == "abcdef" assert entities[0].web_url == "" assert original.model_dump() == before assert storage.mock_calls == [] @@ -313,3 +313,107 @@ async def test_wispr_explicitly_absent_transcript_keeps_available_notes(): with pytest.raises(ProjectionMappingError, match="unavailable or malformed"): async with map_record(record("meeting", payload), "wispr", AsyncMock()): pass + + +@pytest.mark.parametrize("unit", ["codepoints", "utf16"]) +def test_wispr_ranges_preserve_unicode_and_source_whitespace(unit): + from airweave.domains.entities.canonical.projection_mappers import _wispr_text + + first = "स्वीकृत 😀 \n" + offset = len(first) if unit == "codepoints" else len(first.encode("utf-16-le")) // 2 + pages = [ + { + "requested_ranges": {"view_content": {"start_char": 0}}, + "response": { + "content": first + + "\n\n(...truncated, 3 chars remaining; " + + f"continue with view_content.start_char={offset}...)" + }, + }, + { + "requested_ranges": {"view_content": {"start_char": offset}}, + "response": {"content": "done \n\n"}, + }, + ] + assert _wispr_text(pages, "content") == first + "done \n\n" + + +def test_wispr_wrapped_transcript_preserves_midword_and_terminal_whitespace(): + import copy + + from airweave.domains.entities.canonical.projection_mappers import ( + _WISPR_TRANSCRIPT_FOOTER, + _WISPR_TRANSCRIPT_HEADER, + _wispr_text, + ) + + pages = [ + { + "requested_ranges": {"view_transcript": {"start_char": 0}}, + "response": { + "transcript": _WISPR_TRANSCRIPT_HEADER + + "appro\n\n(...truncated, 3 chars remaining; " + "continue with view_transcript.start_char=5...)" + _WISPR_TRANSCRIPT_FOOTER + }, + }, + { + "requested_ranges": {"view_transcript": {"start_char": 5}}, + "response": { + "transcript": _WISPR_TRANSCRIPT_HEADER + "ved \n" + _WISPR_TRANSCRIPT_FOOTER + }, + }, + ] + before = copy.deepcopy(pages) + assert _wispr_text(pages, "transcript") == "approved \n" + assert pages == before + + +@pytest.mark.parametrize("defect", ["separator", "footer", "length", "gap", "unit_change"]) +def test_wispr_inconsistent_range_framing_fails(defect): + from airweave.domains.entities.canonical.projection_mappers import _wispr_text + + pages = [ + { + "requested_ranges": {"view_content": {"start_char": 0}}, + "response": { + "content": "😀a\n\n(...truncated, 2 chars remaining; " + "continue with view_content.start_char=2...)" + }, + }, + {"requested_ranges": {"view_content": {"start_char": 2}}, "response": {"content": "bc"}}, + ] + if defect == "separator": + pages[0]["response"]["content"] = pages[0]["response"]["content"].replace("\n\n", "\n") + elif defect == "footer": + pages[0]["response"]["content"] += "\nUnknown guidance" + elif defect == "length": + pages[0]["response"]["content"] = pages[0]["response"]["content"].replace( + "start_char=2", "start_char=7" + ) + elif defect == "gap": + pages[1]["requested_ranges"]["view_content"]["start_char"] = 3 + else: + # First range requires codepoints, the next would require UTF16 units. + pages[1]["response"]["content"] = ( + "😀a\n\n(...truncated, 1 chars remaining; continue with view_content.start_char=5...)" + ) + with pytest.raises(ProjectionMappingError): + _wispr_text(pages, "content") + + +def test_wispr_terminal_transcript_rejects_incomplete_wrapper(): + from airweave.domains.entities.canonical.projection_mappers import ( + _WISPR_TRANSCRIPT_HEADER, + _wispr_text, + ) + + with pytest.raises(ProjectionMappingError, match="framing"): + _wispr_text( + [ + { + "requested_ranges": {"view_transcript": {"start_char": 0}}, + "response": {"transcript": _WISPR_TRANSCRIPT_HEADER + "intact text"}, + } + ], + "transcript", + ) From d54209ec8a1f211e38b7b8f50c43e354c45df975 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 21:49:24 -0700 Subject: [PATCH 265/266] fix: preserve plain Gmail content when HTML alternative is empty --- .../entities/canonical/gmail_projection.py | 44 ++++++--- .../canonical/tests/test_gmail_projection.py | 91 +++++++++++++++++++ 2 files changed, 124 insertions(+), 11 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/gmail_projection.py b/backend/airweave/domains/entities/canonical/gmail_projection.py index c1e72f9d2e..1e91097e08 100644 --- a/backend/airweave/domains/entities/canonical/gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/gmail_projection.py @@ -34,13 +34,6 @@ def _attachment(part: dict) -> bool: ) -def _has_html(part: dict) -> bool: - """Find an HTML alternative even when wrapped by multipart/related.""" - return part.get("mimeType") == "text/html" or any( - _has_html(child) for child in part.get("parts", []) - ) - - async def _body( part: dict, path: str, @@ -80,6 +73,32 @@ def __init__(self, record: SourceRecord, storage: StorageBackend, directory: Pat self.directory = directory self.attachments: list[ProjectionInput] = [] self.charset_recoveries: list[CharsetRecovery] = [] + self.decoded_text: dict[str, DecodedText] = {} + + async def _decoded_text_for(self, part: dict, path: str) -> DecodedText: + """Decode one MIME body once, whether selection or rendering reaches it first.""" + decoded = self.decoded_text.get(path) + if decoded is None: + content = await _body(part, path, self.record, self.storage) + decoded = _decode_text(content, part, path) + self.decoded_text[path] = decoded + return decoded + + async def _has_nonempty_html(self, part: dict, path: str) -> bool: + """Check decoded HTML content, including HTML nested in related parts.""" + mime = part.get("mimeType", "").lower() + if mime.startswith("multipart/"): + children = list(enumerate(part.get("parts", []))) + if mime == "multipart/alternative": + children.reverse() + for index, child in children: + if await self._has_nonempty_html(child, f"{path}/parts/{index}"): + return True + return False + if mime != "text/html" or _attachment(part): + return False + decoded = await self._decoded_text_for(part, path) + return bool(decoded.text.strip()) async def attachment(self, part: dict, path: str) -> None: """Project retained attachment bytes; keep missing bytes explicit on the original.""" @@ -130,18 +149,21 @@ async def render(self, part: dict, path: str) -> str: if mime.startswith("multipart/"): children = list(enumerate(part.get("parts", []))) if mime == "multipart/alternative": - html_parts = [(i, p) for i, p in children if _has_html(p)] + html_part = None + for index, child in reversed(children): + if await self._has_nonempty_html(child, f"{path}/parts/{index}"): + html_part = (index, child) + break plain_parts = [(i, p) for i, p in children if p.get("mimeType") == "text/plain"] # Prefer one complete alternative, not duplicate plain and HTML text. - children = html_parts[-1:] or plain_parts[-1:] or children[-1:] + children = [html_part] if html_part else plain_parts[-1:] or children[-1:] return "\n".join( [await self.render(child, f"{path}/parts/{index}") for index, child in children] ) if _attachment(part): await self.attachment(part, path) return "" - content = await _body(part, path, self.record, self.storage) - decoded = _decode_text(content, part, path) + decoded = await self._decoded_text_for(part, path) if decoded.recovery is not None: self.charset_recoveries.append(decoded.recovery) text = decoded.text diff --git a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py index 80a950964c..6ba89e4a0a 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py +++ b/backend/airweave/domains/entities/canonical/tests/test_gmail_projection.py @@ -119,6 +119,97 @@ async def test_alternative_prefers_full_html_not_plain_duplicate_or_snippet(tmp_ 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( From b4dcc134cc7e6394d4f7355a823e8c80c8d774b3 Mon Sep 17 00:00:00 2001 From: rohans0509 Date: Thu, 1 Oct 2026 22:34:13 -0700 Subject: [PATCH 266/266] test: cover change-feed continuation after access withdrawal --- .../entities/canonical/tests/test_http.py | 105 ++++++++++++++++-- 1 file changed, 94 insertions(+), 11 deletions(-) diff --git a/backend/airweave/domains/entities/canonical/tests/test_http.py b/backend/airweave/domains/entities/canonical/tests/test_http.py index 1eda304049..3d767f3f67 100644 --- a/backend/airweave/domains/entities/canonical/tests/test_http.py +++ b/backend/airweave/domains/entities/canonical/tests/test_http.py @@ -11,24 +11,16 @@ 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 CaptureBatch +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 -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"))) - ) +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) - owner = fence.organization_id - - async def context(): - return SimpleNamespace(organization=SimpleNamespace(id=owner)) async def session(): async with database() as db: @@ -39,6 +31,21 @@ async def 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: @@ -64,3 +71,79 @@ async def session(): 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"