kamidso-back/tests/kamidso/test_memory_lifecycle.py

"""Memory lifecycle: queues, ReAct scratch, adapter cache, DuckDB caps."""

from __future__ import annotations

import asyncio
import gc
import json
import tracemalloc
from types import SimpleNamespace

import psutil
import pytest

from kamidso.engine.app.openapi.api_v1.react_runtime import (
    REACT_AGENT_MEMORY_CACHE,
    REACT_CONVERSATION_LOCKS,
    acquire_conversation_run,
    conversation_lock_refcount,
    end_react_turn,
    release_conversation_run,
)
from kamidso.engine.core.agent.core.memory.gpts.base import GptsMessage
from kamidso.engine.core.agent.core.memory.gpts.gpts_memory import GptsMemory
from kamidso.engine.core.datasource import duckdb_memory as duck_mem
from kamidso.integrations.engine import events as adapter_events
from kamidso.integrations.engine.engine import KamidsoEngine
from kamidso.integrations.engine.events import (
    adapter_cache_size,
    adapter_run_is_active,
    begin_adapter_run,
    map_engine_event,
    release_adapter_run,
    sweep_adapter_cache,
)


def _message(
    conv_id: str,
    content: str,
    *,
    rounds: int = 0,
    sender: str = "ReAct",
    receiver: str = "Human",
    action_report: str | None = None,
) -> GptsMessage:
    return GptsMessage(
        conv_id=conv_id,
        sender=sender,
        receiver=receiver,
        role="ai" if sender != "Human" else "human",
        content=content,
        rounds=rounds,
        app_code="react_agent",
        app_name="ReAct",
        action_report=action_report,
    )


class _Store:
    def __init__(self) -> None:
        self.rows: list[GptsMessage] = []
        self.deleted: list[str] = []

    def append(self, message: GptsMessage) -> None:
        self.rows.append(message)

    def get_by_conv_id(self, conv_id: str) -> list[GptsMessage]:
        return [row for row in self.rows if row.conv_id == conv_id]

    def get_by_agent(self, conv_id: str, agent: str) -> list[GptsMessage]:
        return [
            row
            for row in self.rows
            if row.conv_id == conv_id and (row.sender == agent or row.receiver == agent)
        ]

    def get_between_agents(self, conv_id: str, agent1: str, agent2: str, current_goal: str | None = None):
        del conv_id, agent1, agent2, current_goal
        return []

    def get_last_message(self, conv_id: str) -> GptsMessage | None:
        rows = self.get_by_conv_id(conv_id)
        return rows[-1] if rows else None

    def delete_by_conv_id(self, conv_id: str) -> None:
        self.deleted.append(conv_id)


class _Plans:
    def __init__(self) -> None:
        self.removed: list[str] = []

    def remove_by_conv_id(self, conv_id: str) -> None:
        self.removed.append(conv_id)


class _LatestStore(_Store):
    """Keeps one row so a repeated-load test does not retain every payload."""

    def append(self, message: GptsMessage) -> None:
        self.rows = [message]


def _memory(store: _Store | None = None) -> tuple[GptsMemory, _Store, _Plans]:
    store = store or _Store()
    plans = _Plans()
    return GptsMemory(plans_memory=plans, message_memory=store), store, plans


def _reset_runtime() -> None:
    REACT_AGENT_MEMORY_CACHE.clear()
    REACT_CONVERSATION_LOCKS.clear()
    from kamidso.engine.app.openapi.api_v1 import react_runtime

    react_runtime._LOCK_REFS.clear()


def test_vis_off_without_reader_does_not_queue_history_snapshot():
    async def _run() -> None:
        memory, store, plans = _memory()
        conv = "conv-quiet"

        async def _forbidden(*_args, **_kwargs):
            raise AssertionError("full history snapshot was built")

        memory.simple_message = _forbidden  # type: ignore[method-assign]
        memory.init(conv, enable_vis_message=False)
        payload = "x" * 50_000
        for index in range(6):
            await memory.append_message(
                conv,
                _message(
                    conv,
                    payload,
                    rounds=index,
                    action_report=json.dumps({"content": f"step-{index}-" + payload[:32]}),
                ),
            )
        queue = memory.queue(conv)
        assert queue is not None and queue.empty()
        assert len(store.rows) == 6
        assert store.deleted == []
        assert plans.removed == []
        memory.dispose_runtime(conv)
        assert memory.queue(conv) is None
        assert store.deleted == []
        assert plans.removed == []
        assert len(store.rows) == 6

    asyncio.run(_run())


def test_vis_off_reader_receives_only_the_latest_message():
    async def _run() -> None:
        memory, _store, _plans = _memory()
        conv = "conv-reader"
        memory.init(conv, enable_vis_message=False)
        memory.attach_queue_reader(conv)
        for index in range(4):
            await memory.append_message(
                conv,
                _message(
                    conv,
                    f"body-{index}",
                    rounds=index,
                    action_report=json.dumps({"content": f"view-{index}"}),
                ),
            )
        queued = []
        queue = memory.queue(conv)
        assert queue is not None
        while not queue.empty():
            queued.append(queue.get_nowait())
        assert [item[0]["markdown"] for item in queued] == [f"view-{i}" for i in range(4)]
        assert all(len(item) == 1 for item in queued)
        memory.detach_queue_reader(conv)
        await memory.complete(conv)
        assert queue.get_nowait() == "[DONE]"
        memory.dispose_runtime(conv)

    asyncio.run(_run())


def test_vis_on_still_publishes_the_legacy_view():
    async def _run() -> None:
        memory, _store, _plans = _memory()
        conv = "conv-vis"

        async def _link(_conv_id: str) -> str:
            return "VIS-SNAPSHOT"

        memory.app_link_chat_message = _link  # type: ignore[method-assign]
        memory.init(conv, enable_vis_message=True)
        await memory.append_message(conv, _message(conv, "hello", rounds=0))
        assert memory.queue(conv).get_nowait() == "VIS-SNAPSHOT"
        memory.clear(conv)
        assert memory._executor is not None
        memory.shutdown_executor()

    asyncio.run(_run())


def test_clear_does_not_delete_persisted_rows_or_owned_executor():
    async def _run() -> None:
        memory, store, plans = _memory()
        conv = "conv-clear"
        memory.init(conv, enable_vis_message=False)
        await memory.append_message(conv, _message(conv, "kept", rounds=0))
        executor = memory._executor
        memory.clear(conv)
        assert store.get_by_conv_id(conv)[0].content == "kept"
        assert store.deleted == []
        assert plans.removed == []
        assert memory._executor is executor
        memory.shutdown_executor()

    asyncio.run(_run())


def test_next_turn_reloads_history_after_dispose():
    async def _run() -> None:
        memory, store, plans = _memory()
        conv = "conv-reload"
        report = json.dumps({"content": "persisted-answer", "is_exe_success": True})
        memory.init(conv, enable_vis_message=False)
        await memory.append_message(
            conv,
            _message(conv, "what is revenue?", rounds=0, sender="Human", receiver="ReAct"),
        )
        await memory.append_message(
            conv,
            _message(conv, "persisted-answer", rounds=1, sender="ReAct", action_report=report),
        )
        memory.dispose_runtime(conv)
        assert memory._executor is None
        assert store.deleted == []
        assert plans.removed == []

        reloaded, _store, plans2 = _memory(store)
        reloaded.init(conv, enable_vis_message=False)
        rows = await reloaded.get_messages(conv)
        assert [row.content for row in rows] == ["what is revenue?", "persisted-answer"]
        history = await reloaded.get_agent_history_memory(conv, "ReAct")
        assert history and history[0].content == "persisted-answer"
        assert plans2.removed == []
        reloaded.dispose_runtime(conv)

    asyncio.run(_run())


def test_metadb_rows_survive_dispose(tmp_path):
    from sqlalchemy.pool import NullPool

    from kamidso.engine.core.storage.metadata.db_manager import db
    from kamidso.engine.serve.agent.agents.db_gpts_memory import (
        MetaDbGptsMessageMemory,
        MetaDbGptsPlansMemory,
    )

    db_file = tmp_path / "lifecycle.db"
    db.init_db(
        f"sqlite:///{db_file.as_posix()}",
        engine_args={"connect_args": {"check_same_thread": False}, "poolclass": NullPool},
    )
    db.create_all()

    async def _run() -> None:
        conv = "conv-metadb-lifecycle"
        memory = GptsMemory(
            plans_memory=MetaDbGptsPlansMemory(),
            message_memory=MetaDbGptsMessageMemory(),
        )
        memory.init(conv, enable_vis_message=False)
        await memory.append_message(conv, _message(conv, "durable", rounds=1, sender="Human"))
        memory.dispose_runtime(conv)
        rows = MetaDbGptsMessageMemory().get_by_conv_id(conv)
        assert rows and rows[0].content == "durable"
        follow = GptsMemory(
            plans_memory=MetaDbGptsPlansMemory(),
            message_memory=MetaDbGptsMessageMemory(),
        )
        follow.init(conv, enable_vis_message=False)
        loaded = await follow.get_messages(conv)
        assert loaded[0].content == "durable"
        follow.dispose_runtime(conv)

    asyncio.run(_run())


def test_same_conversation_is_serialized_and_lock_survives_a_waiter():
    async def _run() -> None:
        _reset_runtime()
        conv = "conv-same"
        order: list[str] = []
        locks: list[asyncio.Lock] = []

        class _Scratch:
            def __init__(self) -> None:
                self.disposed = False

            def dispose_runtime(self, conv_id: str | None = None) -> None:
                del conv_id
                self.disposed = True

        async def _turn(name: str, hold: float) -> None:
            lock = await acquire_conversation_run(conv)
            locks.append(lock)
            order.append(f"enter:{name}")
            try:
                if conv not in REACT_AGENT_MEMORY_CACHE:
                    REACT_AGENT_MEMORY_CACHE[conv] = _Scratch()
                await asyncio.sleep(hold)
                order.append(f"exit:{name}")
            finally:
                await end_react_turn(conv, lock)

        first = asyncio.create_task(_turn("a", 0.08))
        await asyncio.sleep(0.02)
        second = asyncio.create_task(_turn("b", 0.01))
        await asyncio.sleep(0.02)
        assert conversation_lock_refcount(conv) == 2
        assert REACT_CONVERSATION_LOCKS[conv].locked()
        await first
        await second
        assert order == ["enter:a", "exit:a", "enter:b", "exit:b"]
        assert locks[0] is locks[1]
        assert conv not in REACT_CONVERSATION_LOCKS
        assert conv not in REACT_AGENT_MEMORY_CACHE
        assert conversation_lock_refcount(conv) == 0

    asyncio.run(_run())


def test_cancelled_waiter_does_not_pin_the_conversation_lock():
    async def _run() -> None:
        _reset_runtime()
        conv = "conv-cancel-wait"
        holder = await acquire_conversation_run(conv)

        async def _wait() -> None:
            await acquire_conversation_run(conv)

        waiter = asyncio.create_task(_wait())
        await asyncio.sleep(0.05)
        assert conversation_lock_refcount(conv) == 2
        waiter.cancel()
        with pytest.raises(asyncio.CancelledError):
            await waiter
        assert conversation_lock_refcount(conv) == 1
        assert REACT_CONVERSATION_LOCKS[conv] is holder
        await release_conversation_run(conv, holder)
        assert conv not in REACT_CONVERSATION_LOCKS

    asyncio.run(_run())


def test_twenty_different_conversations_run_together():
    async def _run() -> None:
        _reset_runtime()
        active = 0
        peak = 0

        async def _one(index: int) -> None:
            nonlocal active, peak
            conv = f"user-{index}"
            lock = await acquire_conversation_run(conv)
            active += 1
            peak = max(peak, active)
            await asyncio.sleep(0.05)
            active -= 1
            await release_conversation_run(conv, lock)

        await asyncio.gather(*(_one(index) for index in range(20)))
        assert peak == 20
        assert REACT_CONVERSATION_LOCKS == {}

    asyncio.run(_run())


def test_repeated_turns_release_scratch_and_allocated_memory_stabilizes():
    async def _cycle(store: _LatestStore, index: int) -> None:
        memory, _store, _plans = _memory(store)
        conv = "conv-repeat"
        memory.init(conv, enable_vis_message=False)
        payload = ("row-" + str(index) + "-") * 4_000
        for step in range(6):
            await memory.append_message(
                conv,
                _message(
                    conv,
                    payload,
                    rounds=step,
                    action_report=json.dumps({"content": payload[:64]}),
                ),
            )
        assert memory.queue(conv).empty()
        executor = memory._executor
        memory.dispose_runtime(conv)
        assert executor is not None and getattr(executor, "_shutdown", False)
        assert memory._executor is None
        assert memory.queue(conv) is None
        assert conv not in memory.messages_cache

    async def _run() -> None:
        store = _LatestStore()
        tracemalloc.start()
        alloc: list[int] = []
        rss: list[int] = []
        process = psutil.Process()
        try:
            for index in range(100):
                await _cycle(store, index)
                if index % 10 == 9:
                    gc.collect()
                    current, _peak = tracemalloc.get_traced_memory()
                    alloc.append(current)
                    rss.append(process.memory_info().rss)
        finally:
            tracemalloc.stop()
        growth = alloc[-1] - alloc[2]
        print(
            "LIFECYCLE_MEASURE",
            f"alloc_bytes={alloc}",
            f"rss_bytes={rss}",
            f"alloc_growth={growth}",
            f"rss_growth={rss[-1] - rss[2]}",
        )
        assert store.deleted == []
        assert len(store.rows) == 1
        # Warmup is the first samples. Later allocated memory must not climb
        # with the number of turns. RSS may stay above the process baseline.
        assert growth < 8 * 1024 * 1024

    asyncio.run(_run())


def test_adapter_release_covers_done_error_and_abandoned_runs():
    done_id = "run-done"
    error_id = "run-error"
    quiet_id = "run-quiet"
    active_id = "run-active"
    try:
        begin_adapter_run(done_id)
        mapped = map_engine_event(done_id, {"type": "step.start", "id": "s", "title": "sql"})
        assert mapped and mapped[0]["type"] == "run.started"
        assert adapter_run_is_active(done_id)
        release_adapter_run(done_id)
        assert not adapter_run_is_active(done_id)

        begin_adapter_run(error_id)
        try:
            map_engine_event(error_id, {"type": "step.start", "id": "s", "title": "code"})
            raise RuntimeError("stream failed")
        except RuntimeError:
            release_adapter_run(error_id)
        assert not adapter_run_is_active(error_id)

        map_engine_event(quiet_id, {"type": "step.start", "id": "s", "title": "sql"})
        begin_adapter_run(active_id)
        map_engine_event(active_id, {"type": "step.start", "id": "s", "title": "sql"})
        with adapter_events._LOCK:
            adapter_events._TOUCHED[quiet_id] = 0.0
            adapter_events._TOUCHED[active_id] = 0.0
        removed = sweep_adapter_cache(now=10_000.0)
        assert quiet_id in removed
        assert active_id not in removed
        assert adapter_run_is_active(active_id)
    finally:
        release_adapter_run(done_id)
        release_adapter_run(error_id)
        release_adapter_run(quiet_id)
        release_adapter_run(active_id)


def test_release_does_not_drop_another_active_run():
    first = "run-keep"
    second = "run-drop"
    try:
        begin_adapter_run(first)
        begin_adapter_run(second)
        map_engine_event(first, {"type": "step.start", "id": "a", "title": "sql"})
        map_engine_event(second, {"type": "step.start", "id": "b", "title": "sql"})
        release_adapter_run(second)
        assert adapter_run_is_active(first)
        assert not adapter_run_is_active(second)
        assert adapter_cache_size() >= 1
    finally:
        release_adapter_run(first)
        release_adapter_run(second)


def test_engine_releases_adapter_on_success_error_and_disconnect():
    async def _collect(engine: KamidsoEngine, run_id: str):
        run = SimpleNamespace(id=run_id)
        events = []
        async for event in engine.execute(
            user=None,  # type: ignore[arg-type]
            run=run,  # type: ignore[arg-type]
            request=None,  # type: ignore[arg-type]
            materialized=[],
        ):
            events.append(event)
        return events

    async def _run() -> None:
        engine = KamidsoEngine.__new__(KamidsoEngine)
        engine._runtime = SimpleNamespace(ready=False)
        failed = await _collect(engine, "run-unavailable")
        assert failed[0]["error"]["code"] == "ENGINE_UNAVAILABLE"
        assert not adapter_run_is_active("run-unavailable")

        engine._runtime = SimpleNamespace(ready=True)

        async def _ok(**_kwargs):
            yield {"type": "run.delta"}

        engine._execute_ready = _ok  # type: ignore[method-assign]
        events = await _collect(engine, "run-ok")
        assert events == [{"type": "run.delta"}]
        assert not adapter_run_is_active("run-ok")

        async def _boom(**_kwargs):
            yield {"type": "partial"}
            raise RuntimeError("execution failed")

        engine._execute_ready = _boom  # type: ignore[method-assign]
        with pytest.raises(RuntimeError):
            await _collect(engine, "run-boom")
        assert not adapter_run_is_active("run-boom")

        async def _slow(**_kwargs):
            yield {"type": "partial"}
            await asyncio.sleep(30)

        engine._execute_ready = _slow  # type: ignore[method-assign]
        generator = engine.execute(
            user=None,  # type: ignore[arg-type]
            run=SimpleNamespace(id="run-disconnect"),  # type: ignore[arg-type]
            request=None,  # type: ignore[arg-type]
            materialized=[],
        )
        first = await generator.__anext__()
        assert first["type"] == "partial"
        await generator.aclose()
        assert not adapter_run_is_active("run-disconnect")

    asyncio.run(_run())


def _limit_bytes(text: str) -> float:
    number, unit = str(text).split()
    scales = {
        "b": 1,
        "byte": 1,
        "bytes": 1,
        "kib": 1024,
        "mib": 1024**2,
        "gib": 1024**3,
        "tib": 1024**4,
        "kb": 1000,
        "mb": 1000**2,
        "gb": 1000**3,
        "tb": 1000**4,
    }
    return float(number) * scales[unit.lower()]


def test_duckdb_invalid_limit_falls_back_to_two_gigabytes(monkeypatch):
    monkeypatch.setenv("KAMIDSO_DUCKDB_MEMORY_LIMIT", "not-a-size")
    monkeypatch.delenv("KAMIDSO_DUCKDB_AGGREGATE_MEMORY_LIMIT", raising=False)
    assert duck_mem.configured_memory_limit() == "2GB"
    assert duck_mem.configured_aggregate_limit() == "2GB"


def test_duckdb_tool_connections_share_one_aggregate_budget(monkeypatch):
    monkeypatch.setenv("KAMIDSO_DUCKDB_MEMORY_LIMIT", "2GB")
    monkeypatch.delenv("KAMIDSO_DUCKDB_AGGREGATE_MEMORY_LIMIT", raising=False)
    duck_mem.reset_duckdb_memory_state_for_tests()
    duck_mem.cap_duckdb_default_connection()
    first = duck_mem.open_tool_duckdb(":memory:")
    second = None
    try:
        alone = first.execute("SELECT current_setting('memory_limit')").fetchone()[0]
        assert _limit_bytes(alone) == pytest.approx(2_000_000_000, rel=0.15)
        assert len(duck_mem._LIVE) == 1
        second = duck_mem.open_tool_duckdb(":memory:")
        limits = [
            _limit_bytes(conn.execute("SELECT current_setting('memory_limit')").fetchone()[0])
            for conn in (first, second)
        ]
        assert all(700 * 1024**2 < limit < 1.2 * 1024**3 for limit in limits)
        duck_mem.release_tool_duckdb(second)
        second.close()
        second = None
        restored = first.execute("SELECT current_setting('memory_limit')").fetchone()[0]
        assert _limit_bytes(restored) == pytest.approx(2_000_000_000, rel=0.15)
    finally:
        duck_mem.release_tool_duckdb(first)
        first.close()
        if second is not None:
            duck_mem.release_tool_duckdb(second)
            second.close()
        duck_mem.reset_duckdb_memory_state_for_tests()


def test_excel_reader_registers_and_releases_duckdb(tmp_path, monkeypatch):
    monkeypatch.setenv("KAMIDSO_DUCKDB_MEMORY_LIMIT", "2GB")
    monkeypatch.delenv("KAMIDSO_DUCKDB_AGGREGATE_MEMORY_LIMIT", raising=False)
    duck_mem.reset_duckdb_memory_state_for_tests()
    from kamidso.engine.app.scene.chat_data.chat_excel.excel_reader import ExcelReader

    csv_path = tmp_path / "tiny.csv"
    csv_path.write_text("city,amount\nShiraz,10\nTehran,20\n", encoding="utf-8")
    monkeypatch.setattr(
        "kamidso.engine.core.util.file_client.knowledge_upload_root",
        lambda: str(tmp_path),
    )
    reader = ExcelReader("conv-excel", str(csv_path), read_type="df")
    live_after_close = None
    try:
        limit = reader.db.execute("SELECT current_setting('memory_limit')").fetchone()[0]
        assert _limit_bytes(limit) == pytest.approx(2_000_000_000, rel=0.15)
        assert len(duck_mem._LIVE) == 1
        rows = reader.db.execute("SELECT COUNT(*) FROM temp_table").fetchone()[0]
        assert rows == 2
    finally:
        reader.close()
        live_after_close = dict(duck_mem._LIVE)
        duck_mem.reset_duckdb_memory_state_for_tests()
    assert reader.db is None
    assert live_after_close == {}


def test_duckdb_connector_binds_sqlalchemy_connections(tmp_path, monkeypatch):
    monkeypatch.setenv("KAMIDSO_DUCKDB_MEMORY_LIMIT", "2GB")
    monkeypatch.delenv("KAMIDSO_DUCKDB_AGGREGATE_MEMORY_LIMIT", raising=False)
    duck_mem.reset_duckdb_memory_state_for_tests()
    from kamidso.engine.ext.datasource.rdbms.conn_duckdb import DuckDbConnector

    connector = DuckDbConnector.from_file_path(str(tmp_path / "analytics.duckdb"))
    try:
        with connector._engine.connect() as connection:
            limit = connection.exec_driver_sql("SELECT current_setting('memory_limit')").scalar()
        assert _limit_bytes(str(limit)) == pytest.approx(2_000_000_000, rel=0.15)
    finally:
        connector._engine.dispose()
        duck_mem.reset_duckdb_memory_state_for_tests()