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.
"""

from __future__ import annotations

import asyncio
import contextlib
import socket
import struct
import time
from pathlib import Path

import pytest
import uvicorn

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 ..sandbox.test_isolation import USERSNS_AVAILABLE


def _userns_probe() -> None:
    if not USERSNS_AVAILABLE:
        pytest.skip("user namespaces unavailable on this host (probe)")


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


class KillableProxy:
    """TCP forwarder whose data path dies on ``kill()`` — the listener stays.

    Accepts are preserved so the agent's reconnects fail fast (connection
    reset) instead of hanging on a black-holed socket, keeping the outage
    window deterministic. ``revive()`` restores forwarding.
    """

    def __init__(self, target_port: int) -> None:
        self.target_port = target_port
        self.listen_port = _free_port()
        self._listener: socket.socket | None = None
        self._killed = False
        self._conns: list[socket.socket] = []

    async def start(self) -> None:
        self._listener = socket.socket()
        self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._listener.bind(("127.0.0.1", self.listen_port))
        self._listener.listen(16)
        self._listener.setblocking(False)
        loop = asyncio.get_running_loop()
        loop.create_task(self._accept_loop())

    async def _accept_loop(self) -> None:
        loop = asyncio.get_running_loop()
        assert self._listener is not None
        while True:
            try:
                client, _ = await loop.sock_accept(self._listener)
            except OSError:
                return  # listener closed
            client.setblocking(False)  # accepted sockets default to blocking
            if self._killed:
                with contextlib.suppress(OSError):
                    # RST on accept: reconnects fail fast while "down"
                    client.setsockopt(
                        socket.SOL_SOCKET, socket.SO_LINGER,
                        struct.pack("<ii", 1, 0),
                    )
                    client.close()
                continue
            upstream = socket.socket()
            upstream.setblocking(False)
            try:
                await loop.sock_connect(upstream, ("127.0.0.1", self.target_port))
            except OSError:
                with contextlib.suppress(OSError):
                    client.close()
                continue
            self._conns.extend((client, upstream))
            loop.create_task(self._pump(client, upstream))
            loop.create_task(self._pump(upstream, client))

    async def _pump(self, src: socket.socket, dst: socket.socket) -> None:
        """Forward until EOF, error, or the killed flag is observed.

        Only ``dst`` is closed here — the reverse-direction pump owns ``src``
        (closing the peer's socket from this task causes EBADF cascades).
        """
        loop = asyncio.get_running_loop()
        try:
            while True:
                data = await loop.sock_recv(src, 65536)
                if not data:
                    return
                if self._killed:
                    return  # outage: drop the relay, peer sees EOF on dst close
                await loop.sock_sendall(dst, data)
        except OSError:
            return
        finally:
            with contextlib.suppress(OSError):
                dst.close()

    def kill(self) -> None:
        """Sever the data path — flag only.

        The pumps and the accept loop observe ``_killed`` themselves and close
        THEIR OWN sockets from inside the event loop. Closing sockets that the
        loop is currently awaiting from the outside deadlocks the loop's
        child-watcher in this environment (observed: post-kill subprocess
        spawns hung forever) — so no external socket surgery, ever.
        """

    def revive(self) -> None:
        self._killed = False


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)


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

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

    store = SQLiteTraceStore(db_path=tmp_path / "durability.db")
    app = create_app()
    app.state.trace_store = store
    app.state.trace_integrity = TraceIntegrityMap()

    server_port = _free_port()
    server = uvicorn.Server(
        uvicorn.Config(app, host="127.0.0.1", port=server_port, log_level="warning")
    )
    serve_task = asyncio.get_running_loop().create_task(server.serve())

    proxy = KillableProxy(target_port=server_port)
    await proxy.start()

    manager = SandboxManager(backend=UnshareBackend(), settings=Settings())
    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.
            "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}"
            ),
            # Fast reconnect so the probe stays bounded.
            "NC_BACKOFF_BASE_S": "0.1",
            "NC_BACKOFF_MAX_S": "0.5",
        }

    manager._capture_env = capture_env  # noqa: SLF001 - test seam

    try:
        for _ in range(100):
            if server.started:
                break
            await asyncio.sleep(0.1)
        assert server.started, "uvicorn did not start"

        handle = await manager.create(learner, task_id=task)
        try:
            live = manager._handles[handle.id]  # noqa: SLF001 - test seam
            backend: UnshareBackend = manager._backend  # noqa: SLF001 - test seam

            # 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.
            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).
            proxy.revive()
            await asyncio.sleep(1.5)
            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:
        server.should_exit = True
        with contextlib.suppress(Exception):
            await asyncio.wait_for(serve_task, timeout=10.0)
        store.close()