unslothai--unsloth/studio/backend/tests/test_tool_approvals.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Concurrency tests for the per-call tool-call confirmation gate.

``state.tool_approvals`` coordinates two threads: the agentic loop thread
blocked in ``wait_tool_decision`` and the request thread that delivers the
user's choice through ``resolve_tool_decision``. Each gated call carries a
unique ``approval_id`` so a stale or concurrent confirmation can never
resolve the wrong call. These tests exercise that handshake directly --
no model, no server -- so the race windows are fast and deterministic.
"""

import threading
import time

import pytest

from state import tool_approvals
from state.tool_approvals import (
    TOOL_REJECTED_MESSAGE,
    abort_tool_decision,
    begin_tool_decision,
    new_approval_id,
    request_tool_decision,
    resolve_tool_decision,
    wait_tool_decision,
)


@pytest.fixture(autouse = True)
def _clear_pending():
    """Each test starts and ends with an empty ``_pending`` map."""
    with tool_approvals._lock:
        tool_approvals._pending.clear()
    yield
    with tool_approvals._lock:
        tool_approvals._pending.clear()


class _Waiter:
    """Run ``request_tool_decision`` in a thread and capture its result."""

    def __init__(
        self,
        session_id,
        approval_id,
        cancel_event = None,
        timeout = None,
    ):
        self.session_id = session_id
        self.approval_id = approval_id
        self.cancel_event = cancel_event
        self.timeout = timeout
        self.result = None
        self._thread = threading.Thread(target = self._run, daemon = True)

    def _run(self):
        kwargs = {"cancel_event": self.cancel_event}
        if self.timeout is not None:
            kwargs["timeout"] = self.timeout
        self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs)

    def start(self):
        self._thread.start()
        _wait_until(lambda: _has_pending(self.approval_id))
        return self

    def join(self, timeout = 5.0):
        self._thread.join(timeout = timeout)
        assert not self._thread.is_alive(), "waiter thread did not finish"
        return self.result


def _has_pending(approval_id) -> bool:
    with tool_approvals._lock:
        return approval_id in tool_approvals._pending


def _wait_until(
    pred,
    timeout = 2.0,
    interval = 0.005,
) -> bool:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if pred():
            return True
        time.sleep(interval)
    return False


# ── Basic allow / deny ───────────────────────────────────────────────


def test_allow_decision():
    aid = new_approval_id()
    w = _Waiter("sess", aid).start()
    assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
    assert w.join() == "allow"


def test_deny_decision():
    aid = new_approval_id()
    w = _Waiter("sess", aid).start()
    assert resolve_tool_decision(aid, "deny", session_id = "sess") is True
    assert w.join() == "deny"


def test_slot_cleaned_up_after_decision():
    aid = new_approval_id()
    w = _Waiter("sess", aid).start()
    resolve_tool_decision(aid, "allow")
    w.join()
    assert _wait_until(lambda: not _has_pending(aid))


def test_abort_tool_decision_removes_unwaited_slot():
    aid = new_approval_id()
    slot = begin_tool_decision("sess", aid)
    abort_tool_decision(slot, aid)
    assert not _has_pending(aid)
    assert resolve_tool_decision(aid, "allow", session_id = "sess") is False


def test_approval_ids_are_unique():
    ids = {new_approval_id() for _ in range(1000)}
    assert len(ids) == 1000


# ── Pre-registration race (begin before wait) ────────────────────────


def test_resolve_before_wait_is_not_lost():
    """A decision delivered after ``begin`` but before ``wait`` survives.

    The loop registers the slot before it yields ``tool_start``, so even a
    confirmation that races ahead of the blocking ``wait`` is recorded on
    the slot and returned -- never dropped.
    """
    aid = new_approval_id()
    slot = begin_tool_decision("sess", aid)
    assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
    # wait() is only entered now, after the decision already landed.
    assert wait_tool_decision(slot, aid) == "allow"
    assert not _has_pending(aid)


# ── Resolver edge cases ──────────────────────────────────────────────


def test_resolve_unknown_approval_returns_false():
    assert resolve_tool_decision(new_approval_id(), "allow") is False


def test_resolve_empty_approval_returns_false():
    assert resolve_tool_decision("", "allow") is False
    assert resolve_tool_decision(None, "allow") is False


def test_resolve_wrong_session_scope_returns_false():
    aid = new_approval_id()
    w = _Waiter("sess-a", aid).start()
    # Correct approval_id but the wrong session must not resolve it.
    assert resolve_tool_decision(aid, "allow", session_id = "sess-b") is False
    assert _has_pending(aid)
    # The right session still works.
    assert resolve_tool_decision(aid, "allow", session_id = "sess-a") is True
    assert w.join() == "allow"


def test_duplicate_resolve_after_completion_returns_false():
    aid = new_approval_id()
    w = _Waiter("sess", aid).start()
    assert resolve_tool_decision(aid, "allow") is True
    w.join()
    assert _wait_until(lambda: not _has_pending(aid))
    assert resolve_tool_decision(aid, "deny") is False


def test_first_decision_is_immutable():
    """A second confirmation cannot flip an already-recorded decision.

    The waiter reads ``slot["decision"]`` outside the lock and then cleans up,
    so a duplicate or out-of-order POST that lands in that window must be
    rejected and must not overwrite the first decision -- an Allow can never
    become a Deny. Distinct from the after-completion case above: here the slot
    is still pending (no waiter has consumed it yet).
    """
    aid = new_approval_id()
    slot = begin_tool_decision("sess", aid)
    assert resolve_tool_decision(aid, "allow", session_id = "sess") is True
    # Second decision, same id, before any waiter consumes/cleans the slot.
    assert resolve_tool_decision(aid, "deny", session_id = "sess") is False
    assert slot["decision"] == "allow"
    # The waiter still observes the first (immutable) decision.
    assert wait_tool_decision(slot, aid) == "allow"
    assert not _has_pending(aid)


# ── Cancellation and timeout ─────────────────────────────────────────


def test_cancel_event_breaks_wait_as_deny():
    cancel = threading.Event()
    aid = new_approval_id()
    w = _Waiter("sess", aid, cancel_event = cancel).start()
    cancel.set()
    assert w.join(timeout = 3.0) == "deny"
    assert _wait_until(lambda: not _has_pending(aid))


def test_timeout_returns_deny():
    aid = new_approval_id()
    start = time.monotonic()
    result = request_tool_decision("sess", aid, timeout = 0.1)
    assert result == "deny"
    assert time.monotonic() - start < 2.0
    assert not _has_pending(aid)


# ── Independence across concurrent calls ─────────────────────────────


def test_two_pending_calls_same_session_are_independent():
    """Keying on approval_id, not session, keeps concurrent calls distinct.

    Resolving the first call's id must not unblock or alter the second
    call pending in the same session.
    """
    a1, a2 = new_approval_id(), new_approval_id()
    w1 = _Waiter("sess", a1).start()
    w2 = _Waiter("sess", a2).start()

    assert resolve_tool_decision(a1, "deny", session_id = "sess") is True
    assert w1.join() == "deny"
    # w2 is still waiting on its own id.
    assert _has_pending(a2)
    assert resolve_tool_decision(a2, "allow", session_id = "sess") is True
    assert w2.join() == "allow"


def test_concurrent_distinct_calls_route_their_own_decisions():
    n = 25
    waiters = {}
    for i in range(n):
        aid = new_approval_id()
        waiters[aid] = _Waiter(f"s{i}", aid).start()
    expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)}
    for aid, decision in expected.items():
        assert resolve_tool_decision(aid, decision) is True
    for aid, w in waiters.items():
        assert w.join() == expected[aid]


# ── Constants ────────────────────────────────────────────────────────


def test_rejected_message_is_user_facing_text():
    assert isinstance(TOOL_REJECTED_MESSAGE, str)
    assert TOOL_REJECTED_MESSAGE.strip()