from datetime import datetime, timezone
from typing import Any, TYPE_CHECKING, cast
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4

import pytest

from frisk_sdk import DecisionOutcome
from frisk_sdk.adapters.langchain.frisk_tool_middleware import FriskToolMiddleware
from frisk_sdk.core.session_registry import SessionRegistry
from frisk_sdk.core.tool_approval_request import (
    ToolApprovalRequest,
    ToolApprovalRequestStatusUpdateResult,
    ToolApprovalStatus,
)
from frisk_sdk.core.tool_call_decision import ToolCallDecision
from frisk_sdk.errors.invocation_errors import (
    InvalidFriskSessionError,
    MissingFriskSessionContextError,
)


# FriskToolMiddleware expects a `FriskLangchain` instance, but in tests we only
# need the small surface area used by the middleware.
class DummyFrisk:
    def __init__(self, *args: Any, **kwargs: Any):
        self.logger_name = "frisk"
        self._tracer = None
        self._adapter = MagicMock()
        self.normalize_tool_call: MagicMock = MagicMock()
        self.get_or_create_many_tool_approval_requests: AsyncMock = AsyncMock(
            return_value={}
        )
        self.cancel_many_tool_approval_requests: AsyncMock = AsyncMock(return_value={})


if TYPE_CHECKING:
    from typing import Protocol

    class DummyFriskProtocol(Protocol):
        logger_name: str
        _tracer: Any
        _adapter: Any
        normalize_tool_call: MagicMock
        get_or_create_many_tool_approval_requests: AsyncMock
        cancel_many_tool_approval_requests: AsyncMock


class DummyFriskSession:
    def __init__(self, result, session_id=None):
        self._result = result
        self.session_id = session_id or uuid4()
        self.id = str(self.session_id)
        self.decide_tool_call_calls: list[dict[str, Any]] = []

    def decide_tool_call(self, *, tool_call, agent_state):
        self.decide_tool_call_calls.append(
            {"tool_call": tool_call, "agent_state": agent_state}
        )
        return self._resolve_result(tool_call=tool_call, agent_state=agent_state)

    def _resolve_result(self, *, tool_call, agent_state):
        if callable(self._result):
            return self._result(tool_call=tool_call, agent_state=agent_state)
        if isinstance(self._result, list):
            return self._result.pop(0)
        if isinstance(self._result, dict):
            return self._result[tool_call.id]
        return self._result


class DummyToolCall:
    def __init__(self, name="tool", id="id1", args=None):
        self.name = name
        self.id = id
        self.args = args or {}


class DummyToolRuntime:
    def __init__(self, context: dict[str, Any]):
        self.context = context


class DummyProcessToolCallResult:
    def __init__(self, outcome, reason="reason"):
        self.outcome = outcome
        self.reason = reason


class DummyToolMessage:
    def __init__(self, content, tool_call_id, name, status=None):
        self.content = content
        self.tool_call_id = tool_call_id
        self.name = name
        self.status = status


class DummyAIMessage:
    def __init__(self, content="", tool_calls=None):
        self.content = content
        self.tool_calls = tool_calls or []


class DummyInterrupt(Exception):
    def __init__(self, payload: dict[str, Any]):
        super().__init__("interrupt")
        self.payload = payload


class StatefulInterruptMock:
    def __init__(self):
        self.has_pending_response = False
        self.response_value: Any = None
        self.payloads: list[dict[str, Any]] = []

    def resume(self, response: Any):
        self.response_value = response

    def __call__(self, payload: dict[str, Any]):
        self.payloads.append(payload)
        if not self.has_pending_response:
            self.has_pending_response = True
            raise DummyInterrupt(payload)

        response_value = self.response_value
        self.has_pending_response = False
        self.response_value = None
        return response_value


def _build_state_with_tool_calls(tool_calls: list[Any]) -> dict[str, Any]:
    return {"messages": [DummyAIMessage(tool_calls=tool_calls)]}


def _make_decision(
    outcome: DecisionOutcome,
    *,
    reason: str = "reason",
    policy_id: str | None = "policy-1",
    policy_version_id: str | None = "version-1",
) -> ToolCallDecision:
    return ToolCallDecision(
        outcome=outcome,
        rules_matched_count=1,
        reason=reason,
        policy_id=policy_id,
        policy_version_id=policy_version_id,
    )


def _make_approval_request(
    *,
    tool_call_id: str,
    tool_name: str = "tool",
    request_id: str = "approval-1",
    status: ToolApprovalStatus = ToolApprovalStatus.PENDING,
) -> ToolApprovalRequest:
    return ToolApprovalRequest(
        id=request_id,
        tool_call_id=tool_call_id,
        session_id="session-1",
        policy_id="policy-1",
        policy_version_id="version-1",
        tool_name=tool_name,
        agent_id="agent-1",
        status=status,
        organization_id="org-1",
        requested_at=datetime.now(timezone.utc),
        decided_at=None,
        decided_by_user_id=None,
        decision_notes=None,
    )


def _make_cancellation_result(
    approval_request: ToolApprovalRequest,
) -> ToolApprovalRequestStatusUpdateResult:
    return ToolApprovalRequestStatusUpdateResult(
        updated=True,
        result=approval_request,
    )


async def _run_with_resumes(
    middleware: FriskToolMiddleware,
    *,
    state: dict[str, Any],
    runtime: Any,
    interrupt_mock: StatefulInterruptMock,
    responses: list[Any],
):
    response_iter = iter(responses)

    while True:
        try:
            return await middleware.aafter_model(cast(Any, state), cast(Any, runtime))
        except DummyInterrupt:
            try:
                interrupt_mock.resume(next(response_iter))
            except StopIteration:
                raise DummyInterrupt(interrupt_mock.payloads[-1])


@pytest.fixture(autouse=True)
def _patch_langchain_messages(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.AIMessage",
        DummyAIMessage,
    )
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.ToolMessage",
        DummyToolMessage,
    )


def _register_runtime_for_session(session: DummyFriskSession) -> DummyToolRuntime:
    registry = SessionRegistry()
    registry.register(str(session.session_id), cast(Any, session))
    return DummyToolRuntime({"frisk_session_id": str(session.session_id)})


@pytest.mark.asyncio
async def test_aafter_model_deny():
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    result = DummyProcessToolCallResult(outcome=DecisionOutcome.DENY, reason="nope")

    session_id = uuid4()
    session = DummyFriskSession(result, session_id)

    registry = SessionRegistry()
    registry.register(str(session_id), cast(Any, session))

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    state = _build_state_with_tool_calls(["call"])

    mw = FriskToolMiddleware(cast(Any, frisk))
    out = await mw.aafter_model(cast(Any, state), cast(Any, runtime))

    assert out is not None
    assert len(out["messages"]) == 1
    denied_message = out["messages"][0]
    assert isinstance(denied_message, DummyToolMessage)
    assert "denied" in denied_message.content
    assert denied_message.tool_call_id == tool_call.id
    assert denied_message.name == tool_call.name

    registry.unregister(str(session_id))


@pytest.mark.asyncio
async def test_aafter_model_allow_returns_empty_messages_update():
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    result = DummyProcessToolCallResult(outcome=DecisionOutcome.ALLOW)

    session_id = uuid4()
    session = DummyFriskSession(result, session_id)

    registry = SessionRegistry()
    registry.register(str(session_id), cast(Any, session))

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    state = _build_state_with_tool_calls(["call"])

    mw = FriskToolMiddleware(cast(Any, frisk))
    out = await mw.aafter_model(cast(Any, state), cast(Any, runtime))

    assert out == {"messages": []}

    registry.unregister(str(session_id))


@pytest.mark.asyncio
async def test_aafter_model_missing_session():
    frisk = DummyFrisk("dummy_api_key")
    frisk.normalize_tool_call.return_value = DummyToolCall()
    runtime = DummyToolRuntime({})
    state = _build_state_with_tool_calls(["call"])

    mw = FriskToolMiddleware(cast(Any, frisk))
    with pytest.raises(MissingFriskSessionContextError):
        await mw.aafter_model(cast(Any, state), cast(Any, runtime))


@pytest.mark.asyncio
async def test_aafter_model_invalid_session():
    frisk = DummyFrisk("dummy_api_key")
    frisk.normalize_tool_call.return_value = DummyToolCall()
    runtime = DummyToolRuntime({"frisk_session_id": object()})
    state = _build_state_with_tool_calls(["call"])

    mw = FriskToolMiddleware(cast(Any, frisk))
    with pytest.raises(InvalidFriskSessionError):
        await mw.aafter_model(cast(Any, state), cast(Any, runtime))


@pytest.mark.asyncio
async def test_aafter_model_without_tool_calls_returns_none():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))

    runtime = DummyToolRuntime({"frisk_session_id": str(uuid4())})

    assert await mw.aafter_model({"messages": []}, cast(Any, runtime)) is None
    assert (
        await mw.aafter_model(
            cast(Any, {"messages": [DummyAIMessage(tool_calls=[])]}),
            cast(Any, runtime),
        )
        is None
    )


@pytest.mark.asyncio
async def test_aafter_model_escalate_invalid_tool_call_id_reinterrupts_with_error(
    monkeypatch: pytest.MonkeyPatch,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call
    frisk.get_or_create_many_tool_approval_requests.return_value = {
        tool_call.id: _make_approval_request(
            tool_call_id=tool_call.id,
            tool_name=tool_call.name,
        )
    }

    session = DummyFriskSession(_make_decision(DecisionOutcome.ESCALATE))
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        with pytest.raises(DummyInterrupt):
            await _run_with_resumes(
                mw,
                state=state,
                runtime=runtime,
                interrupt_mock=interrupt_mock,
                responses=[{"unknown-tool": "cancel"}],
            )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    assert len(interrupt_mock.payloads) == 3
    assert (
        "invalid interrupt response" in interrupt_mock.payloads[-1]["message"].lower()
    )
    frisk.cancel_many_tool_approval_requests.assert_not_awaited()


@pytest.mark.asyncio
async def test_aafter_model_escalate_cancel_updates_pending_requests(
    monkeypatch: pytest.MonkeyPatch,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call

    pending_request = _make_approval_request(
        tool_call_id=tool_call.id,
        tool_name=tool_call.name,
        request_id="approval-1",
        status=ToolApprovalStatus.PENDING,
    )
    canceled_request = _make_approval_request(
        tool_call_id=tool_call.id,
        tool_name=tool_call.name,
        request_id="approval-1",
        status=ToolApprovalStatus.CANCELED,
    )

    frisk.get_or_create_many_tool_approval_requests.return_value = {
        tool_call.id: pending_request
    }
    frisk.cancel_many_tool_approval_requests.return_value = {
        canceled_request.id: _make_cancellation_result(canceled_request)
    }

    session = DummyFriskSession(_make_decision(DecisionOutcome.ESCALATE))
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        out = await _run_with_resumes(
            mw,
            state=state,
            runtime=runtime,
            interrupt_mock=interrupt_mock,
            responses=[{tool_call.id: "cancel"}],
        )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    frisk.cancel_many_tool_approval_requests.assert_awaited_once_with(
        approval_request_ids=[pending_request.id]
    )
    assert out is not None
    assert len(out["messages"]) == 1
    assert "canceled" in out["messages"][0].content


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "resume_response",
    [
        {"tool-1": "retry"},
        {},
    ],
)
async def test_aafter_model_escalate_retry_or_empty_response_reinterrupts(
    monkeypatch: pytest.MonkeyPatch,
    resume_response: dict[str, str],
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call
    frisk.get_or_create_many_tool_approval_requests.return_value = {
        tool_call.id: _make_approval_request(
            tool_call_id=tool_call.id,
            tool_name=tool_call.name,
            status=ToolApprovalStatus.PENDING,
        )
    }

    session = DummyFriskSession(_make_decision(DecisionOutcome.ESCALATE))
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        with pytest.raises(DummyInterrupt):
            await _run_with_resumes(
                mw,
                state=state,
                runtime=runtime,
                interrupt_mock=interrupt_mock,
                responses=[resume_response],
            )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    assert len(interrupt_mock.payloads) == 3
    frisk.cancel_many_tool_approval_requests.assert_not_awaited()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("resumed_outcome", "expected_message"),
    [
        (DecisionOutcome.ALLOW, None),
        (DecisionOutcome.DENY, "denied"),
    ],
)
async def test_aafter_model_escalate_then_decision_change_on_resume(
    monkeypatch: pytest.MonkeyPatch,
    resumed_outcome: DecisionOutcome,
    expected_message: str | None,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call

    pending_request = _make_approval_request(
        tool_call_id=tool_call.id,
        tool_name=tool_call.name,
    )
    frisk.get_or_create_many_tool_approval_requests.side_effect = [
        {tool_call.id: pending_request},
        {},
    ]

    session = DummyFriskSession(
        [
            _make_decision(DecisionOutcome.ESCALATE),
            _make_decision(resumed_outcome, reason="changed"),
        ]
    )
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        out = await _run_with_resumes(
            mw,
            state=state,
            runtime=runtime,
            interrupt_mock=interrupt_mock,
            responses=[{}],
        )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    if expected_message is None:
        assert out == {"messages": []}
    else:
        assert expected_message in out["messages"][0].content


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "initial_outcome",
    [DecisionOutcome.ALLOW, DecisionOutcome.DENY],
)
async def test_aafter_model_decision_change_from_non_escalate_to_escalate_interrupts(
    monkeypatch: pytest.MonkeyPatch,
    initial_outcome: DecisionOutcome,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call

    pending_request = _make_approval_request(
        tool_call_id=tool_call.id,
        tool_name=tool_call.name,
    )
    frisk.get_or_create_many_tool_approval_requests.side_effect = [
        {},
        {tool_call.id: pending_request},
    ]

    session = DummyFriskSession(
        [
            _make_decision(initial_outcome, reason="initial"),
            _make_decision(DecisionOutcome.ESCALATE),
        ]
    )
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        first_out = await mw.aafter_model(cast(Any, state), cast(Any, runtime))
        with pytest.raises(DummyInterrupt):
            await mw.aafter_model(cast(Any, state), cast(Any, runtime))
    finally:
        SessionRegistry().unregister(str(session.session_id))

    if initial_outcome == DecisionOutcome.ALLOW:
        assert first_out == {"messages": []}
    else:
        assert "denied" in cast(Any, first_out)["messages"][0].content
    assert len(interrupt_mock.payloads) == 1


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("approved_status", "expected_message"),
    [
        (ToolApprovalStatus.APPROVED, None),
        (ToolApprovalStatus.REJECTED, "denied after it was escalated"),
    ],
)
async def test_aafter_model_retry_sees_updated_approval_status(
    monkeypatch: pytest.MonkeyPatch,
    approved_status: ToolApprovalStatus,
    expected_message: str | None,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call

    frisk.get_or_create_many_tool_approval_requests.side_effect = [
        {
            tool_call.id: _make_approval_request(
                tool_call_id=tool_call.id,
                tool_name=tool_call.name,
                status=ToolApprovalStatus.PENDING,
            )
        },
        {
            tool_call.id: _make_approval_request(
                tool_call_id=tool_call.id,
                tool_name=tool_call.name,
                status=approved_status,
            )
        },
    ]

    session = DummyFriskSession(
        [
            _make_decision(DecisionOutcome.ESCALATE),
            _make_decision(DecisionOutcome.ESCALATE),
        ]
    )
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        out = await _run_with_resumes(
            mw,
            state=state,
            runtime=runtime,
            interrupt_mock=interrupt_mock,
            responses=[{tool_call.id: "retry"}],
        )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    if expected_message is None:
        assert out == {"messages": []}
    else:
        assert expected_message in out["messages"][0].content


@pytest.mark.asyncio
async def test_aafter_model_cancel_failure_reinterrupts_with_error(
    monkeypatch: pytest.MonkeyPatch,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call
    pending_request = _make_approval_request(
        tool_call_id=tool_call.id,
        tool_name=tool_call.name,
        request_id="approval-1",
    )
    frisk.get_or_create_many_tool_approval_requests.return_value = {
        tool_call.id: pending_request
    }
    frisk.cancel_many_tool_approval_requests.return_value = {}

    session = DummyFriskSession(_make_decision(DecisionOutcome.ESCALATE))
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        with pytest.raises(DummyInterrupt):
            await _run_with_resumes(
                mw,
                state=state,
                runtime=runtime,
                interrupt_mock=interrupt_mock,
                responses=[{tool_call.id: "cancel"}],
            )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    assert len(interrupt_mock.payloads) == 3
    assert (
        "failed to cancel approval request"
        in interrupt_mock.payloads[-1]["message"].lower()
    )


@pytest.mark.asyncio
async def test_aafter_model_get_or_create_failure_keeps_request_pending(
    monkeypatch: pytest.MonkeyPatch,
):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    frisk.normalize_tool_call.return_value = tool_call
    frisk.get_or_create_many_tool_approval_requests.return_value = {tool_call.id: None}

    session = DummyFriskSession(_make_decision(DecisionOutcome.ESCALATE))
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([tool_call])

    try:
        with pytest.raises(DummyInterrupt):
            await _run_with_resumes(
                mw,
                state=state,
                runtime=runtime,
                interrupt_mock=interrupt_mock,
                responses=[{}],
            )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    assert len(interrupt_mock.payloads) == 3


@pytest.mark.asyncio
async def test_aafter_model_mixed_deny_and_escalate_only_interrupts_pending_tools(
    monkeypatch: pytest.MonkeyPatch,
):
    frisk = DummyFrisk("dummy_api_key")
    denied_tool = DummyToolCall(id="tool-1", name="weather", args={"city": "sf"})
    escalated_tool = DummyToolCall(id="tool-2", name="email", args={"to": "a@b.com"})
    frisk.normalize_tool_call.side_effect = [
        denied_tool,
        escalated_tool,
        denied_tool,
        escalated_tool,
    ]

    pending_request = _make_approval_request(
        tool_call_id=escalated_tool.id,
        tool_name=escalated_tool.name,
        request_id="approval-2",
    )
    canceled_request = _make_approval_request(
        tool_call_id=escalated_tool.id,
        tool_name=escalated_tool.name,
        request_id="approval-2",
        status=ToolApprovalStatus.CANCELED,
    )
    frisk.get_or_create_many_tool_approval_requests.return_value = {
        escalated_tool.id: pending_request
    }
    frisk.cancel_many_tool_approval_requests.return_value = {
        canceled_request.id: _make_cancellation_result(canceled_request)
    }

    session = DummyFriskSession(
        {
            denied_tool.id: _make_decision(
                DecisionOutcome.DENY, reason="policy denied"
            ),
            escalated_tool.id: _make_decision(DecisionOutcome.ESCALATE),
        }
    )
    runtime = _register_runtime_for_session(session)
    interrupt_mock = StatefulInterruptMock()
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.interrupt",
        interrupt_mock,
    )

    mw = FriskToolMiddleware(cast(Any, frisk))
    state = _build_state_with_tool_calls([denied_tool, escalated_tool])

    try:
        out = await _run_with_resumes(
            mw,
            state=state,
            runtime=runtime,
            interrupt_mock=interrupt_mock,
            responses=[{escalated_tool.id: "cancel"}],
        )
    finally:
        SessionRegistry().unregister(str(session.session_id))

    assert set(interrupt_mock.payloads[0]["escalated_tool_calls"]) == {
        escalated_tool.id
    }
    assert len(out["messages"]) == 2
    assert any("policy denied" in message.content for message in out["messages"])
    assert any("canceled" in message.content for message in out["messages"])


def test_get_frisk_session_from_runtime_valid():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))

    session_id = uuid4()
    session = DummyFriskSession(None, session_id)

    registry = SessionRegistry()
    registry.register(str(session_id), cast(Any, session))

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    out = mw.resolve_frisk_session(runtime)
    assert out is session

    registry.unregister(str(session_id))


def test_get_frisk_session_from_runtime_missing():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))
    runtime = DummyToolRuntime({})
    with pytest.raises(MissingFriskSessionContextError):
        mw.resolve_frisk_session(runtime)


def test_get_frisk_session_from_runtime_invalid():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))
    runtime = DummyToolRuntime({"frisk_session_id": object()})
    with pytest.raises(InvalidFriskSessionError):
        mw.resolve_frisk_session(runtime)
