"""
Tests for the five integration guard subclasses:
  - LangChainGuard    (wrap_tool, wrap_tool_async)
  - LlamaIndexGuard   (wrap_query_engine, wrap_query_engine_async)
  - OpenAIAgentsGuard (function_tool — ImportError path + mocked path)
  - BedrockGuard      (wrap_invoke_agent, wrap_invoke_agent_async)
  - GeminiGuard       (call_if_allowed, call_if_allowed_async)

Strategy: no real LangChain / AWS / Gemini SDKs required.
We use simple stub objects for tools, engines, and clients.
Every guard routes through the same base ContextGuard, so the heavy
policy + audit path is already covered; here we verify:
  1. Each guard's _source_type tag ends up in the audit log.
  2. ALLOW path: underlying function is called and result returned.
  3. DENY path: PermissionError raised, underlying function NOT called.
  4. last_event_id is updated after each call.
  5. Guard-specific wrappers (wrap_tool, wrap_query_engine, etc.) work
     correctly including edge-case fallbacks.
"""

import asyncio
import textwrap
import uuid
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from autopil.langchain_guard import LangChainGuard
from autopil.llamaindex_guard import LlamaIndexGuard
from autopil.openai_agents_guard import OpenAIAgentsGuard
from autopil.bedrock_guard import BedrockGuard
from autopil.gemini_guard import GeminiGuard
from autopil.models import SensitivityLevel


# ── shared policy/db fixtures ─────────────────────────────────────────────────

POLICY_YAML = textwrap.dedent("""
    policies:
      - name: analyst_policy
        agent_role: analyst
        allowed_sources:
          - reports
          - market_data
        denied_sources:
          - executive_comms
        max_sensitivity: high
""")

ALLOW_KWARGS = dict(agent_role="analyst", user_id="u1", source_id="reports")
DENY_KWARGS  = dict(agent_role="analyst", user_id="u1", source_id="executive_comms")


@pytest.fixture
def policy_file(tmp_path) -> str:
    p = tmp_path / "policies.yaml"
    p.write_text(POLICY_YAML)
    return str(p)


@pytest.fixture
def audit_db(tmp_path) -> str:
    return str(tmp_path / "audit.db")


def _last_event_source_type(guard, session_id: str) -> str | None:
    """Return source_type of the most recent event for a session."""
    events = guard.get_audit_trail(session_id)
    return events[-1].source_type if events else None


def _last_event_decision(guard, session_id: str) -> str | None:
    events = guard.get_audit_trail(session_id)
    return events[-1].decision.value if events else None


# ── LangChainGuard ────────────────────────────────────────────────────────────

class TestLangChainGuard:

    @pytest.fixture
    def guard(self, policy_file, audit_db):
        return LangChainGuard(policy_path=policy_file, audit_db=audit_db)

    def _make_tool(self, return_value="tool result", has_arun=True):
        """Minimal stub mimicking a LangChain BaseTool."""
        tool = MagicMock()
        tool._run = MagicMock(return_value=return_value)
        if has_arun:
            tool._arun = AsyncMock(return_value=return_value)
        else:
            del tool._arun
        return tool

    def test_source_type_is_langchain(self):
        assert LangChainGuard._source_type == "langchain"

    def test_wrap_tool_allow_calls_underlying(self, guard):
        tool = self._make_tool()
        sid = str(uuid.uuid4())
        guard.wrap_tool(tool, **ALLOW_KWARGS, session_id=sid)
        result = tool._run("what is the revenue?")
        assert result == "tool result"
        tool._run.__wrapped__ if hasattr(tool._run, "__wrapped__") else None
        # underlying MagicMock was called once
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_tool_deny_raises_permission_error(self, guard):
        tool = self._make_tool()
        sid = str(uuid.uuid4())
        guard.wrap_tool(tool, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            tool._run("show executive emails")
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_tool_deny_does_not_call_underlying(self, guard):
        """Underlying _run must not be called on a DENY."""
        call_log = []
        tool = MagicMock()
        tool._run = MagicMock(side_effect=lambda q: call_log.append(q) or "data")
        guard.wrap_tool(tool, **DENY_KWARGS)
        with pytest.raises(PermissionError):
            tool._run("secret query")
        assert call_log == [], "Underlying tool was called despite DENY"

    def test_wrap_tool_audit_source_type(self, guard):
        tool = self._make_tool()
        sid = str(uuid.uuid4())
        guard.wrap_tool(tool, **ALLOW_KWARGS, session_id=sid)
        tool._run("market query")
        assert _last_event_source_type(guard, sid) == "langchain"

    def test_wrap_tool_updates_last_event_id(self, guard):
        tool = self._make_tool()
        guard.wrap_tool(tool, **ALLOW_KWARGS)
        assert guard.last_event_id is None
        tool._run("query")
        assert guard.last_event_id is not None

    def test_wrap_tool_async_allow(self, guard):
        tool = self._make_tool(has_arun=True)
        sid = str(uuid.uuid4())
        guard.wrap_tool_async(tool, **ALLOW_KWARGS, session_id=sid)
        result = asyncio.get_event_loop().run_until_complete(tool._arun("market data"))
        assert result == "tool result"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_tool_async_deny(self, guard):
        tool = self._make_tool(has_arun=True)
        sid = str(uuid.uuid4())
        guard.wrap_tool_async(tool, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            asyncio.get_event_loop().run_until_complete(tool._arun("exec comms"))
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_tool_async_fallback_when_no_arun(self, guard):
        """If the tool has no _arun, wrap_tool_async falls back to wrapping _run."""
        tool = self._make_tool(has_arun=False)
        sid = str(uuid.uuid4())
        guard.wrap_tool_async(tool, **ALLOW_KWARGS, session_id=sid)
        # _arun should now be callable (created by the guard)
        assert callable(tool._arun)
        result = asyncio.get_event_loop().run_until_complete(tool._arun("query"))
        assert result == "tool result"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_tool_audit_source_type_async(self, guard):
        tool = self._make_tool(has_arun=True)
        sid = str(uuid.uuid4())
        guard.wrap_tool_async(tool, **ALLOW_KWARGS, session_id=sid)
        asyncio.get_event_loop().run_until_complete(tool._arun("q"))
        assert _last_event_source_type(guard, sid) == "langchain"


# ── LlamaIndexGuard ───────────────────────────────────────────────────────────

class TestLlamaIndexGuard:

    @pytest.fixture
    def guard(self, policy_file, audit_db):
        return LlamaIndexGuard(policy_path=policy_file, audit_db=audit_db)

    def _make_engine(self, return_value="llama result", has_aquery=True):
        engine = MagicMock()
        engine.query = MagicMock(return_value=return_value)
        if has_aquery:
            engine.aquery = AsyncMock(return_value=return_value)
        else:
            del engine.aquery
        return engine

    def test_source_type_is_llamaindex(self):
        assert LlamaIndexGuard._source_type == "llamaindex"

    def test_wrap_query_engine_allow(self, guard):
        engine = self._make_engine()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine(engine, **ALLOW_KWARGS, session_id=sid)
        result = guarded.query("what are the quarterly results?")
        assert result == "llama result"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_query_engine_deny(self, guard):
        engine = self._make_engine()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine(engine, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            guarded.query("exec comms query")
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_query_engine_deny_does_not_call_underlying(self, guard):
        call_log = []
        engine = MagicMock()
        engine.query = MagicMock(side_effect=lambda q: call_log.append(q) or "data")
        guarded = guard.wrap_query_engine(engine, **DENY_KWARGS)
        with pytest.raises(PermissionError):
            guarded.query("secret")
        assert call_log == []

    def test_wrap_query_engine_passthrough_attributes(self, guard):
        """Non-query attributes on the underlying engine are still accessible."""
        engine = self._make_engine()
        engine.index_name = "my_index"
        guarded = guard.wrap_query_engine(engine, **ALLOW_KWARGS)
        assert guarded.index_name == "my_index"

    def test_wrap_query_engine_audit_source_type(self, guard):
        engine = self._make_engine()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine(engine, **ALLOW_KWARGS, session_id=sid)
        guarded.query("q")
        assert _last_event_source_type(guard, sid) == "llamaindex"

    def test_wrap_query_engine_updates_last_event_id(self, guard):
        engine = self._make_engine()
        guarded = guard.wrap_query_engine(engine, **ALLOW_KWARGS)
        assert guard.last_event_id is None
        guarded.query("q")
        assert guard.last_event_id is not None

    def test_wrap_query_engine_async_allow(self, guard):
        engine = self._make_engine(has_aquery=True)
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine_async(engine, **ALLOW_KWARGS, session_id=sid)
        result = asyncio.get_event_loop().run_until_complete(guarded.aquery("revenue?"))
        assert result == "llama result"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_query_engine_async_deny(self, guard):
        engine = self._make_engine(has_aquery=True)
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine_async(engine, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            asyncio.get_event_loop().run_until_complete(guarded.aquery("exec query"))
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_query_engine_async_sync_query_still_works(self, guard):
        """The async-wrapped engine also exposes a sync .query() method."""
        engine = self._make_engine()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine_async(engine, **ALLOW_KWARGS, session_id=sid)
        result = guarded.query("sync query")
        assert result == "llama result"

    def test_wrap_query_engine_async_fallback_no_aquery(self, guard):
        """If engine has no aquery, guard falls back to wrapping sync .query()."""
        engine = self._make_engine(has_aquery=False)
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine_async(engine, **ALLOW_KWARGS, session_id=sid)
        result = asyncio.get_event_loop().run_until_complete(guarded.aquery("q"))
        assert result == "llama result"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_query_engine_async_audit_source_type(self, guard):
        engine = self._make_engine()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_query_engine_async(engine, **ALLOW_KWARGS, session_id=sid)
        asyncio.get_event_loop().run_until_complete(guarded.aquery("q"))
        assert _last_event_source_type(guard, sid) == "llamaindex"


# ── OpenAIAgentsGuard ─────────────────────────────────────────────────────────

class TestOpenAIAgentsGuard:

    @pytest.fixture
    def guard(self, policy_file, audit_db):
        return OpenAIAgentsGuard(policy_path=policy_file, audit_db=audit_db)

    def test_source_type_is_openai_agents(self):
        assert OpenAIAgentsGuard._source_type == "openai_agents"

    def test_function_tool_raises_import_error_when_sdk_missing(self, guard):
        """If openai-agents is not installed, function_tool raises ImportError with install hint."""
        with patch.dict("sys.modules", {"agents": None}):
            with pytest.raises(ImportError, match="pip install openai-agents"):
                decorator = guard.function_tool(**ALLOW_KWARGS)
                # The ImportError is raised when the decorator is applied
                @decorator
                def my_fn(query: str) -> str:
                    return "result"

    def test_function_tool_allow_with_mocked_sdk(self, guard):
        """With a mocked openai-agents SDK, allowed calls execute and audit correctly."""
        fake_function_tool = lambda fn: fn  # identity decorator
        sid = str(uuid.uuid4())

        with patch("autopil.openai_agents_guard.OpenAIAgentsGuard.function_tool") as mock_ft:
            # Bypass the SDK and test the protect() path directly
            call_log = []

            @guard.protect(**ALLOW_KWARGS, session_id=sid)
            def my_tool(query: str) -> str:
                call_log.append(query)
                return "openai result"

            result = my_tool("market data query")
            assert result == "openai result"
            assert call_log == ["market data query"]
            assert _last_event_decision(guard, sid) == "ALLOW"
            assert _last_event_source_type(guard, sid) == "openai_agents"

    def test_function_tool_deny_with_protect(self, guard):
        """Denied requests raise PermissionError and don't call the underlying fn."""
        sid = str(uuid.uuid4())
        call_log = []

        @guard.protect(**DENY_KWARGS, session_id=sid)
        def my_tool(query: str) -> str:
            call_log.append(query)
            return "secret data"

        with pytest.raises(PermissionError, match="DENIED"):
            my_tool("exec comms")
        assert call_log == []
        assert _last_event_decision(guard, sid) == "DENY"

    def test_function_tool_mocked_decorator_allow(self, guard):
        """Full function_tool path with a mocked agents module."""
        sid = str(uuid.uuid4())
        call_log = []

        fake_agents = MagicMock()
        fake_agents.function_tool = lambda fn: fn  # identity — don't wrap

        with patch.dict("sys.modules", {"agents": fake_agents}):
            decorator = guard.function_tool(**ALLOW_KWARGS, session_id=sid)

            @decorator
            def get_report(query: str) -> str:
                call_log.append(query)
                return "report data"

            result = get_report("Q4 results")
            assert result == "report data"
            assert call_log == ["Q4 results"]
            assert _last_event_decision(guard, sid) == "ALLOW"

    def test_function_tool_mocked_decorator_deny(self, guard):
        sid = str(uuid.uuid4())
        call_log = []

        fake_agents = MagicMock()
        fake_agents.function_tool = lambda fn: fn

        with patch.dict("sys.modules", {"agents": fake_agents}):
            decorator = guard.function_tool(**DENY_KWARGS, session_id=sid)

            @decorator
            def get_comms(query: str) -> str:
                call_log.append(query)
                return "exec comms"

            with pytest.raises(PermissionError, match="DENIED"):
                get_comms("give me exec emails")
            assert call_log == []
            assert _last_event_decision(guard, sid) == "DENY"

    def test_function_tool_audit_source_type(self, guard):
        sid = str(uuid.uuid4())

        @guard.protect(**ALLOW_KWARGS, session_id=sid)
        def my_tool(query: str) -> str:
            return "data"

        my_tool("q")
        assert _last_event_source_type(guard, sid) == "openai_agents"


# ── BedrockGuard ──────────────────────────────────────────────────────────────

class TestBedrockGuard:

    @pytest.fixture
    def guard(self, policy_file, audit_db):
        return BedrockGuard(policy_path=policy_file, audit_db=audit_db)

    def _make_bedrock_client(self, return_value="bedrock response"):
        client = MagicMock()
        client.invoke_agent = MagicMock(return_value=return_value)
        return client

    def test_source_type_is_bedrock(self):
        assert BedrockGuard._source_type == "bedrock"

    def test_wrap_invoke_agent_allow(self, guard):
        client = self._make_bedrock_client()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_invoke_agent(client, **ALLOW_KWARGS, session_id=sid)
        result = guarded.invoke_agent(inputText="summarise Q3 report")
        assert result == "bedrock response"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_invoke_agent_deny(self, guard):
        client = self._make_bedrock_client()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_invoke_agent(client, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            guarded.invoke_agent(inputText="exec comms summary")
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_invoke_agent_deny_does_not_call_underlying(self, guard):
        call_log = []
        client = MagicMock()
        client.invoke_agent = MagicMock(side_effect=lambda **kw: call_log.append(kw) or "data")
        guarded = guard.wrap_invoke_agent(client, **DENY_KWARGS)
        with pytest.raises(PermissionError):
            guarded.invoke_agent(inputText="exec comms")
        assert call_log == []

    def test_wrap_invoke_agent_passthrough_attributes(self, guard):
        """Other boto3 client methods are still accessible via __getattr__."""
        client = self._make_bedrock_client()
        client.get_agent = MagicMock(return_value={"agentId": "abc"})
        guarded = guard.wrap_invoke_agent(client, **ALLOW_KWARGS)
        assert guarded.get_agent() == {"agentId": "abc"}

    def test_wrap_invoke_agent_audit_source_type(self, guard):
        client = self._make_bedrock_client()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_invoke_agent(client, **ALLOW_KWARGS, session_id=sid)
        guarded.invoke_agent(inputText="q")
        assert _last_event_source_type(guard, sid) == "bedrock"

    def test_wrap_invoke_agent_updates_last_event_id(self, guard):
        client = self._make_bedrock_client()
        guarded = guard.wrap_invoke_agent(client, **ALLOW_KWARGS)
        assert guard.last_event_id is None
        guarded.invoke_agent(inputText="q")
        assert guard.last_event_id is not None

    def test_wrap_invoke_agent_async_allow(self, guard):
        client = self._make_bedrock_client()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_invoke_agent_async(client, **ALLOW_KWARGS, session_id=sid)
        result = asyncio.get_event_loop().run_until_complete(
            guarded.invoke_agent(inputText="async bedrock query")
        )
        assert result == "bedrock response"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_wrap_invoke_agent_async_deny(self, guard):
        client = self._make_bedrock_client()
        sid = str(uuid.uuid4())
        guarded = guard.wrap_invoke_agent_async(client, **DENY_KWARGS, session_id=sid)
        with pytest.raises(PermissionError, match="DENIED"):
            asyncio.get_event_loop().run_until_complete(
                guarded.invoke_agent(inputText="exec emails")
            )
        assert _last_event_decision(guard, sid) == "DENY"

    def test_wrap_invoke_agent_async_passthrough(self, guard):
        client = self._make_bedrock_client()
        client.list_agents = MagicMock(return_value=[])
        guarded = guard.wrap_invoke_agent_async(client, **ALLOW_KWARGS)
        assert guarded.list_agents() == []


# ── GeminiGuard ───────────────────────────────────────────────────────────────

class TestGeminiGuard:

    @pytest.fixture
    def guard(self, policy_file, audit_db):
        return GeminiGuard(policy_path=policy_file, audit_db=audit_db)

    def test_source_type_is_gemini(self):
        assert GeminiGuard._source_type == "gemini"

    def test_call_if_allowed_allow(self, guard):
        sid = str(uuid.uuid4())
        call_log = []

        def get_data(customer_id: str) -> dict:
            call_log.append(customer_id)
            return {"balance": 10000}

        result = guard.call_if_allowed(
            get_data, **ALLOW_KWARGS, session_id=sid, fn_args={"customer_id": "c123"}
        )
        assert result == {"balance": 10000}
        assert call_log == ["c123"]
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_call_if_allowed_deny(self, guard):
        sid = str(uuid.uuid4())
        call_log = []

        def get_comms(customer_id: str) -> str:
            call_log.append(customer_id)
            return "exec emails"

        with pytest.raises(PermissionError, match="DENIED"):
            guard.call_if_allowed(
                get_comms, **DENY_KWARGS, session_id=sid, fn_args={"customer_id": "c1"}
            )
        assert call_log == []
        assert _last_event_decision(guard, sid) == "DENY"

    def test_call_if_allowed_no_fn_args(self, guard):
        """fn_args defaults to {} — fn called with no arguments."""
        sid = str(uuid.uuid4())

        def get_summary() -> str:
            return "summary"

        result = guard.call_if_allowed(get_summary, **ALLOW_KWARGS, session_id=sid)
        assert result == "summary"
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_call_if_allowed_audit_source_type(self, guard):
        sid = str(uuid.uuid4())
        guard.call_if_allowed(lambda: "ok", **ALLOW_KWARGS, session_id=sid)
        assert _last_event_source_type(guard, sid) == "gemini"

    def test_call_if_allowed_updates_last_event_id(self, guard):
        assert guard.last_event_id is None
        guard.call_if_allowed(lambda: "ok", **ALLOW_KWARGS)
        assert guard.last_event_id is not None

    def test_call_if_allowed_async_allow(self, guard):
        sid = str(uuid.uuid4())
        call_log = []

        async def get_data_async(customer_id: str) -> dict:
            call_log.append(customer_id)
            return {"limit": 50000}

        result = asyncio.get_event_loop().run_until_complete(
            guard.call_if_allowed_async(
                get_data_async, **ALLOW_KWARGS, session_id=sid,
                fn_args={"customer_id": "c999"},
            )
        )
        assert result == {"limit": 50000}
        assert call_log == ["c999"]
        assert _last_event_decision(guard, sid) == "ALLOW"

    def test_call_if_allowed_async_deny(self, guard):
        sid = str(uuid.uuid4())
        call_log = []

        async def get_comms_async(customer_id: str) -> str:
            call_log.append(customer_id)
            return "exec email"

        with pytest.raises(PermissionError, match="DENIED"):
            asyncio.get_event_loop().run_until_complete(
                guard.call_if_allowed_async(
                    get_comms_async, **DENY_KWARGS, session_id=sid,
                    fn_args={"customer_id": "c1"},
                )
            )
        assert call_log == []
        assert _last_event_decision(guard, sid) == "DENY"

    def test_call_if_allowed_async_audit_source_type(self, guard):
        sid = str(uuid.uuid4())

        async def fn() -> str:
            return "ok"

        asyncio.get_event_loop().run_until_complete(
            guard.call_if_allowed_async(fn, **ALLOW_KWARGS, session_id=sid)
        )
        assert _last_event_source_type(guard, sid) == "gemini"

    def test_call_if_allowed_async_updates_last_event_id(self, guard):
        """
        last_event_id is stored in a ContextVar — it's task-local in async.
        We must read it from within the same coroutine context where the call ran.
        """
        async def run() -> str | None:
            async def fn() -> str:
                return "ok"
            await guard.call_if_allowed_async(fn, **ALLOW_KWARGS)
            return guard.last_event_id  # read inside same async context

        event_id = asyncio.get_event_loop().run_until_complete(run())
        assert event_id is not None


# ── cross-guard: each guard is a separate ContextGuard instance ───────────────

class TestGuardIsolation:
    """Separate guard instances don't share audit state."""

    def test_each_guard_has_independent_audit_log(self, policy_file, audit_db, tmp_path):
        db2 = str(tmp_path / "audit2.db")
        g1 = LangChainGuard(policy_path=policy_file, audit_db=audit_db)
        g2 = LlamaIndexGuard(policy_path=policy_file, audit_db=db2)

        sid = str(uuid.uuid4())

        tool = MagicMock()
        tool._run = MagicMock(return_value="data")
        g1.wrap_tool(tool, **ALLOW_KWARGS, session_id=sid)
        tool._run("q")

        # g2 has its own DB — should see no events for this session
        assert g2.get_audit_trail(sid) == []

    def test_source_types_are_distinct_across_guards(self, policy_file, tmp_path):
        sid = str(uuid.uuid4())
        dbs = {name: str(tmp_path / f"{name}.db") for name in ("lc", "li", "bk", "gm")}

        lc_guard = LangChainGuard(policy_path=policy_file, audit_db=dbs["lc"])
        li_guard = LlamaIndexGuard(policy_path=policy_file, audit_db=dbs["li"])
        bk_guard = BedrockGuard(policy_path=policy_file, audit_db=dbs["bk"])
        gm_guard = GeminiGuard(policy_path=policy_file, audit_db=dbs["gm"])

        tool = MagicMock()
        tool._run = MagicMock(return_value="r")
        lc_guard.wrap_tool(tool, **ALLOW_KWARGS, session_id=sid)
        tool._run("q")

        engine = MagicMock()
        engine.query = MagicMock(return_value="r")
        guarded_engine = li_guard.wrap_query_engine(engine, **ALLOW_KWARGS, session_id=sid)
        guarded_engine.query("q")

        bk_client = MagicMock()
        bk_client.invoke_agent = MagicMock(return_value="r")
        guarded_client = bk_guard.wrap_invoke_agent(bk_client, **ALLOW_KWARGS, session_id=sid)
        guarded_client.invoke_agent(inputText="q")

        gm_guard.call_if_allowed(lambda: "r", **ALLOW_KWARGS, session_id=sid)

        assert _last_event_source_type(lc_guard, sid) == "langchain"
        assert _last_event_source_type(li_guard, sid) == "llamaindex"
        assert _last_event_source_type(bk_guard, sid) == "bedrock"
        assert _last_event_source_type(gm_guard, sid) == "gemini"
