LangBot/tests/unit_tests/test_operation_trace_tenant_scope.py

"""Regression tests for the operation-trace tenant scope.

Cloud runtime refuses every database call made without an explicit scope, and the
audit table is row-level-secured on the bound Workspace. The writer drains its
queue outside any request lifetime, so these tests drive the real service through
persistence doubles that enforce exactly that contract: without a bound scope the
call fails. A writer (or a retention pass) that forgets to bind a scope records
nothing, which is the defect pinned here.
"""

from __future__ import annotations

import asyncio
import contextlib
import contextvars
import datetime
import logging
import typing
from types import SimpleNamespace

import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine

from langbot.pkg.api.http.authz import WorkspaceRole
from langbot.pkg.api.http.context import PrincipalType
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence import metadata as metadata_module
from langbot.pkg.entity.persistence import operation_log as operation_log_module
from langbot.pkg.operation_trace import service as service_module
from langbot.pkg.persistence.tenant_uow import CrossScopeTransactionError, TenantScopeRequiredError

pytestmark = pytest.mark.asyncio


@pytest.fixture(autouse=True)
def _closed_gate(monkeypatch):
    """The gate latch is process-wide: start every test from a disabled instance."""

    monkeypatch.setattr(service_module, '_global_enabled', False)


TRACED_WORKSPACE = '00000000-0000-0000-0000-00000000000a'
IDLE_WORKSPACE = '00000000-0000-0000-0000-00000000000b'
MODEL = operation_log_module.WorkspaceOperationLog
LEVEL_KEY = service_module.OPERATION_LEVEL_KEY


class _CloudRuntimePersistence:
    """Persistence double mirroring the cloud-runtime scope contract."""

    def __init__(self, engine) -> None:
        self.engine = engine
        self._bound: contextvars.ContextVar[tuple[str, typing.Any] | None] = contextvars.ContextVar(
            'bound_workspace', default=None
        )

    def current_scope(self) -> str | None:
        bound = self._bound.get()
        if bound is None:
            return None
        workspace_uuid, owner_task = bound
        if owner_task is not asyncio.current_task():
            # The real manager refuses a boundary a child task inherited.
            raise CrossScopeTransactionError('scoped boundaries cannot be inherited by child tasks')
        return workspace_uuid

    @contextlib.asynccontextmanager
    async def tenant_scope(self, workspace_uuid: str):
        token = self._bound.set((workspace_uuid, asyncio.current_task()))
        try:
            yield self
        finally:
            self._bound.reset(token)

    async def execute_async(self, *args, **kwargs):
        if self.current_scope() is None:
            raise TenantScopeRequiredError('cloud persistence access requires an explicit Workspace scope')
        # The database carries the scope as an RLS setting, so a statement that
        # writes another Workspace's rows would be rejected as well.
        async with self.engine.connect() as connection:
            result = await connection.execute(*args, **kwargs)
            await connection.commit()
            return result


class _SingleDatabasePersistence:
    """Persistence double for the unscoped single-database deployment."""

    def __init__(self, engine) -> None:
        self.engine = engine

    async def execute_async(self, *args, **kwargs):
        async with self.engine.connect() as connection:
            result = await connection.execute(*args, **kwargs)
            await connection.commit()
            return result


@pytest.fixture
async def trace_env(tmp_path):
    """A real service against SQLite, with the enforcement double installed."""

    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "trace-scope.db"}')
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.create_all)
    persistence = _CloudRuntimePersistence(engine)
    application = SimpleNamespace(
        persistence_mgr=persistence,
        logger=logging.getLogger('operation-trace-scope-test'),
        user_service=None,
        operation_trace_active=False,
    )
    try:
        yield service_module.WorkspaceSettingsService(application), application, engine
    finally:
        await engine.dispose()


async def _set_level(engine, workspace_uuid: str, level: int) -> None:
    async with engine.begin() as connection:
        await connection.execute(
            sqlalchemy.insert(metadata_module.WorkspaceMetadata).values(
                workspace_uuid=workspace_uuid,
                key=LEVEL_KEY,
                value=str(level),
            )
        )


async def _rows(engine) -> list:
    async with engine.begin() as connection:
        result = await connection.execute(sqlalchemy.select(MODEL))
    return list(result.all())


async def test_writer_binds_workspace_scope_when_recording(trace_env):
    """A queued trace is written by the background writer holding its own scope."""

    service, application, engine = trace_env
    rule = service_module.classify('POST', '/api/v1/pipelines')

    queued = await service.record(
        TRACED_WORKSPACE,
        rule=rule,
        level=service_module.OPERATION_LEVEL_MUTATION,
        http_method='POST',
        route='/api/v1/pipelines',
    )
    assert queued is True

    await service.flush_pending(timeout=5.0)

    rows = await _rows(engine)
    assert len(rows) == 1
    assert rows[0].workspace_uuid == TRACED_WORKSPACE
    assert service._write_failures == 0


async def test_retention_prune_binds_workspace_scope(trace_env):
    """The maintenance loop prunes without a request, so prune binds its own scope."""

    service, application, engine = trace_env
    stale = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) - datetime.timedelta(days=90)
    async with engine.begin() as connection:
        await connection.execute(
            sqlalchemy.insert(MODEL).values(
                workspace_uuid=TRACED_WORKSPACE,
                level=service_module.OPERATION_LEVEL_MUTATION,
                outcome='ok',
                created_at=stale,
            )
        )
    await _set_level(engine, TRACED_WORKSPACE, service_module.OPERATION_LEVEL_MUTATION)

    report = await service.prune(TRACED_WORKSPACE, retention_days=30)

    assert report['expired'] == 1
    assert await _rows(engine) == []


async def test_request_wrapper_records_after_the_handler_released_its_scope(trace_env):
    """The route wrapper traces outside the request scope and must still record.

    Level resolution and the queued write both happen after the handler released
    its tenant scope. Each of them has to bind the Workspace itself: an unscoped
    read reported level 0, so every request was dropped without a trace.
    """

    service, application, engine = trace_env
    await _set_level(engine, TRACED_WORKSPACE, service_module.OPERATION_LEVEL_MUTATION)
    context = SimpleNamespace(
        workspace_uuid=TRACED_WORKSPACE,
        workspace=SimpleNamespace(role=WorkspaceRole.OWNER.value),
        principal=SimpleNamespace(
            principal_type=PrincipalType.ACCOUNT,
            account_uuid=None,
            actor_account_uuid=None,
            api_key_uuid=None,
        ),
        auth_type='user_token',
        request_id='req-scope-test',
    )

    recorded = await service.record_request(
        context,
        method='POST',
        route=f'/api/v1/workspaces/{TRACED_WORKSPACE}/invitations',
        status_code=200,
    )

    assert recorded is True
    await service.flush_pending(timeout=5.0)
    rows = await _rows(engine)
    assert len(rows) == 1
    assert rows[0].action == 'member_invite'


async def test_writer_started_inside_a_request_scope_still_records(trace_env):
    """The writer outlives the request that first enqueued a trace.

    A task created while its caller held a Workspace boundary inherits it, and a
    child task may not use an inherited boundary. Starting the writer with the
    caller's context made it reject every later write instead of recording it.
    """

    service, application, engine = trace_env
    await _set_level(engine, TRACED_WORKSPACE, service_module.OPERATION_LEVEL_MUTATION)
    rule = service_module.classify('POST', '/api/v1/pipelines')

    async with application.persistence_mgr.tenant_scope(TRACED_WORKSPACE):
        queued = await service.record(
            TRACED_WORKSPACE,
            rule=rule,
            level=service_module.OPERATION_LEVEL_MUTATION,
            http_method='POST',
            route='/api/v1/pipelines',
        )
        assert queued is True
        await service.flush_pending(timeout=5.0)

    rows = await _rows(engine)
    assert len(rows) == 1
    assert service._write_failures == 0


async def test_prime_opens_gate_from_discovered_workspace_bindings(trace_env):
    """A restart resumes tracing: the gate opens from the Workspaces this instance owns."""

    service, application, engine = trace_env
    await _set_level(engine, IDLE_WORKSPACE, service_module.OPERATION_LEVEL_NONE)
    await _set_level(engine, TRACED_WORKSPACE, service_module.OPERATION_LEVEL_MUTATION)

    async def list_active_execution_bindings():
        return [
            SimpleNamespace(workspace_uuid=IDLE_WORKSPACE),
            SimpleNamespace(workspace_uuid=TRACED_WORKSPACE),
        ]

    application.workspace_service = SimpleNamespace(list_active_execution_bindings=list_active_execution_bindings)

    await service.prime_global_flag()

    assert application.operation_trace_active is True


async def test_prime_keeps_gate_closed_when_no_workspace_traces(trace_env):
    """A disabled instance still pays nothing on the request path."""

    service, application, engine = trace_env
    await _set_level(engine, IDLE_WORKSPACE, service_module.OPERATION_LEVEL_NONE)

    async def list_active_execution_bindings():
        return [SimpleNamespace(workspace_uuid=IDLE_WORKSPACE)]

    application.workspace_service = SimpleNamespace(list_active_execution_bindings=list_active_execution_bindings)

    await service.prime_global_flag()

    assert application.operation_trace_active is False


async def test_prime_falls_back_to_metadata_scan_without_discovery(tmp_path):
    """A single-database deployment without a Workspace directory still resumes."""

    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "trace-single.db"}')
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.create_all)
    application = SimpleNamespace(
        persistence_mgr=_SingleDatabasePersistence(engine),
        logger=logging.getLogger('operation-trace-single-test'),
        user_service=None,
        operation_trace_active=False,
    )
    try:
        await _set_level(engine, TRACED_WORKSPACE, service_module.OPERATION_LEVEL_MUTATION)
        service = service_module.WorkspaceSettingsService(application)

        await service.prime_global_flag()

        assert application.operation_trace_active is True
    finally:
        await engine.dispose()