test_service.py

"""Tests for TelemetryWriterService.

The service is exercised end-to-end against a real in-memory SQLite database
(no mocking of the persistence layer). We bypass ``start()`` in most tests so
we don't have to bring up the entire langflow service manager — instead each
test wires up the writer's internals directly against the test
``async_session`` fixture's engine.
"""

from __future__ import annotations

import asyncio
import shutil
import tempfile
import warnings
from pathlib import Path
from uuid import uuid4

import pytest
from langflow.services.database.models.transactions.model import TransactionBase, TransactionTable
from langflow.services.database.models.vertex_builds.model import VertexBuildBase, VertexBuildTable
from langflow.services.telemetry_writer.service import (
    _FAILURE_ESCALATION_THRESHOLD,
    TelemetryWriterService,
    _write_owner_file,
)
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession as SQLModelAsyncSession


class _FakeSettings:
    """Mutable settings stand-in.

    Real ``Settings`` is a pydantic-settings model whose constructor kwargs are
    overridden by env-var sources; tests can't reliably set fields through it.
    A plain object with the attributes the writer reads via ``getattr`` is
    sufficient and side-steps that quirk.
    """

    def __init__(self, **overrides):
        defaults = {
            "telemetry_writer_enabled": True,
            "telemetry_writer_batch_size": 200,
            "telemetry_writer_flush_interval_s": 0.5,
            "telemetry_writer_cleanup_interval_s": 60,
            "telemetry_writer_max_queue": 100_000,
            "telemetry_writer_outbox_dir": None,
            "telemetry_writer_shutdown_drain_s": 5.0,
            "telemetry_writer_size_strategy": "count",
            "telemetry_writer_batch_size_bytes": 262_144,
            "telemetry_writer_max_queue_bytes": 209_715_200,
            "max_transactions_to_keep": 3000,
            "max_vertex_builds_to_keep": 3000,
            "max_vertex_builds_per_vertex": 50,
        }
        for key, value in {**defaults, **overrides}.items():
            setattr(self, key, value)


class _FakeSettingsService:
    def __init__(self, **overrides):
        self.settings = _FakeSettings(**overrides)


def _build_writer(settings_overrides=None) -> TelemetryWriterService:
    return TelemetryWriterService(_FakeSettingsService(**(settings_overrides or {})))


def _make_transaction_row(flow_id, vertex_id="v1") -> dict:
    base = TransactionBase(
        vertex_id=vertex_id,
        target_id=None,
        inputs={"x": 1},
        outputs={"y": 2},
        status="success",
        error=None,
        flow_id=flow_id,
    )
    return TransactionTable(**base.model_dump()).model_dump(mode="python")


def _make_vertex_build_row(flow_id, vertex_id="v1") -> dict:
    base = VertexBuildBase(
        id=vertex_id,
        flow_id=flow_id,
        valid=True,
        params="p",
        data={"r": 1},
        artifacts=None,
    )
    return VertexBuildTable(**base.model_dump()).model_dump(mode="python")


@pytest.fixture
async def writer_with_engine():
    """Yield (writer, engine) wired together with a fresh in-memory SQLite."""
    engine = create_async_engine(
        "sqlite+aiosqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    writer = _build_writer()
    writer._engine = engine
    writer._session_maker = async_sessionmaker(engine, class_=SQLModelAsyncSession, expire_on_commit=False)
    writer._started = True
    writer._shutdown_event = asyncio.Event()
    try:
        yield writer, engine
    finally:
        writer._started = False
        async with engine.begin() as conn:
            await conn.run_sync(SQLModel.metadata.drop_all)
        await engine.dispose()


def test_is_enabled_reads_settings() -> None:
    enabled = _build_writer({"telemetry_writer_enabled": True})
    disabled = _build_writer({"telemetry_writer_enabled": False})
    assert enabled.is_enabled() is True
    assert disabled.is_enabled() is False


def test_enqueue_when_not_running_returns_false() -> None:
    writer = _build_writer()
    assert writer.enqueue_transaction({"any": "thing"}) is False
    assert writer.enqueue_vertex_build({"any": "thing"}) is False


def test_enqueue_when_running_buffers_payload(writer_with_engine) -> None:
    writer, _ = writer_with_engine
    payload = {"flow_id": str(uuid4()), "vertex_id": "v1"}
    assert writer.enqueue_transaction(payload) is True
    assert len(writer._tx_buffer) == 1
    assert writer.enqueued_transactions == 1


def test_overflow_drops_oldest(writer_with_engine) -> None:
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_max_queue = 3
    for i in range(5):
        writer.enqueue_transaction({"i": i})
    # Cap is 3 → oldest two dropped, newest three retained.
    assert len(writer._tx_buffer) == 3
    assert writer.dropped_transactions == 2
    assert [item["i"] for item in writer._tx_buffer] == [2, 3, 4]


def test_count_strategy_ignores_byte_caps(writer_with_engine) -> None:
    """Default 'count' strategy must not track or cap bytes — preserves legacy semantics."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "count"
    writer.settings_service.settings.telemetry_writer_max_queue_bytes = 1  # tiny — would trip on bytes
    for i in range(5):
        writer.enqueue_transaction({"big": "x" * 1000, "i": i})
    assert len(writer._tx_buffer) == 5
    assert writer.dropped_transactions == 0
    # Byte tracking is disabled under 'count'.
    assert writer._tx_bytes == 0
    assert len(writer._tx_sizes) == 0


def test_bytes_strategy_drops_oldest_when_byte_cap_exceeded(writer_with_engine) -> None:
    """'bytes' strategy bounds memory regardless of row count."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "bytes"
    writer.settings_service.settings.telemetry_writer_max_queue = 10_000  # high — must not be the trigger
    writer.settings_service.settings.telemetry_writer_max_queue_bytes = 200
    # {"payload": "x"*100, "i": N} encodes to 123 bytes. Two of them (246) bust
    # the 200B cap, so on each subsequent enqueue the oldest is dropped first.
    for i in range(3):
        writer.enqueue_transaction({"payload": "x" * 100, "i": i})
    assert writer._tx_bytes == 123
    assert len(writer._tx_buffer) == 1
    assert writer.dropped_transactions == 2
    assert writer.dropped_transactions_bytes == 246


def test_either_strategy_trips_on_first_threshold(writer_with_engine) -> None:
    """'either' fires whichever cap (rows or bytes) is breached first."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "either"
    writer.settings_service.settings.telemetry_writer_max_queue = 3
    writer.settings_service.settings.telemetry_writer_max_queue_bytes = 10_000_000
    for i in range(5):
        writer.enqueue_transaction({"i": i})
    # Row cap trips first.
    assert len(writer._tx_buffer) == 3
    assert writer.dropped_transactions == 2


def test_drain_batch_respects_byte_budget(writer_with_engine) -> None:
    """Writer flush must split batches by bytes when strategy uses bytes."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "bytes"
    for i in range(10):
        writer.enqueue_transaction({"payload": "x" * 100, "i": i})
    # Each row encodes to 123 bytes. With a 250B budget the loop pops 3 rows
    # (123, 246, 369) and breaks on the third when batch_bytes >= max_bytes.
    batch = writer._drain_batch("transactions", max_n=10, max_bytes=250)
    assert len(batch) == 3
    assert len(writer._tx_buffer) == 7
    assert len(writer._tx_sizes) == 7
    assert writer._tx_bytes == 7 * 123


def test_drain_batch_always_emits_one_row_even_when_oversized(writer_with_engine) -> None:
    """A single row larger than the byte budget must still make progress."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "bytes"
    writer.enqueue_transaction({"big": "x" * 10_000})
    batch = writer._drain_batch("transactions", max_n=200, max_bytes=10)
    assert len(batch) == 1


def test_return_batch_to_buffer_restores_sizes(writer_with_engine) -> None:
    """Cancel/retry path must put rows AND their sizes back so accounting stays consistent."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "either"
    writer.enqueue_transaction({"a": 1})
    writer.enqueue_transaction({"b": 2})
    bytes_before = writer._tx_bytes
    batch = writer._drain_batch("transactions", max_n=2, max_bytes=10_000)
    assert writer._tx_bytes == 0
    writer._return_batch_to_buffer("transactions", batch)
    assert writer._tx_bytes == bytes_before
    assert len(writer._tx_sizes) == len(writer._tx_buffer) == 2


def test_bytes_strategy_round_trip_through_spill_and_restore(tmp_path: Path) -> None:
    """Spill must reset size tracking; restore must rebuild it via _enqueue."""
    own_dir = tmp_path / "pid"
    own_dir.mkdir()
    writer = _build_writer({"telemetry_writer_size_strategy": "either"})
    writer._started = True
    writer.enqueue_transaction({"flow": "a", "i": 1})
    writer.enqueue_transaction({"flow": "b", "i": 2})
    original_bytes = writer._tx_bytes
    assert original_bytes > 0
    assert len(writer._tx_sizes) == 2

    writer._spill_to_disk(own_dir, kind="transactions", buffer=writer._tx_buffer)
    # Spill drained the buffer and must zero size tracking.
    assert writer._tx_bytes == 0
    assert len(writer._tx_sizes) == 0
    assert not writer._tx_buffer

    # Restore into a fresh writer with the same strategy.
    restored = _build_writer({"telemetry_writer_size_strategy": "either"})
    restored._started = True
    restored._restore_from_disk(own_dir, kind="transactions", buffer=restored._tx_buffer)
    assert len(restored._tx_buffer) == 2
    assert len(restored._tx_sizes) == 2
    # Bytes must match the pre-spill total — restore re-encodes via _enqueue.
    assert restored._tx_bytes == original_bytes


def test_bytes_strategy_round_trip_through_adoption(tmp_path: Path) -> None:
    """Adopting an orphan outbox under 'bytes' strategy must populate size deques."""
    dead_pid = _find_dead_pid()
    dead_dir = tmp_path / str(dead_pid)
    dead_dir.mkdir()
    _write_owner_file(dead_dir)
    spill_writer = _build_writer({"telemetry_writer_size_strategy": "either"})
    spill_writer._started = True
    spill_writer.enqueue_transaction({"orphan": True, "i": 1})
    spill_writer.enqueue_transaction({"orphan": True, "i": 2})
    spilled_bytes = spill_writer._tx_bytes
    spill_writer._spill_to_disk(dead_dir, kind="transactions", buffer=spill_writer._tx_buffer)

    adopter = _build_writer({"telemetry_writer_size_strategy": "either"})
    adopter._adopt_orphan_outboxes(tmp_path, own_pid=1)
    assert len(adopter._tx_buffer) == 2
    assert len(adopter._tx_sizes) == 2
    assert adopter._tx_bytes == spilled_bytes


async def test_sweeper_calls_heartbeat_and_prune(tmp_path: Path) -> None:
    """One sweeper tick must touch our owner file and prune stale foreign dirs."""
    import json as _json
    import os as _os
    import time as _time

    own_dir = tmp_path / "1234"
    own_dir.mkdir()
    _write_owner_file(own_dir)
    owner_file = own_dir / "owner.json"
    aged = _time.time() - 7200
    _os.utime(owner_file, (aged, aged))
    own_before = owner_file.stat().st_mtime

    foreign_dir = tmp_path / "9999"
    foreign_dir.mkdir()
    (foreign_dir / "owner.json").write_text(
        _json.dumps({"host": "dead-foreign-pod", "boot": "x", "pid": 9999, "started_at": 0})
    )
    _os.utime(foreign_dir / "owner.json", (aged, aged))

    writer = _build_writer(
        {
            "telemetry_writer_outbox_dir": str(tmp_path),
            "telemetry_writer_orphan_max_age_s": 3600.0,
            # Skip the real retention pass — it needs a DB session.
            "telemetry_writer_cleanup_interval_s": 0.01,
        }
    )
    writer._own_outbox_dir = own_dir
    writer._shutdown_event = asyncio.Event()
    # Bypass _run_retention_pass by clearing the session_maker; the sweeper
    # catches the exception path either way, but None makes it a no-op.
    writer._session_maker = None

    async def _stop_after_one_tick() -> None:
        await asyncio.sleep(0.1)
        writer._shutdown_event.set()

    stopper = asyncio.create_task(_stop_after_one_tick())
    await writer._run_sweeper()
    await stopper

    assert owner_file.stat().st_mtime > own_before
    assert not foreign_dir.exists()


async def test_flush_inserts_transactions_and_vertex_builds(writer_with_engine) -> None:
    writer, engine = writer_with_engine
    flow_id = uuid4()
    tx_batch = [_make_transaction_row(flow_id) for _ in range(5)]
    vb_batch = [_make_vertex_build_row(flow_id) for _ in range(3)]

    await writer._flush(tx_batch, vb_batch)

    async with AsyncSession(engine) as session:
        tx_count = await session.scalar(select(func.count()).select_from(TransactionTable))
        vb_count = await session.scalar(select(func.count()).select_from(VertexBuildTable))
    assert tx_count == 5
    assert vb_count == 3
    # Sweeper-driven retention requires dirty-flow tracking; flush populates it.
    assert str(flow_id) in writer._dirty_tx_flows
    assert str(flow_id) in writer._dirty_vb_flows


async def test_flush_and_retention_do_not_call_deprecated_session_execute(writer_with_engine) -> None:
    writer, _ = writer_with_engine
    writer.settings_service.settings.max_transactions_to_keep = 1
    writer.settings_service.settings.max_vertex_builds_to_keep = 1
    flow_id = uuid4()

    tx_batch = [_make_transaction_row(flow_id) for _ in range(3)]
    vb_batch = [_make_vertex_build_row(flow_id, vertex_id="v1") for _ in range(3)]

    with warnings.catch_warnings():
        warnings.filterwarnings("error", message="(?s).*session\\.exec\\(\\).*", category=DeprecationWarning)
        await writer._flush(tx_batch, vb_batch)
        await writer._run_retention_pass()


async def test_retention_sweep_caps_transactions_per_flow(writer_with_engine) -> None:
    writer, engine = writer_with_engine
    writer.settings_service.settings.max_transactions_to_keep = 3
    flow_id = uuid4()
    tx_batch = [_make_transaction_row(flow_id) for _ in range(10)]
    await writer._flush(tx_batch, [])
    async with AsyncSession(engine) as session:
        before = await session.scalar(select(func.count()).select_from(TransactionTable))
    assert before == 10

    await writer._run_retention_pass()

    async with AsyncSession(engine) as session:
        after = await session.scalar(select(func.count()).select_from(TransactionTable))
    assert after == 3


async def test_retention_sweep_caps_vertex_builds_globally(writer_with_engine) -> None:
    writer, engine = writer_with_engine
    writer.settings_service.settings.max_vertex_builds_to_keep = 4
    writer.settings_service.settings.max_vertex_builds_per_vertex = 50
    flow_id = uuid4()
    vb_batch = [_make_vertex_build_row(flow_id, vertex_id=f"v{i}") for i in range(8)]
    await writer._flush([], vb_batch)
    await writer._run_retention_pass()

    async with AsyncSession(engine) as session:
        count = await session.scalar(select(func.count()).select_from(VertexBuildTable))
    assert count == 4


async def test_writer_loop_drains_buffers(writer_with_engine) -> None:
    writer, engine = writer_with_engine
    writer.settings_service.settings.telemetry_writer_batch_size = 2
    writer.settings_service.settings.telemetry_writer_flush_interval_s = 0.01
    flow_id = uuid4()
    for _ in range(5):
        writer.enqueue_transaction(_make_transaction_row(flow_id))
        writer.enqueue_vertex_build(_make_vertex_build_row(flow_id))

    task = asyncio.create_task(writer._run_writer())
    # Give the writer time to flush all batches.
    for _ in range(20):
        await asyncio.sleep(0.05)
        if not writer._tx_buffer and not writer._vb_buffer:
            break
    writer._shutdown_event.set()
    await asyncio.wait_for(task, timeout=2)

    async with AsyncSession(engine) as session:
        tx_count = await session.scalar(select(func.count()).select_from(TransactionTable))
        vb_count = await session.scalar(select(func.count()).select_from(VertexBuildTable))
    assert tx_count == 5
    assert vb_count == 5
    assert writer.failed_batches == 0


def test_spill_and_restore_round_trip(tmp_path: Path) -> None:
    writer = _build_writer()
    own_dir = tmp_path / str(12345)
    own_dir.mkdir()

    # Producer-side fills in-memory buffer.
    writer._tx_buffer.extend([{"i": 1}, {"i": 2}])
    writer._vb_buffer.extend([{"j": 1}])

    writer._spill_to_disk(own_dir, kind="transactions", buffer=writer._tx_buffer)
    writer._spill_to_disk(own_dir, kind="vertex_builds", buffer=writer._vb_buffer)
    assert not writer._tx_buffer
    assert not writer._vb_buffer

    # New writer comes up against the same outbox directory and restores.
    writer2 = _build_writer()
    writer2._restore_from_disk(own_dir, kind="transactions", buffer=writer2._tx_buffer)
    writer2._restore_from_disk(own_dir, kind="vertex_builds", buffer=writer2._vb_buffer)
    assert list(writer2._tx_buffer) == [{"i": 1}, {"i": 2}]
    assert list(writer2._vb_buffer) == [{"j": 1}]


async def test_spill_restore_round_trip_writes_realistic_payload_to_db(writer_with_engine, tmp_path: Path) -> None:
    """Realistic payloads (UUID, datetime) survive JSON spill/restore and INSERT.

    Guards against the SQLite outbox's ``json.dumps(default=str)`` encoding
    breaking SQLAlchemy core inserts on the UUID/datetime columns of
    ``transaction`` / ``vertex_build`` once a row has been spilled to disk.
    """
    writer, engine = writer_with_engine
    flow_id = uuid4()
    writer._tx_buffer.append(_make_transaction_row(flow_id))
    writer._vb_buffer.append(_make_vertex_build_row(flow_id))

    own_dir = tmp_path / "pid"
    own_dir.mkdir()
    writer._spill_to_disk(own_dir, kind="transactions", buffer=writer._tx_buffer)
    writer._spill_to_disk(own_dir, kind="vertex_builds", buffer=writer._vb_buffer)
    writer._restore_from_disk(own_dir, kind="transactions", buffer=writer._tx_buffer)
    writer._restore_from_disk(own_dir, kind="vertex_builds", buffer=writer._vb_buffer)

    tx_batch = list(writer._tx_buffer)
    vb_batch = list(writer._vb_buffer)
    writer._tx_buffer.clear()
    writer._vb_buffer.clear()
    await writer._flush(tx_batch, vb_batch)

    async with AsyncSession(engine) as session:
        tx_count = await session.scalar(select(func.count()).select_from(TransactionTable))
        vb_count = await session.scalar(select(func.count()).select_from(VertexBuildTable))
    assert tx_count == 1
    assert vb_count == 1


def _find_dead_pid(start: int = 99_999, max_checks: int = 50_000) -> int:
    from langflow.services.telemetry_writer.service import _pid_alive

    pid = start
    for _ in range(max_checks):
        if not _pid_alive(pid):
            return pid
        pid += 1
    pytest.fail("Unable to find an unused PID for orphan outbox setup")
    return 0  # unreachable, for type-checkers


def test_adopt_orphan_outboxes(tmp_path: Path) -> None:
    # Simulate a dead worker that left rows in a sibling outbox.
    dead_pid = _find_dead_pid()
    dead_dir = tmp_path / str(dead_pid)
    dead_dir.mkdir()
    _write_owner_file(dead_dir)
    spill_writer = _build_writer()
    spill_writer._tx_buffer.extend([{"orphan": True}])
    spill_writer._vb_buffer.extend([{"orphan_vb": True}])
    spill_writer._spill_to_disk(dead_dir, kind="transactions", buffer=spill_writer._tx_buffer)
    spill_writer._spill_to_disk(dead_dir, kind="vertex_builds", buffer=spill_writer._vb_buffer)

    own_writer = _build_writer()
    own_writer._adopt_orphan_outboxes(tmp_path, own_pid=1)
    assert list(own_writer._tx_buffer) == [{"orphan": True}]
    assert list(own_writer._vb_buffer) == [{"orphan_vb": True}]
    # Orphan directory was cleaned up.
    assert not dead_dir.exists()


def test_adopt_orphan_outboxes_skips_unknown_host(tmp_path: Path) -> None:
    """Cross-host or pre-reboot orphan dirs must not be adopted.

    Guards against PID reuse on container restart pulling in a stranger's
    spill data — the spill on disk is only adopted when the owner file's
    host + boot identity matches the current host.
    """
    import json as _json

    dead_pid = _find_dead_pid()
    dead_dir = tmp_path / str(dead_pid)
    dead_dir.mkdir()
    # Stamp a mismatched owner file.
    (dead_dir / "owner.json").write_text(
        _json.dumps({"host": "some-other-host", "boot": "unrelated-boot", "pid": dead_pid, "started_at": 0})
    )
    spill_writer = _build_writer()
    spill_writer._tx_buffer.extend([{"orphan": True}])
    spill_writer._spill_to_disk(dead_dir, kind="transactions", buffer=spill_writer._tx_buffer)

    own_writer = _build_writer()
    own_writer._adopt_orphan_outboxes(tmp_path, own_pid=1)
    assert list(own_writer._tx_buffer) == []
    # The orphan directory is left in place so a forensic operator can inspect it.
    assert dead_dir.exists()


def test_adopt_orphan_outboxes_skips_when_owner_missing(tmp_path: Path) -> None:
    """Pre-owner-file directories from older runs must not be silently adopted."""
    dead_pid = _find_dead_pid()
    dead_dir = tmp_path / str(dead_pid)
    dead_dir.mkdir()
    # No owner.json on purpose.
    spill_writer = _build_writer()
    spill_writer._tx_buffer.extend([{"orphan": True}])
    spill_writer._spill_to_disk(dead_dir, kind="transactions", buffer=spill_writer._tx_buffer)

    own_writer = _build_writer()
    own_writer._adopt_orphan_outboxes(tmp_path, own_pid=1)
    assert list(own_writer._tx_buffer) == []


def test_adopt_orphan_outboxes_honors_max_queue(tmp_path: Path) -> None:
    """A pathologically large orphan spill must not OOM the current worker.

    Regression for the case where ``_restore_from_disk`` and
    ``_adopt_orphan_outboxes`` appended directly to the deque, bypassing the
    queue cap.
    """
    dead_pid = _find_dead_pid()
    dead_dir = tmp_path / str(dead_pid)
    dead_dir.mkdir()
    _write_owner_file(dead_dir)
    spill_writer = _build_writer()
    # Spill more rows than the receiving writer's cap.
    spill_writer._tx_buffer.extend([{"i": i} for i in range(100)])
    spill_writer._spill_to_disk(dead_dir, kind="transactions", buffer=spill_writer._tx_buffer)

    own_writer = _build_writer({"telemetry_writer_max_queue": 10})
    own_writer._adopt_orphan_outboxes(tmp_path, own_pid=1)

    # Buffer capped at 10; older rows dropped, drop counter reflects loss.
    assert len(own_writer._tx_buffer) == 10
    assert own_writer.dropped_transactions == 90
    # The newest rows (90..99) survive.
    assert [r["i"] for r in own_writer._tx_buffer] == list(range(90, 100))


def test_prune_stale_foreign_outboxes_deletes_aged_cross_host_dirs(tmp_path: Path) -> None:
    """Cross-host orphan dirs older than max_age must be deleted from a shared volume."""
    import json as _json
    import os as _os
    import time as _time

    foreign_dir = tmp_path / "9999"
    foreign_dir.mkdir()
    owner_file = foreign_dir / "owner.json"
    owner_file.write_text(
        _json.dumps({"host": "some-other-host", "boot": "unrelated-boot", "pid": 9999, "started_at": 0})
    )
    # Backdate the owner file beyond the prune threshold.
    aged = _time.time() - 7200  # 2h ago, default max_age is 1h
    _os.utime(owner_file, (aged, aged))

    writer = _build_writer({"telemetry_writer_orphan_max_age_s": 3600.0})
    writer._prune_stale_foreign_outboxes(tmp_path)
    assert not foreign_dir.exists()


def test_prune_stale_foreign_outboxes_preserves_fresh_cross_host_dirs(tmp_path: Path) -> None:
    """A live foreign pod's outbox must not be deleted while it's actively heartbeating."""
    import json as _json

    foreign_dir = tmp_path / "1234"
    foreign_dir.mkdir()
    (foreign_dir / "owner.json").write_text(
        _json.dumps({"host": "live-foreign-pod", "boot": "unrelated-boot", "pid": 1234, "started_at": 0})
    )
    # Owner file mtime is "now" — pod is actively heartbeating.

    writer = _build_writer({"telemetry_writer_orphan_max_age_s": 3600.0})
    writer._prune_stale_foreign_outboxes(tmp_path)
    assert foreign_dir.exists()


def test_prune_stale_foreign_outboxes_leaves_same_host_dirs_alone(tmp_path: Path) -> None:
    """Same-host orphans are the adoption path's job, not the pruner's — even when aged."""
    import os as _os
    import time as _time

    own_dir = tmp_path / "5555"
    own_dir.mkdir()
    _write_owner_file(own_dir)
    owner_file = own_dir / "owner.json"
    aged = _time.time() - 7200
    _os.utime(owner_file, (aged, aged))

    writer = _build_writer({"telemetry_writer_orphan_max_age_s": 3600.0})
    writer._prune_stale_foreign_outboxes(tmp_path)
    assert own_dir.exists()


def test_heartbeat_owner_file_refreshes_mtime(tmp_path: Path) -> None:
    """The heartbeat must bump the owner file's mtime so foreign hosts can age us out."""
    import os as _os
    import time as _time

    own_dir = tmp_path / "1"
    own_dir.mkdir()
    _write_owner_file(own_dir)
    owner_file = own_dir / "owner.json"
    aged = _time.time() - 7200
    _os.utime(owner_file, (aged, aged))
    before = owner_file.stat().st_mtime

    writer = _build_writer()
    writer._own_outbox_dir = own_dir
    writer._heartbeat_owner_file()
    after = owner_file.stat().st_mtime
    assert after > before


def test_spill_caps_at_max_queue(tmp_path: Path) -> None:
    """A backlogged buffer at shutdown must not spill unbounded rows to disk.

    Older rows beyond ``telemetry_writer_max_queue`` are dropped (matching
    the producer-side overflow policy) and the drop counter is incremented.
    """
    own_dir = tmp_path / "pid"
    own_dir.mkdir()
    writer = _build_writer({"telemetry_writer_max_queue": 20})
    writer._tx_buffer.extend([{"i": i} for i in range(100)])
    writer._spill_to_disk(own_dir, kind="transactions", buffer=writer._tx_buffer)

    assert writer.dropped_transactions == 80
    # Round-trip what hit disk: only the newest 20 rows survive.
    restore = _build_writer({"telemetry_writer_max_queue": 1000})
    restore._restore_from_disk(own_dir, kind="transactions", buffer=restore._tx_buffer)
    assert [r["i"] for r in restore._tx_buffer] == list(range(80, 100))


async def test_sanitization_survives_writer_round_trip(writer_with_engine) -> None:
    """A sensitive value passed through the producer must land redacted in the DB."""
    writer, engine = writer_with_engine
    flow_id = uuid4()
    # TransactionBase.__init__ runs sanitize_data on inputs/outputs. The
    # writer path serializes via model_dump and bulk-inserts; verify the
    # sanitization isn't lost along the way.
    base = TransactionBase(
        vertex_id="v1",
        target_id=None,
        inputs={"api_key": "sk-very-secret-token-12345"},  # pragma: allowlist secret
        outputs={"password": "hunter2"},  # pragma: allowlist secret
        status="success",
        error=None,
        flow_id=flow_id,
    )
    row = TransactionTable(**base.model_dump()).model_dump(mode="python")
    await writer._flush([row], [])

    async with AsyncSession(engine) as session:
        result = await session.execute(select(TransactionTable.inputs, TransactionTable.outputs))
        inputs, outputs = result.one()
    assert "sk-very-secret-token-12345" not in str(inputs)
    assert "hunter2" not in str(outputs)
    assert "api_key" in inputs
    assert "password" in outputs


async def test_retention_failure_preserves_dirty_flows(writer_with_engine) -> None:
    """A sweep that crashes before commit must leave the dirty-flow set intact."""
    writer, _ = writer_with_engine
    tx_flow = uuid4()
    vb_flow = uuid4()
    writer._dirty_tx_flows.add(str(tx_flow))
    writer._dirty_vb_flows.add(str(vb_flow))

    # Wrap the real session_maker so commit raises but exec() works
    # normally — this drives the retention pass through every query and only
    # fails at the commit boundary, which is exactly the scenario the snapshot/
    # restore logic guards against.
    real_session_maker = writer._session_maker

    class _CommitFailsSessionMaker:
        def __call__(self_inner):  # noqa: N805
            session = real_session_maker()

            class _CommitFailsCM:
                async def __aenter__(self):
                    self._session = await session.__aenter__()
                    original_commit = self._session.commit

                    async def _raise_on_commit():
                        await original_commit()  # exercise the path
                        msg = "boom"
                        raise RuntimeError(msg)

                    self._session.commit = _raise_on_commit
                    return self._session

                async def __aexit__(self, *args):
                    return await session.__aexit__(*args)

            return _CommitFailsCM()

    writer._session_maker = _CommitFailsSessionMaker()
    with pytest.raises(RuntimeError, match="boom"):
        await writer._run_retention_pass()

    # Dirty set must still hold both ids since the commit never landed.
    assert str(tx_flow) in writer._dirty_tx_flows
    assert str(vb_flow) in writer._dirty_vb_flows


async def test_in_flight_batch_returned_on_cancel(writer_with_engine) -> None:
    """If the writer is cancelled mid-flush, popped rows must go back to the buffer."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_batch_size = 100
    writer.settings_service.settings.telemetry_writer_flush_interval_s = 0.01
    flow_id = uuid4()
    for _ in range(20):
        writer.enqueue_transaction(_make_transaction_row(flow_id))
    assert len(writer._tx_buffer) == 20

    # Replace the session_maker with one that hangs forever inside exec,
    # so the writer is guaranteed to be blocked mid-flush when we cancel.
    hang = asyncio.Event()  # never set

    class _HangingSessionMaker:
        def __call__(self):
            return self

        async def __aenter__(self):
            return self

        async def __aexit__(self, *_):
            return False

        async def exec(self, *_, **__):
            await hang.wait()

        async def commit(self):
            pass

    writer._session_maker = _HangingSessionMaker()
    task = asyncio.create_task(writer._run_writer())
    # Give the writer a tick to drain into a batch and start the flush.
    for _ in range(20):
        await asyncio.sleep(0.02)
        if len(writer._tx_buffer) == 0:
            break
    assert len(writer._tx_buffer) == 0  # the batch is in-flight

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    # Cancelled mid-flush — the 20 rows should be back in the buffer.
    assert len(writer._tx_buffer) == 20


async def test_lifecycle_idempotent_when_disabled() -> None:
    writer = _build_writer({"telemetry_writer_enabled": False})
    # start() is a no-op when disabled.
    await writer.start()
    assert writer.is_running() is False
    # teardown() on a never-started writer is harmless.
    await writer.teardown()


async def test_teardown_spills_remaining_buffer(tmp_path: Path) -> None:
    outbox = tmp_path / "outbox"
    writer = _build_writer(
        {
            "telemetry_writer_outbox_dir": str(outbox),
        }
    )
    # Manually wire enough state to exercise teardown's spill path without
    # going through start() (which needs the full langflow service stack).
    writer._started = True
    writer._shutdown_event = asyncio.Event()
    writer._shutdown_event.set()
    writer._own_outbox_dir = outbox / "1"
    writer._own_outbox_dir.mkdir(parents=True)
    writer._tx_buffer.append({"to_disk": 1})

    await writer.teardown()

    # A second writer pointed at the same dir should see the spilled row.
    writer2 = _build_writer()
    writer2._restore_from_disk(outbox / "1", kind="transactions", buffer=writer2._tx_buffer)
    assert list(writer2._tx_buffer) == [{"to_disk": 1}]

    # Clean up the outbox directory tree.
    shutil.rmtree(outbox, ignore_errors=True)
    # Also clean up our potential tempfile fallback if used.
    fallback = Path(tempfile.gettempdir()) / "langflow_telemetry_outbox"
    if fallback.exists() and fallback.is_dir() and not any(fallback.iterdir()):
        fallback.rmdir()


async def test_retention_sweep_caps_vertex_builds_per_vertex(writer_with_engine) -> None:
    """Per-vertex cap must prune the oldest builds for a single vertex."""
    writer, engine = writer_with_engine
    writer.settings_service.settings.max_vertex_builds_per_vertex = 3
    writer.settings_service.settings.max_vertex_builds_to_keep = 10_000
    flow_id = uuid4()
    # Insert 8 builds for the same vertex — only 3 should survive.
    vb_batch = [_make_vertex_build_row(flow_id, vertex_id="v1") for _ in range(8)]
    await writer._flush([], vb_batch)
    await writer._run_retention_pass()

    async with AsyncSession(engine) as session:
        count = await session.scalar(select(func.count()).select_from(VertexBuildTable))
    assert count == 3


def test_either_strategy_trips_on_bytes_first(writer_with_engine) -> None:
    """'either' strategy must enforce the byte cap when it trips before the row cap."""
    writer, _ = writer_with_engine
    writer.settings_service.settings.telemetry_writer_size_strategy = "either"
    writer.settings_service.settings.telemetry_writer_max_queue = 10_000  # high — must not trigger
    writer.settings_service.settings.telemetry_writer_max_queue_bytes = 200
    # Each row encodes to ~123 bytes; two rows (246B) exceed the 200B cap.
    for i in range(3):
        writer.enqueue_transaction({"payload": "x" * 100, "i": i})
    # Byte cap trips first — same drop semantics as the bytes-only strategy.
    assert writer.dropped_transactions >= 1
    assert writer._tx_bytes <= 200


async def test_writer_retries_on_batch_failure(writer_with_engine) -> None:
    """Failed flush must increment failed_batches, return rows to buffer, then succeed on retry."""
    writer, engine = writer_with_engine
    writer.settings_service.settings.telemetry_writer_batch_size = 100
    writer.settings_service.settings.telemetry_writer_flush_interval_s = 0.01
    flow_id = uuid4()
    for _ in range(5):
        writer.enqueue_transaction(_make_transaction_row(flow_id))

    fail_count = 0

    real_flush = writer._flush.__func__

    async def _fail_twice(self, tx_batch, vb_batch):
        nonlocal fail_count
        if fail_count < 2:
            fail_count += 1
            msg = "injected failure"
            raise RuntimeError(msg)
        return await real_flush(self, tx_batch, vb_batch)

    import types

    writer._flush = types.MethodType(_fail_twice, writer)

    task = asyncio.create_task(writer._run_writer())
    # The two injected failures each return the batch to the buffer and back off
    # via _wait_or_shutdown. Setting the shutdown event short-circuits those
    # backoffs so the writer reaches its successful retry and drains — the
    # shutdown, not wall-clock waiting, is what lands the rows.
    writer._shutdown_event.set()
    await asyncio.wait_for(task, timeout=5)

    assert writer.failed_batches == 2
    assert writer.flushed_rows == 5
    async with AsyncSession(engine) as session:
        final = await session.scalar(select(func.count()).select_from(TransactionTable))
    assert final == 5


async def test_writer_escalates_after_threshold_failures(writer_with_engine) -> None:
    """Consecutive failures past the threshold must keep retrying (exercises the `>=` branch, not `==`)."""
    import types

    writer, engine = writer_with_engine
    writer.settings_service.settings.telemetry_writer_batch_size = 100
    writer.settings_service.settings.telemetry_writer_flush_interval_s = 0.01
    flow_id = uuid4()
    for _ in range(5):
        writer.enqueue_transaction(_make_transaction_row(flow_id))

    # Fail one more than the threshold so the escalation branch runs at counts
    # both equal to and greater than the threshold — exactly what `>=` covers
    # and `==` would not.
    fail_target = _FAILURE_ESCALATION_THRESHOLD + 1
    fail_count = 0
    real_flush = writer._flush.__func__

    async def _fail_n_times(self, tx_batch, vb_batch):
        nonlocal fail_count
        if fail_count < fail_target:
            fail_count += 1
            msg = "injected failure"
            raise RuntimeError(msg)
        return await real_flush(self, tx_batch, vb_batch)

    writer._flush = types.MethodType(_fail_n_times, writer)

    task = asyncio.create_task(writer._run_writer())
    # Short-circuit every backoff so the failures (and the final success) run fast.
    writer._shutdown_event.set()
    await asyncio.wait_for(task, timeout=5)

    # Every injected failure was counted (driving consecutive_failures past the
    # threshold), then the buffer drained on the success.
    assert writer.failed_batches == fail_target
    assert writer.flushed_rows == 5
    async with AsyncSession(engine) as session:
        final = await session.scalar(select(func.count()).select_from(TransactionTable))
    assert final == 5


async def test_teardown_cancelled_still_spills_buffer(writer_with_engine, tmp_path: Path) -> None:
    """If teardown() is itself cancelled, the in-memory buffer must still spill to disk."""
    writer, _ = writer_with_engine
    own_dir = tmp_path / "pid"
    own_dir.mkdir()
    writer._own_outbox_dir = own_dir
    # Long drain so teardown blocks in the writer-drain await while we cancel it.
    writer.settings_service.settings.telemetry_writer_shutdown_drain_s = 30.0

    for i in range(3):
        writer.enqueue_transaction({"cancelled_spill": i})
    assert len(writer._tx_buffer) == 3

    # A writer task that never finishes on its own, so teardown stays parked in
    # the drain await until we cancel it.
    async def _never() -> None:
        await asyncio.Event().wait()

    writer._writer_task = asyncio.create_task(_never())

    teardown_task = asyncio.create_task(writer.teardown())
    await asyncio.sleep(0.05)  # let teardown reach the drain await
    teardown_task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await teardown_task

    # The cancellation must not have dropped the buffer: a fresh reader pointed
    # at the same outbox sees the spilled rows.
    reader = _build_writer()
    reader._restore_from_disk(own_dir, kind="transactions", buffer=reader._tx_buffer)
    assert [row["cancelled_spill"] for row in reader._tx_buffer] == [0, 1, 2]