From 3997fe905ce2c011f4c33a115f5326bfdfd4141f Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Tue, 26 May 2026 09:49:39 -0700 Subject: [PATCH] fix: audit log retention policy and cleanup --- .../7c8d9e0f1a2b_authz_foundations.py | 16 +- .../services/authorization/service.py | 18 ++ .../langflow/services/authorization/utils.py | 259 +++++++++++++++--- src/backend/base/langflow/services/utils.py | 44 +++ .../authorization/test_audit_retention.py | 151 ++++++++++ .../unit/services/authorization/test_utils.py | 153 +++++------ src/lfx/src/lfx/services/settings/auth.py | 11 + 7 files changed, 533 insertions(+), 119 deletions(-) create mode 100644 src/backend/tests/unit/services/authorization/test_audit_retention.py diff --git a/src/backend/base/langflow/alembic/versions/7c8d9e0f1a2b_authz_foundations.py b/src/backend/base/langflow/alembic/versions/7c8d9e0f1a2b_authz_foundations.py index 1ddd6d4a47..897a4ee77c 100644 --- a/src/backend/base/langflow/alembic/versions/7c8d9e0f1a2b_authz_foundations.py +++ b/src/backend/base/langflow/alembic/versions/7c8d9e0f1a2b_authz_foundations.py @@ -1,14 +1,25 @@ """Authz foundations: tables, workspace_id columns, and built-in system roles. +Phase: EXPAND Revision ID: 7c8d9e0f1a2b Revises: mb01b2c3d4e5 Create Date: 2026-05-20 + +This migration is purely additive (new tables + nullable ``workspace_id`` +columns on flow/folder/deployment) and follows the expand-contract pattern +under the EXPAND phase: no destructive operations on existing schema. + +Downgrade is intentionally lossy: every authz_* row, every casbin_rule, and +every authz_share / authz_audit_log entry is dropped. Operators running +``alembic downgrade`` past this revision will lose all RBAC policy and audit +state. There is no in-place rollback path for enterprise deployments — back +up the authz tables before downgrading if you intend to roll forward again. """ from __future__ import annotations -from collections.abc import Sequence from datetime import datetime, timezone +from typing import TYPE_CHECKING from uuid import uuid4 import sqlalchemy as sa @@ -16,6 +27,9 @@ import sqlmodel from alembic import op from langflow.utils import migration +if TYPE_CHECKING: + from collections.abc import Sequence + revision: str = "7c8d9e0f1a2b" # pragma: allowlist secret down_revision: str | None = "mb01b2c3d4e5" # pragma: allowlist secret branch_labels: str | Sequence[str] | None = None diff --git a/src/backend/base/langflow/services/authorization/service.py b/src/backend/base/langflow/services/authorization/service.py index 10d2adb463..9a28605812 100644 --- a/src/backend/base/langflow/services/authorization/service.py +++ b/src/backend/base/langflow/services/authorization/service.py @@ -26,6 +26,24 @@ class LangflowAuthorizationService(BaseAuthorizationService): self.settings_service = settings_service self.set_ready() logger.debug("Langflow authorization service initialized") + # Loud, operator-visible warning when AUTHZ_ENABLED is on but the + # registered service is the OSS pass-through. Without the enterprise + # Casbin plugin registered via ``lfx.toml``, ``enforce()`` returns True + # for every request — so flipping the env var alone changes nothing + # except audit-log emission. The dry-run CLI helps verify policy + # decisions, but ops engineers can be misled by a "True" flag. + try: + authz_enabled = bool(getattr(settings_service.auth_settings, "AUTHZ_ENABLED", False)) + except Exception: # noqa: BLE001 — never break startup on a warning probe + authz_enabled = False + if authz_enabled and not self.SUPPORTS_CROSS_USER_FETCH: + logger.warning( + "LANGFLOW_AUTHZ_ENABLED=true but the OSS pass-through authorization service is " + "registered (no enterprise enforcement plugin found). Every enforce() call will " + "return True; route guards still run and audit rows still write, but no policy is " + "applied. Register the enterprise Casbin plugin via lfx.toml or set " + "LANGFLOW_AUTHZ_ENABLED=false to silence this warning." + ) @property def name(self) -> str: diff --git a/src/backend/base/langflow/services/authorization/utils.py b/src/backend/base/langflow/services/authorization/utils.py index 50d2ac97d1..4600a030e2 100644 --- a/src/backend/base/langflow/services/authorization/utils.py +++ b/src/backend/base/langflow/services/authorization/utils.py @@ -95,23 +95,192 @@ def _split_obj(obj: str) -> tuple[str | None, UUID | None]: return resource_type, None -# Strong refs for in-flight audit tasks (awaited on shutdown). +# ----------------------------------------------------------------------------- # +# Batched audit pipeline. +# +# An earlier revision did ``asyncio.create_task(_write())`` per authorization +# decision, with each write opening its own ``session_scope()``. That works on +# light traffic, but on a real workload (every authenticated request emits at +# least one audit row) it turns into a connection-pool storm: per request you +# pay for the request session + a fresh audit session, and the unbounded +# ``_pending_audit_tasks`` set grows until shutdown drain. +# +# This version routes every decision through a bounded queue drained by a +# single long-lived writer task. The writer batches up to ``_AUDIT_BATCH_MAX`` +# rows per ``session_scope()`` and commits them in one INSERT. ``audit_decision`` +# stays a non-blocking ``put_nowait``: if the queue is saturated the row is +# dropped (with a periodic warning) so the request path is never gated on the +# audit DB. +# ----------------------------------------------------------------------------- # + +_AUDIT_QUEUE_MAX = 10_000 +_AUDIT_BATCH_MAX = 100 + + +class _AuditEntry: + """One pending audit row. + + A plain class (not a dataclass) so it can be instantiated cheaply from the + request path without dataclass overhead. The fields mirror + ``AuthzAuditLog`` columns plus the raw ``obj`` string — the writer splits + ``obj`` into ``(resource_type, resource_id)`` once per batch. + """ + + __slots__ = ("action", "details", "obj", "result", "user_id") + + def __init__( + self, + *, + user_id: UUID | None, + action: str, + obj: str, + result: str, + details: dict[str, Any] | None, + ) -> None: + self.user_id = user_id + self.action = action + self.obj = obj + self.result = result + self.details = details + + +# Module-level state. Bound to whichever event loop is running when the first +# ``audit_decision`` call happens. ``_audit_queue_loop`` lets us detect a fresh +# loop (e.g. between pytest test cases) and restart the writer in the new loop +# instead of writing to a queue tied to a dead loop. +_audit_queue: asyncio.Queue[_AuditEntry] | None = None +_audit_queue_loop: asyncio.AbstractEventLoop | None = None +_audit_writer_task: asyncio.Task[None] | None = None +_audit_dropped_count: int = 0 +# Kept as a vestigial public name for backward compatibility with downstream +# callers (and the existing drain helper). The new pipeline tracks the single +# writer task here so ``drain_pending_audit_writes`` can await it. _pending_audit_tasks: set[asyncio.Task[None]] = set() -async def drain_pending_audit_writes(timeout: float = 5.0) -> None: - """Await in-flight audit writes during shutdown (bounded by timeout).""" - pending = {task for task in _pending_audit_tasks if not task.done()} - if not pending: +def _ensure_audit_writer_started() -> asyncio.Queue[_AuditEntry] | None: + """Lazily start the audit writer task in the current event loop. + + Returns the queue, or ``None`` if no event loop is running (audit is + skipped entirely in that case — there's no place to schedule the writer). + """ + global _audit_queue, _audit_queue_loop, _audit_writer_task # noqa: PLW0603 + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return None + + # A fresh event loop replaces the previous queue+writer. Without this, + # a subsequent ``audit_decision`` call (e.g. in a new pytest test) would + # ``put_nowait`` into a queue that no live task is consuming. + if _audit_queue_loop is not loop: + _audit_queue = asyncio.Queue(maxsize=_AUDIT_QUEUE_MAX) + _audit_queue_loop = loop + _audit_writer_task = None + _pending_audit_tasks.clear() + + if _audit_writer_task is None or _audit_writer_task.done(): + _audit_writer_task = loop.create_task(_audit_writer_loop()) + _pending_audit_tasks.add(_audit_writer_task) + _audit_writer_task.add_done_callback(_pending_audit_tasks.discard) + + return _audit_queue + + +async def _audit_writer_loop() -> None: + """Drain the audit queue and write batches to the DB. + + Loops until cancelled. Each iteration blocks on the first row, then greedily + pulls everything else already enqueued up to ``_AUDIT_BATCH_MAX`` and + commits them as a single batch insert. DB exceptions are logged and + swallowed — an audit-table outage must never crash the request path that + triggered the row. + """ + while True: + queue = _audit_queue + if queue is None: + return + try: + first = await queue.get() + except asyncio.CancelledError: + return + + batch: list[_AuditEntry] = [first] + try: + while len(batch) < _AUDIT_BATCH_MAX: + batch.append(queue.get_nowait()) + except asyncio.QueueEmpty: + pass + + try: + await _flush_audit_batch(batch) + except Exception: # noqa: BLE001 — never let the writer die quietly + logger.exception("Authz audit writer batch flush failed for %d row(s)", len(batch)) + finally: + for _ in batch: + queue.task_done() + + +async def _flush_audit_batch(batch: list[_AuditEntry]) -> None: + """Insert a batch of ``_AuditEntry`` rows in a single session.""" + if not batch: return - done, still_pending = await asyncio.wait(pending, timeout=timeout) - if still_pending: - logger.warning("drain_pending_audit_writes timed out with %d pending tasks", len(still_pending)) - # Log task failures; never block shutdown. - for task in done: - exc = task.exception() if not task.cancelled() else None - if exc is not None: - logger.warning("Audit task raised during drain: %s", exc) + # Imported lazily so the request path doesn't pull DB modules until the + # writer first runs (matches the lazy import in the old per-row path). + from langflow.services.database.models.auth import AuthzAuditLog + from langflow.services.deps import session_scope + + async with session_scope() as session: + for entry in batch: + resource_type, resource_id = _split_obj(entry.obj) + session.add( + AuthzAuditLog( + user_id=entry.user_id, + action=entry.action, + resource_type=resource_type, + resource_id=resource_id, + result=entry.result, + details=entry.details, + ) + ) + + +async def drain_pending_audit_writes(timeout: float = 5.0) -> None: + """Flush the audit queue and stop the writer (bounded by ``timeout``). + + Safe to call multiple times; safe to call when no audit traffic has run. + Splits the timeout between draining the queue and awaiting writer + cancellation so neither side can hang shutdown indefinitely. + """ + global _audit_writer_task # noqa: PLW0603 + + queue = _audit_queue + writer = _audit_writer_task + if queue is None or writer is None: + return + + drain_budget = max(0.1, timeout * 0.8) + cancel_budget = max(0.1, timeout - drain_budget) + + try: + await asyncio.wait_for(queue.join(), timeout=drain_budget) + except asyncio.TimeoutError: + logger.warning( + "drain_pending_audit_writes timed out after %.2fs with %d row(s) pending", + drain_budget, + queue.qsize(), + ) + + if not writer.done(): + writer.cancel() + from contextlib import suppress + + with suppress(asyncio.CancelledError): + await asyncio.wait_for(writer, timeout=cancel_budget) + + _pending_audit_tasks.discard(writer) + _audit_writer_task = None async def audit_decision( @@ -122,37 +291,49 @@ async def audit_decision( result: str, details: dict[str, Any] | None = None, ) -> None: - """Schedule an AuthzAuditLog write (fire-and-forget).""" + """Enqueue an AuthzAuditLog row for batched background insertion. + + Non-blocking from the caller's perspective. When the queue is saturated the + row is dropped and a sample-rate warning is emitted so a stuck pipeline is + operator-visible. Audit is fully bypassed when + ``AUTHZ_AUDIT_ENABLED=False`` (the default). + """ + global _audit_dropped_count # noqa: PLW0603 + settings = get_settings_service() auth_settings = settings.auth_settings - # Audit can run while AUTHZ_ENABLED is false. - if not getattr(auth_settings, "AUTHZ_AUDIT_ENABLED", True): + # Audit is independent of enforcement. ``AuthSettings.AUTHZ_AUDIT_ENABLED`` + # defaults to ``False`` (see lfx/services/settings/auth.py) because the + # background writer still consumes a DB connection; operators opt in. + if not getattr(auth_settings, "AUTHZ_AUDIT_ENABLED", False): return - resource_type, resource_id = _split_obj(obj) + queue = _ensure_audit_writer_started() + if queue is None: + # No running event loop — nothing to schedule against. The caller is + # likely outside an async context (e.g. a sync test); silently skip. + return - async def _write() -> None: - try: - from langflow.services.database.models.auth import AuthzAuditLog - from langflow.services.deps import session_scope - - async with session_scope() as session: - session.add( - AuthzAuditLog( - user_id=user_id, - action=action, - resource_type=resource_type, - resource_id=resource_id, - result=result, - details=details, - ) - ) - except Exception: # noqa: BLE001 — audit must never raise into the request path - logger.exception("Failed to write AuthzAuditLog row") - - task = asyncio.create_task(_write()) - _pending_audit_tasks.add(task) - task.add_done_callback(_pending_audit_tasks.discard) + entry = _AuditEntry( + user_id=user_id, + action=action, + obj=obj, + result=result, + details=details, + ) + try: + queue.put_nowait(entry) + except asyncio.QueueFull: + _audit_dropped_count += 1 + # Sample-log so a saturated pipeline surfaces in ops dashboards without + # spamming the log on every drop. First drop + every 1000th drop. + if _audit_dropped_count == 1 or _audit_dropped_count % 1000 == 0: + logger.warning( + "AuthzAuditLog queue full (%d/%d); dropped %d row(s) total. DB writer is likely behind or stalled.", + queue.qsize(), + _AUDIT_QUEUE_MAX, + _audit_dropped_count, + ) async def ensure_permission( diff --git a/src/backend/base/langflow/services/utils.py b/src/backend/base/langflow/services/utils.py index b42c0a5053..e8eafa6c57 100644 --- a/src/backend/base/langflow/services/utils.py +++ b/src/backend/base/langflow/services/utils.py @@ -410,6 +410,49 @@ async def clean_transactions(settings_service: SettingsService, session: AsyncSe # Don't re-raise since this is a cleanup task +async def clean_authz_audit_log(settings_service: SettingsService, session: AsyncSession) -> int: + """Delete authz_audit_log rows older than ``AUTHZ_AUDIT_RETENTION_DAYS``. + + Retention is configured via ``AuthSettings.AUTHZ_AUDIT_RETENTION_DAYS``; + setting it to ``0`` disables pruning so operators relying on an external + archival pipeline (Postgres partitioning, SIEM export) can opt out without + losing rows here. The function is intentionally best-effort: failures are + logged and swallowed so startup never fails because the audit table is + transiently unreachable. + + Returns the number of rows deleted (best-effort; ``-1`` when the rowcount + is unavailable from the driver). + """ + try: + retention_days = int(getattr(settings_service.auth_settings, "AUTHZ_AUDIT_RETENTION_DAYS", 90)) + except Exception: # noqa: BLE001 — settings shape can vary in tests/stubs + retention_days = 90 + if retention_days <= 0: + logger.debug("authz_audit_log retention disabled (AUTHZ_AUDIT_RETENTION_DAYS=%d)", retention_days) + return 0 + + from datetime import datetime, timedelta, timezone + + from langflow.services.database.models.auth import AuthzAuditLog + + cutoff = datetime.now(timezone.utc) - timedelta(days=retention_days) + try: + delete_stmt = delete(AuthzAuditLog).where(col(AuthzAuditLog.timestamp) < cutoff) + result = await session.exec(delete_stmt) + deleted = getattr(result, "rowcount", None) + deleted_count = int(deleted) if deleted is not None and deleted >= 0 else -1 + logger.debug( + "authz_audit_log cleanup removed %s rows older than %d days", + deleted_count if deleted_count >= 0 else "?", + retention_days, + ) + except (sqlalchemy_exc.SQLAlchemyError, asyncio.TimeoutError) as exc: + logger.warning("authz_audit_log cleanup failed: %s", exc) + return -1 + else: + return deleted_count + + async def clean_vertex_builds(settings_service: SettingsService, session: AsyncSession) -> None: """Clean up old vertex builds from the database. @@ -568,3 +611,4 @@ async def initialize_services(*, fix_migration: bool = False) -> None: async with session_scope() as session: await clean_transactions(settings_service, session) await clean_vertex_builds(settings_service, session) + await clean_authz_audit_log(settings_service, session) diff --git a/src/backend/tests/unit/services/authorization/test_audit_retention.py b/src/backend/tests/unit/services/authorization/test_audit_retention.py new file mode 100644 index 0000000000..b9ad8662cb --- /dev/null +++ b/src/backend/tests/unit/services/authorization/test_audit_retention.py @@ -0,0 +1,151 @@ +"""Tests for the authz_audit_log retention helper and pass-through startup warning.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from langflow.services.authorization.service import LangflowAuthorizationService +from langflow.services.database.models.auth import AuthzAuditLog +from langflow.services.utils import clean_authz_audit_log +from sqlalchemy.ext.asyncio import create_async_engine +from sqlmodel import SQLModel, select +from sqlmodel.ext.asyncio.session import AsyncSession + + +@pytest.fixture(name="authz_engine") +def authz_engine(): + engine = create_async_engine("sqlite+aiosqlite:///:memory:", echo=False) + yield engine + engine.sync_engine.dispose() + + +@pytest.fixture(name="authz_session") +async def authz_session(authz_engine): + async with authz_engine.begin() as conn: + await conn.run_sync(SQLModel.metadata.create_all) + async with AsyncSession(authz_engine, expire_on_commit=False) as session: + yield session + async with authz_engine.begin() as conn: + await conn.run_sync(SQLModel.metadata.drop_all) + await authz_engine.dispose() + + +def _audit_row(*, days_old: int) -> AuthzAuditLog: + timestamp = datetime.now(timezone.utc) - timedelta(days=days_old) + return AuthzAuditLog( + user_id=uuid4(), + action="flow:read", + resource_type="flow", + resource_id=uuid4(), + result="allow", + details={"domain": "*"}, + timestamp=timestamp, + ) + + +def _settings(*, retention_days: int) -> SimpleNamespace: + return SimpleNamespace( + auth_settings=SimpleNamespace( + AUTHZ_AUDIT_RETENTION_DAYS=retention_days, + ), + ) + + +@pytest.mark.anyio +async def test_clean_authz_audit_log_prunes_rows_past_retention(authz_session): + """Rows older than the retention window are deleted; fresh rows survive.""" + authz_session.add(_audit_row(days_old=120)) # outside 90-day window + authz_session.add(_audit_row(days_old=95)) # outside 90-day window + authz_session.add(_audit_row(days_old=5)) # inside window + await authz_session.commit() + + await clean_authz_audit_log(_settings(retention_days=90), authz_session) + await authz_session.commit() + + surviving = (await authz_session.exec(select(AuthzAuditLog))).all() + # Only the 5-day-old row should remain. + assert len(surviving) == 1 + + +@pytest.mark.anyio +async def test_clean_authz_audit_log_zero_disables_retention(authz_session): + """A retention of 0 days is a no-op so external archival pipelines can manage the table.""" + for days_old in (1, 30, 200, 365): + authz_session.add(_audit_row(days_old=days_old)) + await authz_session.commit() + + await clean_authz_audit_log(_settings(retention_days=0), authz_session) + await authz_session.commit() + + surviving = (await authz_session.exec(select(AuthzAuditLog))).all() + assert len(surviving) == 4 + + +@pytest.mark.anyio +async def test_clean_authz_audit_log_empty_table(authz_session): + """Pruning an empty table is a no-op and must not raise.""" + deleted = await clean_authz_audit_log(_settings(retention_days=30), authz_session) + assert deleted in (0, -1) # -1 acceptable when driver doesn't report rowcount + + +class _RecordingLogger: + """Stand-in for loguru's logger that records each call by level.""" + + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + def warning(self, message: str, *args, **kwargs) -> None: # noqa: ARG002 + self.calls.append(("warning", message)) + + def debug(self, message: str, *args, **kwargs) -> None: # noqa: ARG002 + self.calls.append(("debug", message)) + + def info(self, message: str, *args, **kwargs) -> None: # noqa: ARG002 + self.calls.append(("info", message)) + + +@pytest.mark.anyio +async def test_passthrough_warning_emitted_when_authz_enabled(monkeypatch): + """LangflowAuthorizationService warns when AUTHZ_ENABLED=True but plugin is missing.""" + from langflow.services.authorization import service as authz_service_module + + recorder = _RecordingLogger() + monkeypatch.setattr(authz_service_module, "logger", recorder) + + settings = SimpleNamespace( + auth_settings=SimpleNamespace( + AUTHZ_ENABLED=True, + AUTHZ_SUPERUSER_BYPASS=True, + ) + ) + LangflowAuthorizationService(settings) + + warning_messages = [msg for level, msg in recorder.calls if level == "warning"] + assert any("OSS pass-through" in msg for msg in warning_messages), ( + f"Expected a WARNING about the OSS pass-through service; got {warning_messages}" + ) + + +@pytest.mark.anyio +async def test_passthrough_warning_silent_when_authz_disabled(monkeypatch): + """No warning is emitted when AUTHZ_ENABLED=False.""" + from langflow.services.authorization import service as authz_service_module + + recorder = _RecordingLogger() + monkeypatch.setattr(authz_service_module, "logger", recorder) + + settings = SimpleNamespace( + auth_settings=SimpleNamespace( + AUTHZ_ENABLED=False, + AUTHZ_SUPERUSER_BYPASS=True, + ) + ) + LangflowAuthorizationService(settings) + + warning_messages = [msg for level, msg in recorder.calls if level == "warning"] + assert not any("OSS pass-through" in msg for msg in warning_messages), ( + f"No pass-through warning expected when AUTHZ_ENABLED=False; got {warning_messages}" + ) diff --git a/src/backend/tests/unit/services/authorization/test_utils.py b/src/backend/tests/unit/services/authorization/test_utils.py index 7b9911400f..a958ba4b68 100644 --- a/src/backend/tests/unit/services/authorization/test_utils.py +++ b/src/backend/tests/unit/services/authorization/test_utils.py @@ -471,115 +471,110 @@ def test_split_obj_non_uuid_suffix_returns_none_id(): assert resource_id is None +@pytest.fixture +def patched_audit_flush(monkeypatch): + """Replace _flush_audit_batch with a recorder so we exercise the writer without touching the DB.""" + flushed: list[list[object]] = [] + + async def _record(batch): + flushed.append(list(batch)) + + monkeypatch.setattr(authz_utils, "_flush_audit_batch", _record) + return flushed + + +async def _reset_audit_pipeline() -> None: + """Best-effort teardown so each test starts with a clean audit pipeline.""" + await authz_utils.drain_pending_audit_writes(timeout=0.5) + authz_utils._audit_queue = None + authz_utils._audit_queue_loop = None + authz_utils._audit_writer_task = None + authz_utils._audit_dropped_count = 0 + authz_utils._pending_audit_tasks.clear() + + @pytest.mark.anyio -async def test_audit_decision_runs_when_authz_disabled_but_audit_on(monkeypatch): - """Audit is independent of enforcement now. +async def test_audit_decision_runs_when_authz_disabled_but_audit_on(monkeypatch, patched_audit_flush): + """Audit is independent of enforcement. Previously ``audit_decision`` short-circuited when ``AUTHZ_ENABLED=False``, which meant share CRUD writes left no audit trail on default installs. The new contract gates only on ``AUTHZ_AUDIT_ENABLED`` so operators can observe traffic before flipping enforcement on. """ + await _reset_audit_pipeline() _install_settings(monkeypatch, authz_enabled=False, audit_enabled=True) - scheduled: list[object] = [] - - class _FakeTask: - def add_done_callback(self, cb): - self._cb = cb - - def done(self) -> bool: - return False - - def _capture(coro): - coro.close() - task = _FakeTask() - scheduled.append(task) - return task - - monkeypatch.setattr("asyncio.create_task", _capture) await authz_utils.audit_decision(user_id=uuid4(), action="flow:read", obj="flow:x", result="allow") - try: - assert len(scheduled) == 1 - finally: - # Don't pollute the module-global pending set for downstream tests. - for task in scheduled: - authz_utils._pending_audit_tasks.discard(task) + await authz_utils.drain_pending_audit_writes(timeout=1.0) + + assert sum(len(b) for b in patched_audit_flush) == 1 @pytest.mark.anyio -async def test_audit_decision_noop_when_audit_disabled(monkeypatch): +async def test_audit_decision_noop_when_audit_disabled(monkeypatch, patched_audit_flush): """AUTHZ_AUDIT_ENABLED=False suppresses audit writes.""" + await _reset_audit_pipeline() _install_settings(monkeypatch, authz_enabled=True, audit_enabled=False) - scheduled: list[object] = [] - monkeypatch.setattr("asyncio.create_task", lambda coro: scheduled.append(coro) or coro) await authz_utils.audit_decision(user_id=uuid4(), action="flow:read", obj="flow:x", result="allow") - assert scheduled == [] + await authz_utils.drain_pending_audit_writes(timeout=0.5) + + assert patched_audit_flush == [] @pytest.mark.anyio -async def test_audit_decision_schedules_task_when_enabled(monkeypatch): - """A scheduled coroutine is produced when both flags are on. - - The new implementation tracks the task in ``_pending_audit_tasks`` so the - event loop cannot GC it mid-write and so shutdown can drain it. Our mock - returns a lightweight stand-in that supports ``add_done_callback`` exactly - as a real ``asyncio.Task`` does. - """ +async def test_audit_decision_enqueues_when_enabled(monkeypatch, patched_audit_flush): + """When both flags are on, audit_decision enqueues an entry that the background writer flushes.""" + await _reset_audit_pipeline() _install_settings(monkeypatch, authz_enabled=True, audit_enabled=True) - scheduled: list[object] = [] - class _FakeTask: - def __init__(self, coro): - self._coro = coro - self._callbacks: list = [] + user_id = uuid4() + await authz_utils.audit_decision(user_id=user_id, action="flow:read", obj="flow:abc", result="allow") + # Drain forces the writer to flush before we inspect. + await authz_utils.drain_pending_audit_writes(timeout=1.0) - def add_done_callback(self, cb): - self._callbacks.append(cb) - - def done(self) -> bool: - return False - - def _capture(coro): - coro.close() # don't actually run — the function imports from DB modules. - task = _FakeTask(coro) - scheduled.append(task) - return task - - monkeypatch.setattr("asyncio.create_task", _capture) - - await authz_utils.audit_decision(user_id=uuid4(), action="flow:read", obj="flow:x", result="allow") - try: - assert len(scheduled) == 1 - # The implementation must hold a reference so the GC can't drop the task. - assert scheduled[0] in authz_utils._pending_audit_tasks - # And it must register a done-callback that removes itself from the set. - assert scheduled[0]._callbacks, "audit_decision must register a done-callback to clean up _pending_audit_tasks" - finally: - # The fake task is not a real ``asyncio.Task`` — clean up the - # module-global set so the next test starts from an empty state. - authz_utils._pending_audit_tasks.discard(scheduled[0]) + assert len(patched_audit_flush) == 1 + batch = patched_audit_flush[0] + assert len(batch) == 1 + entry = batch[0] + assert entry.user_id == user_id + assert entry.action == "flow:read" + assert entry.obj == "flow:abc" + assert entry.result == "allow" @pytest.mark.anyio -async def test_drain_pending_audit_writes_awaits_tasks(): - """drain_pending_audit_writes waits for tracked tasks to finish.""" - import asyncio +async def test_audit_decision_batches_multiple_entries(monkeypatch, patched_audit_flush): + """Multiple concurrent audit_decision calls coalesce into a single DB batch. - finished = asyncio.Event() + This is the contract we want — the writer should pull every entry already + in the queue when it wakes up, so we make N decisions before yielding and + expect ONE batch of N rows, not N separate session_scope opens. + """ + await _reset_audit_pipeline() + _install_settings(monkeypatch, authz_enabled=True, audit_enabled=True) - async def _slow() -> None: - await asyncio.sleep(0.01) - finished.set() - - task = asyncio.create_task(_slow()) - authz_utils._pending_audit_tasks.add(task) - task.add_done_callback(authz_utils._pending_audit_tasks.discard) + for _ in range(5): + await authz_utils.audit_decision(user_id=uuid4(), action="flow:read", obj="flow:x", result="allow") await authz_utils.drain_pending_audit_writes(timeout=1.0) - assert finished.is_set() - assert task not in authz_utils._pending_audit_tasks + + total_rows = sum(len(batch) for batch in patched_audit_flush) + assert total_rows == 5 + # All entries are emitted before the first await, so they should land in a single batch. + assert len(patched_audit_flush) == 1, ( + f"Expected 1 batch of 5 rows, got {len(patched_audit_flush)} batches " + f"with sizes {[len(b) for b in patched_audit_flush]}" + ) + + +@pytest.mark.anyio +async def test_drain_pending_audit_writes_is_safe_when_idle(): + """drain_pending_audit_writes is a no-op when no audit traffic has run.""" + await _reset_audit_pipeline() + # Must not raise. + await authz_utils.drain_pending_audit_writes(timeout=0.1) @pytest.mark.anyio diff --git a/src/lfx/src/lfx/services/settings/auth.py b/src/lfx/src/lfx/services/settings/auth.py index e713cfc600..a1d3db5178 100644 --- a/src/lfx/src/lfx/services/settings/auth.py +++ b/src/lfx/src/lfx/services/settings/auth.py @@ -152,6 +152,17 @@ class AuthSettings(BaseSettings): "contend with concurrent write transactions ('database is locked')." ), ) + AUTHZ_AUDIT_RETENTION_DAYS: int = Field( + default=90, + ge=0, + description=( + "Number of days to retain rows in ``authz_audit_log``. Older rows are deleted on " + "startup and by the periodic cleanup task. Set to 0 to disable retention pruning " + "(append-only — the table will grow without bound; pair with an external archival " + "or partitioning job). The default of 90 days bounds steady-state size for typical " + "enterprise deployments." + ), + ) pwd_context: CryptContext = CryptContext(schemes=["bcrypt"], deprecated="auto")