OpenOPC/tests/test_office_shutdown_lifecycle.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock

from opc.core.active_task_runs import ActiveTaskRunRegistry
from opc.core.models import ExecutionCheckpoint, Task
from opc.plugins.office_ui.ws_handler import WSHandler


def _make_handler(root_engine: SimpleNamespace | None = None, **overrides: object) -> WSHandler:
    """Build a real WSHandler over stub dependencies.

    Constructed through ``__init__`` (rather than ``__new__`` plus a
    hand-copied attribute list) so every attribute ``shutdown()`` touches is
    initialized by the constructor itself; each test only overrides what its
    scenario controls. ``__init__`` is side-effect free for stub inputs: it
    assigns state, builds Dispatcher/OfficeServices holders, and wires engine
    callbacks via defensive setattr/getattr.
    """
    engine = root_engine if root_engine is not None else SimpleNamespace()
    if not hasattr(engine, "project_id"):
        engine.project_id = "default"
    handler = WSHandler(
        engine=engine,
        agent_store=SimpleNamespace(),
        chat_store=None,
        event_adapter=SimpleNamespace(),
    )
    for name, value in overrides.items():
        setattr(handler, name, value)
    return handler


def test_ws_shutdown_checkpoints_before_cancelling_and_awaiting_sessions() -> None:
    async def scenario() -> None:
        events: list[str] = []
        started = asyncio.Event()

        async def execution() -> None:
            started.set()
            try:
                await asyncio.Event().wait()
            finally:
                events.append("execution_finally")

        async def prepare() -> list[dict]:
            events.append("checkpoint")
            assert not execution_task.done()
            return []

        execution_task = asyncio.create_task(execution())
        await started.wait()

        handler = _make_handler(
            SimpleNamespace(prepare_active_company_runtimes_for_shutdown=prepare),
            _background_tasks={execution_task},
            _task_bg_context={execution_task: {"task_id": "runtime-task"}},
            _task_bg_map={"runtime-task": {execution_task}},
        )

        await handler.shutdown(timeout=1.0)

        assert events == ["checkpoint", "execution_finally"]
        assert execution_task.done()

    asyncio.run(scenario())


def test_ws_shutdown_checkpoint_failure_does_not_cancel_execution_or_close_the_gap() -> None:
    async def scenario() -> None:
        released = asyncio.Event()

        async def execution() -> None:
            try:
                await asyncio.Event().wait()
            finally:
                released.set()

        execution_task = asyncio.create_task(execution())
        await asyncio.sleep(0)

        handler = _make_handler(
            SimpleNamespace(
                prepare_active_company_runtimes_for_shutdown=AsyncMock(
                    side_effect=RuntimeError("checkpoint unavailable")
                ),
            ),
            _background_tasks={execution_task},
            _task_bg_context={execution_task: {"task_id": "runtime-task"}},
            _task_bg_map={"runtime-task": {execution_task}},
        )

        try:
            await handler.shutdown(timeout=1.0)
        except RuntimeError as exc:
            assert str(exc) == "checkpoint unavailable"
        else:
            raise AssertionError("shutdown must fail closed when checkpointing fails")

        assert not released.is_set()
        assert not execution_task.done()
        execution_task.cancel()
        await asyncio.gather(execution_task, return_exceptions=True)

    asyncio.run(scenario())


def test_ws_shutdown_rejects_background_work_scheduled_by_late_ingress() -> None:
    async def scenario() -> None:
        entered = False

        async def late_work() -> None:
            nonlocal entered
            entered = True

        handler = _make_handler(_shutting_down=True)
        task = handler._track(late_work())
        await asyncio.gather(task, return_exceptions=True)
        await asyncio.sleep(0)

        assert task.cancelled()
        assert entered is False
        assert task not in handler._background_tasks

    asyncio.run(scenario())


def test_ws_shutdown_drains_queued_duplicate_handoff_before_checkpointing() -> None:
    async def scenario() -> None:
        registry = ActiveTaskRunRegistry()
        runtime_lock = asyncio.Lock()
        execution_registered = asyncio.Event()
        prepare_called = asyncio.Event()
        execution_released = asyncio.Event()

        async def prepare() -> list[dict]:
            assert registry.is_active("project-a", "runtime-task")
            assert registry.pending_handoff_count == 0
            prepare_called.set()
            return []

        root_engine = SimpleNamespace(
            _active_task_run_registry=registry,
            prepare_active_company_runtimes_for_shutdown=prepare,
        )
        handler = _make_handler(root_engine)

        async def execution() -> None:
            async with runtime_lock:
                attempt_token = registry.register("project-a", "runtime-task")
                execution_registered.set()
                try:
                    await asyncio.Event().wait()
                finally:
                    registry.unregister("project-a", "runtime-task", attempt_token)
                    execution_released.set()

        async def queued_duplicate() -> None:
            async with runtime_lock:
                attempt_token = registry.register("project-a", "runtime-task")
                try:
                    await asyncio.Event().wait()
                finally:
                    registry.unregister("project-a", "runtime-task", attempt_token)

        first_handoff = registry.reserve_handoff()
        with registry.bind_handoff(first_handoff):
            first = handler._track_session(
                "runtime-task",
                execution(),
                project_id="project-a",
                engine=root_engine,
            )
        registry.release_handoff(first_handoff)
        await execution_registered.wait()

        second_handoff = registry.reserve_handoff()
        with registry.bind_handoff(second_handoff):
            second = handler._track_session(
                "runtime-task",
                queued_duplicate(),
                project_id="project-a",
                engine=root_engine,
            )
        registry.release_handoff(second_handoff)
        await asyncio.sleep(0)
        assert registry.pending_handoff_count == 1

        await asyncio.wait_for(handler.shutdown(timeout=1.0), timeout=1.0)

        assert prepare_called.is_set()
        assert execution_released.is_set()
        assert first.cancelled()
        assert second.cancelled()
        assert registry.pending_handoff_count == 0
        assert not registry.is_active("project-a", "runtime-task")

    asyncio.run(scenario())


def test_ws_shutdown_fails_closed_while_execution_cleanup_is_still_running() -> None:
    async def scenario() -> None:
        cancellation_started = asyncio.Event()
        allow_cleanup = asyncio.Event()

        async def execution() -> None:
            try:
                await asyncio.Event().wait()
            except asyncio.CancelledError:
                cancellation_started.set()
                await allow_cleanup.wait()

        execution_task = asyncio.create_task(execution())
        await asyncio.sleep(0)

        handler = _make_handler(
            SimpleNamespace(
                prepare_active_company_runtimes_for_shutdown=AsyncMock(return_value=[]),
            ),
            _background_tasks={execution_task},
            _task_bg_context={
                execution_task: {"task_id": "runtime-task", "execution_handoff": True}
            },
            _task_bg_map={"runtime-task": {execution_task}},
        )

        try:
            await handler.shutdown(timeout=0.01)
        except RuntimeError as exc:
            assert "execution task(s)" in str(exc)
        else:
            raise AssertionError("shutdown must not close resources before execution cleanup")

        assert cancellation_started.is_set()
        assert not execution_task.done()
        allow_cleanup.set()
        await execution_task

    asyncio.run(scenario())


def test_ws_shutdown_cancels_execution_before_waiting_for_client_close() -> None:
    async def scenario() -> None:
        execution_released = asyncio.Event()
        close_entered = asyncio.Event()
        allow_close = asyncio.Event()

        async def execution() -> None:
            try:
                await asyncio.Event().wait()
            finally:
                execution_released.set()

        class BlockingWebSocket:
            async def close(self, **_kwargs: object) -> None:
                close_entered.set()
                await allow_close.wait()

        execution_task = asyncio.create_task(execution())
        await asyncio.sleep(0)
        client = BlockingWebSocket()
        handler = _make_handler(
            SimpleNamespace(
                prepare_active_company_runtimes_for_shutdown=AsyncMock(return_value=[]),
            ),
            _clients={client},
            _background_tasks={execution_task},
            _task_bg_context={execution_task: {"task_id": "runtime-task"}},
            _task_bg_map={"runtime-task": {execution_task}},
        )

        shutdown_task = asyncio.create_task(handler.shutdown(timeout=1.0))
        # close_entered only fires from inside shutdown's client-close loop.
        # Bound the wait so a shutdown crash before that loop fails the test
        # instead of wedging the whole suite.
        await asyncio.wait_for(close_entered.wait(), timeout=5.0)

        assert execution_released.is_set()
        assert execution_task.done()
        allow_close.set()
        await shutdown_task

    asyncio.run(scenario())


def test_duplicate_resume_does_not_leave_shutdown_handoff_barrier_queued() -> None:
    async def scenario() -> None:
        registry = ActiveTaskRunRegistry()
        execution_started = asyncio.Event()
        execution_released = asyncio.Event()

        async def prepare() -> list[dict]:
            assert registry.is_active("project-a", "runtime-task")
            return []

        root_engine = SimpleNamespace(
            project_id="project-a",
            _active_task_run_registry=registry,
            prepare_active_company_runtimes_for_shutdown=prepare,
        )
        handler = _make_handler(root_engine)

        checkpoint = ExecutionCheckpoint(
            checkpoint_id="checkpoint-1",
            project_id="project-a",
            session_id="runtime-session",
            checkpoint_type="company_runtime_interrupted",
            status="pending",
        )
        task = Task(
            id="ui-task",
            title="Company chat",
            project_id="project-a",
            session_id="runtime-session",
            metadata={"exec_mode": "company"},
        )
        target = {
            "runtime_session_id": "runtime-session",
            "checkpoint": checkpoint,
        }
        handler._resolve_company_runtime_target = AsyncMock(return_value=target)

        async def fake_resume(**_kwargs: object) -> None:
            attempt = registry.register("project-a", "runtime-task")
            execution_started.set()
            try:
                await asyncio.Event().wait()
            finally:
                registry.unregister("project-a", "runtime-task", attempt)
                execution_released.set()

        handler._process_company_suspend_reply = fake_resume

        async def route_once() -> bool:
            handoff = registry.reserve_handoff()
            try:
                with registry.bind_handoff(handoff):
                    return await handler._route_company_suspend_reply_if_pending(
                        task_id=task.id,
                        content="continue",
                        session_id=task.session_id,
                        task=task,
                        attachment_refs=None,
                        message_metadata=None,
                        user_message_id=None,
                        user_message_created_at=None,
                        run_engine=root_engine,
                        run_project_id="project-a",
                    )
            finally:
                registry.release_handoff(handoff)

        assert await route_once() is True
        await execution_started.wait()
        assert await route_once() is True
        assert registry.pending_handoff_count == 0
        assert len(handler._background_tasks) == 1

        await asyncio.wait_for(handler.shutdown(timeout=1.0), timeout=1.0)

        assert execution_released.is_set()
        assert registry.pending_handoff_count == 0

    asyncio.run(scenario())