nextcraft/apps/ai-service/tests/telemetry/test_durability.py

"""Dropped-connection durability probe (Task 2-4-01, REQ-3-003).

At-least-once delivery + server-side (learner,task,seq) dedup = every event
stored exactly once, in order — proven against a REAL uvicorn ingest and a
REAL namespace sandbox, with the connection severed mid-stream by a killable
TCP proxy (deterministic "network down" window).

Probe-guarded: skips (not fails) on hosts without user namespaces.

P2 (Task 2-1-04, deterministic waits) — per-site decisions:
  * MOVED: the post-revive fixed settles (was ``asyncio.sleep(1.5)`` /
    ``asyncio.sleep(1.0)``) → ``_await_file_diffs`` deadline-bounded polls of
    the durable store. A genuine hook exists: mid-outage file_diffs are only
    durable after reconnect+flush, so their presence IS the condition the
    sleeps settled blindly.
  * STAYED: the 0.3s inter-write pacing inside the outage window (Phase B of
    test_disconnect_reconnect_loses_nothing). By construction nothing is
    observable while the link is severed — the pacing shapes the scenario
    (distinct watcher poll ticks per write), it does not wait for a condition.
  * NOT TOUCHED: test_sandbox_agent.py's poll loop (G-102 — already
    event-driven, no fixed sleep exists there).
"""

from __future__ import annotations

import asyncio
import time
from pathlib import Path

import pytest

from ai_service.config import Settings
from ai_service.main import create_app
from ai_service.sandbox import SandboxManager
from ai_service.sandbox.unshare_backend import UnshareBackend
from ai_service.telemetry.ingest import TraceIntegrityMap
from ai_service.telemetry.store import SQLiteTraceStore
from ai_service.telemetry.tokens import IngestToken
from tests.helpers import AsyncKillableProxy, free_port, serve, skip_if_no_userns

# P1 tier marker (Task 1-1-02): both tests boot a REAL uvicorn ingest +
# REAL namespace sandboxes behind a killable TCP proxy — real_server
# tier, 180s timeout via the conftest collection hook.
pytestmark = [pytest.mark.real_server]


async def _await_events(
    store: SQLiteTraceStore, learner: str, task: str, *, minimum: int, timeout_s: float = 30.0
) -> list:
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        events = store.get_trace(learner, task)
        if len(events) >= minimum:
            return events
        await asyncio.sleep(0.25)
    return store.get_trace(learner, task)


async def _await_file_diffs(
    store: SQLiteTraceStore,
    learner: str,
    task: str,
    paths: set[str],
    timeout_s: float = 30.0,
) -> None:
    """P2 (Task 2-1-04): poll the DURABLE store until every ``paths``
    file_diff has landed — the event-driven replacement for the fixed
    post-revive sleeps this file used to run.

    Genuine hook: mid-outage writes are emitted while the connection is
    severed, so their file_diff events can only become durable after the
    agent reconnected AND flushed its spool — the exact condition the old
    ``asyncio.sleep(1.5)`` / ``asyncio.sleep(1.0)`` settled blindly. Raises on
    deadline (the ``_wait_until`` model: fail at the wait site, not 30s later
    in the final assertions).
    """
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        stored = _stored_diff_paths(store, learner, task)
        if paths <= stored:
            return
        await asyncio.sleep(0.25)
    stored = _stored_diff_paths(store, learner, task)
    assert paths <= stored, (
        f"spool flush incomplete {timeout_s}s after revive; "
        f"missing file_diffs: {sorted(paths - stored)}"
    )


def _stored_diff_paths(store: SQLiteTraceStore, learner: str, task: str) -> set[str]:
    paths: set[str] = set()
    for e in store.get_trace(learner, task):
        if e.kind != "file_diff":
            continue
        path = e.payload.get("path")
        if isinstance(path, str):  # P5: payload.get is object|None — filter to str
            paths.add(path)
    return paths


@pytest.mark.asyncio
async def test_disconnect_reconnect_loses_nothing(tmp_path: Path, sandbox_dir: Path) -> None:
    """Sever the agent's WS mid-stream; every event lands exactly once, ordered."""
    skip_if_no_userns()

    learner = "durability-learner"
    task = "task-durability"

    store = SQLiteTraceStore(db_path=tmp_path / "durability.db")
    # Hermetic (v0.3.6): pin state dirs; sandbox workdirs go to the
    # bind-mount-safe sandbox_dir fixture (overlayfs /tmp breaks userns binds).
    app = create_app(
        Settings(
            provider="mock",
            db_path=tmp_path / "durability-app.db",
            sandbox_dir=sandbox_dir,
        )
    )
    app.state.trace_store = store
    app.state.trace_integrity = TraceIntegrityMap()
    # v0.4.7: pre-set BEFORE the lifespan (state-injection) so the manager
    # and the endpoint share one instance from the first dial.
    app.state.ingest_token = IngestToken("durability-secret")

    server_port = free_port()
    proxy = AsyncKillableProxy(target_port=server_port)
    await proxy.start()

    manager = SandboxManager(
        backend=UnshareBackend(),
        settings=Settings(sandbox_dir=sandbox_dir),
        ingest_token=app.state.ingest_token,
    )
    app.state.sandbox_manager = manager

    def capture_env(sandbox_id: str, learner_id: str, task_id: str) -> dict[str, str]:
        return {
            "NC_LEARNER_ID": learner_id,
            "NC_TASK_ID": task_id,
            "NC_SANDBOX_ID": sandbox_id,
            # Agent dials the PROXY; the proxy forwards to the real server.
            # v0.4.7: URL carries the app's minted ingest token.
            "NC_INGEST_URL": (
                f"ws://127.0.0.1:{proxy.listen_port}/v1/telemetry/ingest"
                f"?learner_id={learner_id}&task_id={task_id}&sandbox_id={sandbox_id}"
                f"&token={app.state.ingest_token.mint(learner_id, task_id, sandbox_id)}"
            ),
            # Fast reconnect so the probe stays bounded.
            "NC_BACKOFF_BASE_S": "0.1",
            "NC_BACKOFF_MAX_S": "0.5",
        }

    # P5 scoped: deliberate test seam — the proxy URL replaces the real
    # capture agent's env (the durability probe dials the KillableProxy).
    manager._capture_env = capture_env  # type: ignore[method-assign]

    try:
        async with serve(app, "127.0.0.1", server_port):
            handle = await manager.create(learner, task_id=task)
            try:
                live = manager._handles[handle.id]
                backend = manager._backend

                # Phase A — connected: a write streams through.
                result = await backend.exec(live, ["sh", "-c", "echo one > a.txt"])
                assert result.returncode == 0, result.stderr
                events = await _await_events(store, learner, task, minimum=1)
                assert events, "no events arrived before the outage"

                # Phase B — network down: sever mid-stream, keep generating.
                # The 0.3s pacing between writes is scenario shaping INSIDE a
                # hard outage: by construction nothing is observable while the
                # link is down (that is the point), and the spacing keeps the
                # agent's watcher from coalescing the three files into one
                # poll tick. P2 (Task 2-1-04): stays a fixed sleep — no
                # condition hook exists inside a severed connection.
                proxy.kill()
                for n in ("two", "three", "four"):
                    result = await backend.exec(live, ["sh", "-c", f"echo {n} > {n}.txt"])
                    assert result.returncode == 0, result.stderr
                    await asyncio.sleep(0.3)

                # Phase C — revive: the agent reconnects (fast backoff) and
                # flushes the spool. Server dedups on (learner,task,seq).
                # P2 (Task 2-1-04): the old fixed 1.5s settle is now a
                # deadline-bounded poll of the DURABLE store — the Phase-B
                # file_diffs are only stored after reconnect+flush, so their
                # presence is the observable "spool flushed" condition.
                proxy.revive()
                await _await_file_diffs(store, learner, task, {"two.txt", "three.txt", "four.txt"})
                for n in ("five", "six"):
                    result = await backend.exec(live, ["sh", "-c", f"echo {n} > {n}.txt"])
                    assert result.returncode == 0, result.stderr

                events = await _await_events(
                    store, learner, task, minimum=6, timeout_s=30.0
                )
                seqs = [e.seq for e in events]
                assert seqs == sorted(seqs), f"out of order after reconnect: {seqs}"
                assert len(set(seqs)) == len(seqs), f"duplicates stored: {seqs}"
                assert len(events) >= 6, f"lost events across the outage: {seqs}"
                # Exactly-once storage despite at-least-once delivery: every spool
                # flush may re-send the replay-margin line; dedup collapses it.
                stored_files = {
                    e.payload.get("path") for e in events if e.kind == "file_diff"
                }
                assert stored_files, "file_diff events missing from the stored trace"
            finally:
                await manager.destroy(handle.id)
    finally:
        await proxy.stop()  # P3: close relay/listener — no socket left to the GC
        store.close()


@pytest.mark.asyncio
async def test_midburst_disconnect_loses_nothing(tmp_path: Path, sandbox_dir: Path) -> None:
    """REQ-5-007 / D-045 / MH-1d: sever MID-BURST — immediately after a rapid
    multi-frame send, WITHOUT waiting for server observation of the burst —
    then revive and require every emitted seq stored exactly once, in order.

    This is the exact scenario the P07 de-flake documented as uncovered by the
    one-line replay margin (kill-timing is racy, but the OUTCOME is invariant
    under the seq-ack protocol: all interleavings converge to exactly-once
    via spool replay + server dedup + ack-based trimming)."""
    skip_if_no_userns()

    learner = "midburst-learner"
    task = "task-midburst"

    store = SQLiteTraceStore(db_path=tmp_path / "midburst.db")
    app = create_app(
        Settings(provider="mock", db_path=tmp_path / "midburst-app.db", sandbox_dir=sandbox_dir)
    )
    app.state.trace_store = store
    app.state.trace_integrity = TraceIntegrityMap()
    # v0.4.7: pre-set BEFORE the lifespan (state-injection) so the manager
    # and the endpoint share one instance from the first dial.
    app.state.ingest_token = IngestToken("durability-secret")

    server_port = free_port()
    proxy = AsyncKillableProxy(target_port=server_port)
    await proxy.start()

    manager = SandboxManager(
        backend=UnshareBackend(),
        settings=Settings(sandbox_dir=sandbox_dir),
        ingest_token=app.state.ingest_token,
    )
    app.state.sandbox_manager = manager

    def capture_env(sandbox_id: str, learner_id: str, task_id: str) -> dict[str, str]:
        return {
            "NC_LEARNER_ID": learner_id,
            "NC_TASK_ID": task_id,
            "NC_SANDBOX_ID": sandbox_id,
            # v0.4.7: URL carries the app's minted ingest token.
            "NC_INGEST_URL": (
                f"ws://127.0.0.1:{proxy.listen_port}/v1/telemetry/ingest"
                f"?learner_id={learner_id}&task_id={task_id}&sandbox_id={sandbox_id}"
                f"&token={app.state.ingest_token.mint(learner_id, task_id, sandbox_id)}"
            ),
            "NC_BACKOFF_BASE_S": "0.1",
            "NC_BACKOFF_MAX_S": "0.5",
        }

    # P5 scoped: deliberate test seam — the proxy URL replaces the real
    # capture agent's env (the durability probe dials the KillableProxy).
    manager._capture_env = capture_env  # type: ignore[method-assign]

    try:
        async with serve(app, "127.0.0.1", server_port):
            handle = await manager.create(learner, task_id=task)
            try:
                live = manager._handles[handle.id]
                backend = manager._backend

                # Phase A — connected: confirm the pipe works (1 event minimum).
                result = await backend.exec(live, ["sh", "-c", "echo warm > warm.txt"])
                assert result.returncode == 0, result.stderr
                events = await _await_events(store, learner, task, minimum=1)
                assert events, "warm-up event never arrived"

                # Phase B — MID-BURST: rapid-fire writes with NO wait between
                # them, then kill IMMEDIATELY (before any server observation).
                # Frames are in TCP flight when the link dies — the one-line
                # replay margin loses 1..N-1 of them pre-D-045.
                proxy.kill()
                burst = [f"burst-{n}" for n in range(10)]
                for name in burst:
                    result = await backend.exec(live, ["sh", "-c", f"echo {name} > {name}.txt"])
                    assert result.returncode == 0, result.stderr

                # Phase C — revive: agent reconnects (fast backoff), flushes the
                # spool, acks trim it; server dedups any replayed margin.
                # P2 (Task 2-1-04): the old fixed 1.0s settle is now a
                # deadline-bounded poll of the DURABLE store — the mid-burst
                # file_diffs were emitted into the severed link and only
                # become durable after reconnect+flush, so their presence is
                # the observable "spool flushed" condition.
                proxy.revive()
                await _await_file_diffs(
                    store, learner, task, {f"{name}.txt" for name in burst}
                )
                result = await backend.exec(live, ["sh", "-c", "echo after > after.txt"])
                assert result.returncode == 0, result.stderr

                events = await _await_events(store, learner, task, minimum=12, timeout_s=30.0)
                seqs = [e.seq for e in events]
                assert seqs == sorted(seqs), f"out of order after mid-burst reconnect: {seqs}"
                assert len(set(seqs)) == len(seqs), f"duplicates stored: {seqs}"
                # Contiguity: warm-up + 10 burst writes + watcher events + the
                # after-write — no GAPS allowed (the ack protocol must deliver
                # every in-flight frame via replay).
                assert seqs == list(range(seqs[0], seqs[-1] + 1)), (
                    f"gaps in seq chain after mid-burst kill: {seqs}"
                )
                stored_files = {
                    e.payload.get("path") for e in events if e.kind == "file_diff"
                }
                assert "after.txt" in stored_files, "post-revive write missing from the trace"
            finally:
                await manager.destroy(handle.id)
    finally:
        await proxy.stop()  # P3: close relay/listener — no socket left to the GC
        store.close()