diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index cce95fc34c..38ba635627 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4577,6 +4577,7 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4587,6 +4588,12 @@ class LlamaCppBackend: {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) """ from core.inference.tools import execute_tool + from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + begin_tool_decision, + new_approval_id, + wait_tool_decision, + ) if not self.is_loaded: raise RuntimeError("llama-server is not loaded") @@ -5197,20 +5204,50 @@ class LlamaCppBackend: status_text = f"Calling: {tool_name}" yield {"type": "status", "text": status_text} - yield { - "type": "tool_start", - "tool_name": tool_name, - "tool_call_id": tc.get("id", ""), - "arguments": arguments, - } - # ── Duplicate call detection ────────────── # str(dict) is stable here: arguments always comes from # json.loads on the same model output within one request, # so insertion order is deterministic (Python 3.7+). _tc_key = tool_name + str(arguments) _prev = _tool_call_history[-1] if _tool_call_history else None - if _prev and _prev[0] == _tc_key and not _prev[1]: + _is_duplicate = bool(_prev) and _prev[0] == _tc_key and not _prev[1] + # Guard against the model emitting a tool not in the + # per-request advertised set: filtered MCP names, a + # built-in the caller opted out of, or a stale name + # from a prior turn. Mirrors the safetensors loop's + # allowed_tool_names check. + _allowed = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if (t.get("function") or {}).get("name") + } + _is_disabled = bool(_allowed) and tool_name not in _allowed + # Only gate calls that would actually run: duplicate or + # disabled calls are short-circuited below and never + # execute, so prompting for them would be noise. + # Registering the slot before tool_start closes the race + # where the confirmation could arrive before the waiter. + _needs_confirm = ( + confirm_tool_calls and not _is_duplicate and not _is_disabled + ) + _approval_id = new_approval_id() if _needs_confirm else "" + _decision_slot = ( + begin_tool_decision(session_id, _approval_id) + if _needs_confirm + else None + ) + + yield { + "type": "tool_start", + "tool_name": tool_name, + "tool_call_id": tc.get("id", ""), + "arguments": arguments, + "approval_id": _approval_id, + "awaiting_confirmation": _needs_confirm, + } + + _denied = False + if _is_duplicate: result = ( "You already made this exact call. " "Do not repeat the same tool call. " @@ -5219,27 +5256,28 @@ class LlamaCppBackend: "process data you already have, or " "provide your final answer now." ) - else: - _effective_timeout = ( - None if tool_call_timeout >= 9999 else tool_call_timeout + elif _is_disabled: + result = ( + f"Error: tool '{tool_name}' is not enabled " + "for this request. Use one of the enabled " + "tools or provide a final answer." ) - # Guard against the model emitting a tool not in the - # per-request advertised set: filtered MCP names, a - # built-in the caller opted out of, or a stale name - # from a prior turn. Mirrors the safetensors loop's - # allowed_tool_names check. - _allowed = { - (t.get("function") or {}).get("name") - for t in (tools or []) - if (t.get("function") or {}).get("name") - } - if _allowed and tool_name not in _allowed: - result = ( - f"Error: tool '{tool_name}' is not enabled " - "for this request. Use one of the enabled " - "tools or provide a final answer." + else: + _denied = ( + _decision_slot is not None + and wait_tool_decision( + _decision_slot, + _approval_id, + cancel_event = cancel_event, ) + == "deny" + ) + if _denied: + result = TOOL_REJECTED_MESSAGE else: + _effective_timeout = ( + None if tool_call_timeout >= 9999 else tool_call_timeout + ) result = execute_tool( tool_name, arguments, @@ -5269,7 +5307,11 @@ class LlamaCppBackend: _is_error = isinstance(result, str) and result.lstrip().startswith( _error_prefixes ) - _tool_call_history.append((_tc_key, _is_error)) + # A user-denied call never executed, so it must not count + # toward duplicate detection — otherwise re-issuing and + # approving the same call would be rejected as a duplicate. + if not _denied: + _tool_call_history.append((_tc_key, _is_error)) # Strip image sentinel before feeding result to the LLM # (the full result with sentinel is still yielded via # tool_end so the frontend can extract image paths). diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 7e7d7026f6..cfead128f0 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -839,6 +839,7 @@ class InferenceOrchestrator: tool_call_timeout: int = 300, session_id: Optional[str] = None, use_adapter: Optional[Union[bool, str]] = None, + confirm_tool_calls: bool = False, **_unused, ): """Run the safetensors agentic tool loop in this (parent) @@ -895,6 +896,7 @@ class InferenceOrchestrator: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + confirm_tool_calls = confirm_tool_calls, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 73bb3d090a..24b598a2a6 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -24,6 +24,13 @@ from urllib.parse import urlparse from loggers import get_logger +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) + from core.inference.tool_call_parser import ( BUDGET_EXHAUSTED_NUDGE, DUPLICATE_CALL_NUDGE, @@ -105,6 +112,7 @@ def run_safetensors_tool_loop( max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -308,27 +316,55 @@ def run_safetensors_tool_loop( tool_name = tool_name, ) + tc_key = tool_name + str(arguments) + is_disabled = ( + bool(allowed_tool_names) and tool_name not in allowed_tool_names + ) + already_ran_ok = any( + k == tc_key and not err for k, err in tool_call_history + ) + # Only gate calls that would actually run: a disabled or + # duplicate call is short-circuited below and never executes, so + # asking the user to approve it would be noise. Registering the + # approval slot *before* tool_start closes the race where the + # confirmation could arrive before the waiter exists. + needs_confirm = ( + confirm_tool_calls and not is_disabled and not already_ran_ok + ) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = ( + begin_tool_decision(session_id, approval_id) if needs_confirm else None + ) + yield {"type": "status", "text": _status_for_tool(tool_name, arguments)} yield { "type": "tool_start", "tool_name": tool_name, "tool_call_id": tc.get("id", ""), "arguments": arguments, + "approval_id": approval_id, + "awaiting_confirmation": needs_confirm, } - tc_key = tool_name + str(arguments) - if allowed_tool_names and tool_name not in allowed_tool_names: + denied = False + if is_disabled: result = ( f"Error: tool '{tool_name}' is not enabled for this " "request. Use one of the enabled tools or provide a " "final answer." ) + elif already_ran_ok: + result = DUPLICATE_CALL_NUDGE else: - already_ran_ok = any( - k == tc_key and not err for k, err in tool_call_history + denied = ( + decision_slot is not None + and wait_tool_decision( + decision_slot, approval_id, cancel_event = cancel_event + ) + == "deny" ) - if already_ran_ok: - result = DUPLICATE_CALL_NUDGE + if denied: + result = TOOL_REJECTED_MESSAGE else: eff_timeout = ( None if tool_call_timeout >= 9999 else tool_call_timeout @@ -355,7 +391,11 @@ def run_safetensors_tool_loop( is_error = isinstance(result, str) and result.lstrip().startswith( TOOL_ERROR_PREFIXES ) - tool_call_history.append((tc_key, is_error)) + # A user-denied call never executed, so it must not count toward + # duplicate detection — otherwise re-issuing and approving the + # same call would be wrongly rejected as a duplicate. + if not denied: + tool_call_history.append((tc_key, is_error)) # Strip frontend image sentinel from the model's view. # Cut at the first occurrence so leading and consecutive diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index bb1bd394d1..d6f24b5159 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -720,6 +720,10 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.", ) + confirm_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.", + ) auto_heal_tool_calls: Optional[bool] = Field( True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", @@ -945,6 +949,12 @@ class ChatCompletionRequest(BaseModel): return self +class ToolConfirmRequest(BaseModel): + session_id: Optional[str] = None + approval_id: Optional[str] = None + decision: Literal["allow", "deny"] = "deny" + + # ── OpenAI shell-tool container management ───────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7d1c7b2488..d87875ca01 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -186,6 +186,7 @@ from models.inference import ( CompletionUsage, ValidateModelRequest, ValidateModelResponse, + ToolConfirmRequest, TextContentPart, ImageContentPart, ImageUrl, @@ -1260,6 +1261,26 @@ async def cancel_inference( return {"cancelled": n} +@studio_router.post("/tool-confirm") +async def confirm_tool_call( + body: ToolConfirmRequest, + current_subject: str = Depends(get_current_subject), +): + """Allow or deny a tool call awaiting user confirmation. + + Identified by ``approval_id`` (echoed from the ``tool_start`` event); + ``session_id`` is a scope check. Returns {"resolved": bool}. ``False`` + means no matching call was waiting (e.g. a stale or duplicate + confirmation, or a mismatched session). + """ + from state.tool_approvals import resolve_tool_decision + + resolved = resolve_tool_decision( + body.approval_id, body.decision, session_id = body.session_id + ) + return {"resolved": resolved} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, @@ -2790,6 +2811,7 @@ async def openai_chat_completions( if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + confirm_tool_calls = bool(payload.confirm_tool_calls), ) _tool_sentinel = object() @@ -3310,6 +3332,7 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, use_adapter = payload.use_adapter, + confirm_tool_calls = bool(payload.confirm_tool_calls), ) _sf_tool_sentinel = object() diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py new file mode 100644 index 0000000000..1fc7981eb7 --- /dev/null +++ b/studio/backend/state/tool_approvals.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Per-call tool-call confirmation gate. + +When a chat request sets ``confirm_tool_calls``, the agentic loop pauses +before executing each tool and waits here for the user's decision, which +arrives via ``POST /api/inference/tool-confirm`` on a separate connection. + +Each gated call is identified by a unique ``approval_id`` (minted with +``new_approval_id``) that the loop both registers here and echoes in the +``tool_start`` stream event. The frontend sends that exact id back, so a +stale or duplicate confirmation -- or a second tool awaiting a decision in +the same session -- can never resolve the wrong call. ``session_id`` is +kept alongside purely as a scope check. + +The slot is registered with ``begin_tool_decision`` *before* the loop +yields ``tool_start``, closing the race where a fast confirmation (or an +auto "Always allow") could otherwise arrive before the waiter exists. +``wait_tool_decision`` then blocks and cleans up its own slot. +""" + +import secrets +import threading +from typing import Optional + +# Generous ceiling so a user can deliberate; cancellation (stop button / +# disconnect) still breaks the wait early via ``cancel_event``. +_DECISION_TIMEOUT = 3600.0 + +# Fed to the model as the tool result when the user denies a call, so it +# can adapt and keep responding instead of the turn ending abruptly. +TOOL_REJECTED_MESSAGE = "The user declined to run this tool call." + +_lock = threading.Lock() +# approval_id -> {"event": threading.Event, "decision": str|None, "session": str} +_pending: dict[str, dict] = {} + + +def new_approval_id() -> str: + """Mint an unguessable id for one pending tool-call confirmation.""" + return secrets.token_urlsafe(16) + + +def begin_tool_decision(session_id, approval_id) -> dict: + """Register a pending decision slot and return it. + + Call this *before* yielding the ``tool_start`` event so the waiter + always exists by the time the user's confirmation can arrive. + """ + slot = { + "event": threading.Event(), + "decision": None, + "session": session_id or "", + } + with _lock: + _pending[approval_id] = slot + return slot + + +def wait_tool_decision(slot, approval_id, cancel_event = None, timeout = _DECISION_TIMEOUT): + """Block on a slot from ``begin_tool_decision`` until the user decides. + + Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait + times out or generation is cancelled before the user decides. Always + removes its own slot on exit. + """ + try: + waited = 0.0 + while not slot["event"].wait(timeout = 0.5): + if cancel_event is not None and cancel_event.is_set(): + return "deny" + waited += 0.5 + if waited >= timeout: + return "deny" + return slot["decision"] or "deny" + finally: + with _lock: + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) + + +def request_tool_decision( + session_id, approval_id, cancel_event = None, timeout = _DECISION_TIMEOUT +): + """Register and wait in one call (when the slot is not needed early).""" + slot = begin_tool_decision(session_id, approval_id) + return wait_tool_decision( + slot, approval_id, cancel_event = cancel_event, timeout = timeout + ) + + +def resolve_tool_decision(approval_id, decision, session_id = None) -> bool: + """Record the user's "allow"/"deny" decision and unblock the loop. + + Returns ``True`` if a pending call matched, ``False`` otherwise (e.g. a + stale or duplicate confirmation, or a session-scope mismatch). + """ + if not approval_id: + return False + with _lock: + slot = _pending.get(approval_id) + if not slot: + return False + if session_id is not None and slot["session"] != (session_id or ""): + return False + slot["decision"] = decision + slot["event"].set() + return True diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py new file mode 100644 index 0000000000..10db93ea90 --- /dev/null +++ b/studio/backend/tests/test_tool_approvals.py @@ -0,0 +1,224 @@ +# 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, + 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_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 + + +# ── 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() diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py new file mode 100644 index 0000000000..060d26cdef --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration tests for the confirmation gate inside the real tool loop. + +These drive ``run_safetensors_tool_loop`` (no model -- hand-crafted fake +generators) with ``confirm_tool_calls=True`` and resolve each pending +decision inline. The slot is registered before ``tool_start`` is yielded, +so resolving right after receiving that event always lands before the +loop blocks. Covers: allow executes once, deny skips execution and feeds +back the rejection, disabled/duplicate calls are not prompted, and a +denied call does not pollute duplicate detection. +""" + +import pytest + +from core.inference import safetensors_agentic +from core.inference.safetensors_agentic import run_safetensors_tool_loop +from core.inference.tool_call_parser import DUPLICATE_CALL_NUDGE +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision + +_SESSION = "loop-session" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _FakeExecuteTool: + def __init__(self): + self.calls = [] + + def __call__( + self, name, arguments, *, cancel_event = None, timeout = None, session_id = None + ): + self.calls.append((name, arguments)) + return f"RESULT[{name}]" + + +def _tool_call(name, args_json): + return f'{{"name": "{name}", "arguments": {args_json}}}' + + +def _multi_turn(turns): + """A single_turn generator that yields one full snapshot per turn.""" + turn_iter = iter(turns) + + def _gen(_messages): + try: + yield next(turn_iter) + except StopIteration: + return + + return _gen + + +_DEFAULT_TOOLS = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "web_search"}}, +] + + +def _drive(turns, decisions, *, tools = None): + """Run the loop, resolving each gated tool_start with the next decision. + + The advertised ``tools`` list drives the loop's enabled-tool filter + (pass a list omitting a tool to make a call to it "disabled"). + Returns (events, execute_calls). + """ + decision_iter = iter(decisions) + exec_fn = _FakeExecuteTool() + gen = run_safetensors_tool_loop( + single_turn = _multi_turn(turns), + messages = [{"role": "user", "content": "hi"}], + tools = _DEFAULT_TOOLS if tools is None else tools, + execute_tool = exec_fn, + session_id = _SESSION, + confirm_tool_calls = True, + ) + events = [] + for ev in gen: + events.append(ev) + if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): + # Slot is already registered (begin ran before this yield), so + # the decision lands before the loop enters its blocking wait. + resolve_tool_decision( + ev["approval_id"], next(decision_iter), session_id = _SESSION + ) + return events, exec_fn.calls + + +def _tool_starts(events): + return [e for e in events if e["type"] == "tool_start"] + + +def _tool_ends(events): + return [e for e in events if e["type"] == "tool_end"] + + +def test_allow_executes_the_tool_once(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["allow"], + ) + starts = _tool_starts(events) + assert len(starts) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert starts[0]["approval_id"] + assert calls == [("python", {"code": "print(1)"})] + assert _tool_ends(events)[0]["result"] == "RESULT[python]" + + +def test_deny_skips_execution_and_feeds_rejection(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["deny"], + ) + assert calls == [] # tool never ran + assert _tool_ends(events)[0]["result"] == TOOL_REJECTED_MESSAGE + + +def test_disabled_tool_is_not_prompted(): + # python is not advertised -> short-circuited, no approval asked. + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + [], # no decisions consumed + tools = [{"type": "function", "function": {"name": "web_search"}}], + ) + starts = _tool_starts(events) + assert starts[0]["awaiting_confirmation"] is False + assert starts[0]["approval_id"] == "" + assert calls == [] + assert "not enabled" in _tool_ends(events)[0]["result"] + + +def test_duplicate_call_is_not_prompted(): + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["allow"]) + starts = _tool_starts(events) + assert len(starts) == 2 + # First call gated + executed; second is a duplicate -> no prompt. + assert starts[0]["awaiting_confirmation"] is True + assert starts[1]["awaiting_confirmation"] is False + assert calls == [("python", {"code": "print(1)"})] + assert _tool_ends(events)[1]["result"] == DUPLICATE_CALL_NUDGE + + +def test_denied_call_can_be_reissued_and_approved(): + # Deny, then the model re-issues the identical call -> approving it must + # execute, not get suppressed as a duplicate (denied calls are not added + # to the duplicate-detection history). + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["deny", "allow"]) + starts = _tool_starts(events) + assert len(starts) == 2 + assert starts[0]["awaiting_confirmation"] is True + assert starts[1]["awaiting_confirmation"] is True # not treated as dup + assert calls == [("python", {"code": "print(1)"})] # ran once, on approve + ends = _tool_ends(events) + assert ends[0]["result"] == TOOL_REJECTED_MESSAGE + assert ends[1]["result"] == "RESULT[python]" diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py new file mode 100644 index 0000000000..986f0cb2fd --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""End-to-end handshake test for the tool-confirmation gate, no model. + +The real Studio stream wrappers in ``routes/inference.py`` drive the +synchronous agentic generator with ``await asyncio.to_thread(next, gen, +...)`` so the blocking ``threading.Event`` wait runs off the event loop. +This test rebuilds that exact pattern around the real +``state.tool_approvals`` functions, served by a real uvicorn process on +loopback (the same server Studio uses), and proves the load-bearing +property: + +* ``tool_start`` reaches the client before the gate blocks, and +* the separate ``/tool-confirm`` POST is served *while* the stream + connection is blocked, after which the stream resumes with the executed + (allow) or rejected (deny) result -- i.e. no deadlock. + +Each scenario runs under a socket-level timeout, so a regression that +reintroduces a deadlock fails fast instead of hanging the suite. +""" + +import asyncio +import json +import socket +import threading +import time + +import httpx +import pytest +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + begin_tool_decision, + new_approval_id, + resolve_tool_decision, + wait_tool_decision, +) + +_EXECUTED_RESULT = "tool executed: 2" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +def _build_app() -> FastAPI: + """Minimal app mirroring the real stream/confirm wiring.""" + app = FastAPI() + + def agentic_gen(session_id, cancel_event): + # Same shape as the real loops: register the approval slot, announce + # the call (echoing approval_id), gate on the decision, then either + # execute or feed back the rejection. + approval_id = new_approval_id() + slot = begin_tool_decision(session_id, approval_id) + yield { + "type": "tool_start", + "tool_name": "python", + "approval_id": approval_id, + "awaiting_confirmation": True, + } + denied = ( + wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny" + ) + result = TOOL_REJECTED_MESSAGE if denied else _EXECUTED_RESULT + yield {"type": "tool_end", "tool_name": "python", "result": result} + + @app.post("/stream") + async def stream(req: Request): + body = await req.json() + session_id = body.get("session_id") + cancel_event = threading.Event() + sentinel = object() + + async def wrapper(): + gen = agentic_gen(session_id, cancel_event) + while True: + event = await asyncio.to_thread(next, gen, sentinel) + if event is sentinel: + break + yield f"data: {json.dumps(event)}\n\n" + + return StreamingResponse(wrapper(), media_type = "text/event-stream") + + @app.post("/tool-confirm") + async def tool_confirm(req: Request): + body = await req.json() + resolved = resolve_tool_decision( + body.get("approval_id"), + body.get("decision"), + session_id = body.get("session_id"), + ) + return {"resolved": resolved} + + return app + + +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class _Server: + """Run a uvicorn server in a background thread for the test's lifetime.""" + + def __init__(self, app): + self.port = _free_port() + config = uvicorn.Config( + app, host = "127.0.0.1", port = self.port, log_level = "warning" + ) + self.server = uvicorn.Server(config) + self._thread = threading.Thread(target = self.server.run, daemon = True) + + def __enter__(self): + self._thread.start() + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + if self.server.started: + return self + time.sleep(0.02) + raise AssertionError("uvicorn did not start in time") + + def __exit__(self, *exc): + self.server.should_exit = True + self._thread.join(timeout = 10.0) + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +async def _gate_is_blocking(approval_id) -> None: + """Wait until the stream thread is parked on this approval's slot. + + The slot is registered before ``tool_start`` is yielded, so it exists + by the time the client receives the event -- exactly as in reality, + where the confirm POST only arrives after the card renders. + """ + for _ in range(400): + with tool_approvals._lock: + slot = tool_approvals._pending.get(approval_id) + if slot is not None and not slot["event"].is_set(): + return + await asyncio.sleep(0.005) + raise AssertionError("gate never started waiting") + + +async def _drive(base_url, session_id, decision): + events = [] + resolved = None + timeout = httpx.Timeout(10.0) + async with httpx.AsyncClient(base_url = base_url, timeout = timeout) as client: + async with client.stream( + "POST", "/stream", json = {"session_id": session_id} + ) as resp: + assert resp.status_code == 200 + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + event = json.loads(line[len("data: ") :]) + events.append(event) + if event["type"] == "tool_start": + # The stream is now blocked on the gate; the confirm + # POST (echoing approval_id) must still be served over a + # second connection. + approval_id = event["approval_id"] + await _gate_is_blocking(approval_id) + r = await client.post( + "/tool-confirm", + json = { + "session_id": session_id, + "approval_id": approval_id, + "decision": decision, + }, + ) + resolved = r.json()["resolved"] + return events, resolved + + +def _run(session_id, decision): + with _Server(_build_app()) as srv: + return asyncio.run( + asyncio.wait_for(_drive(srv.base_url, session_id, decision), timeout = 15.0) + ) + + +def _types(events): + return [e["type"] for e in events] + + +def test_allow_resumes_stream_with_executed_result(): + events, resolved = _run("sess-allow", "allow") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == _EXECUTED_RESULT + + +def test_deny_resumes_stream_with_rejection_result(): + events, resolved = _run("sess-deny", "deny") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == TOOL_REJECTED_MESSAGE + + +def test_tool_start_precedes_the_block_and_carries_approval_id(): + # The first streamed event is always tool_start, proving the buttons + # can render before the backend pauses for the decision -- and it + # carries the approval_id / awaiting_confirmation the UI needs. + events, _ = _run("sess-order", "allow") + assert events[0]["type"] == "tool_start" + assert events[0]["awaiting_confirmation"] is True + assert events[0]["approval_id"] diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 51cf000863..f3b0e31d9d 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -20,6 +20,7 @@ import { thinkEffortAriaLabel, thinkToggleAriaLabel, } from "@/components/assistant-ui/think-aria-label"; +import { ToolConfirmationControls } from "@/components/assistant-ui/tool-confirmation-controls"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution"; @@ -62,6 +63,7 @@ import { ErrorPrimitive, MessagePrimitive, ThreadPrimitive, + type ToolCallMessagePartComponent, useAui, useAuiEvent, useAuiState, @@ -1293,6 +1295,39 @@ const CancelledIndicator: FC = () => { ); }; +// Render Allow / Always allow / Deny controls under every tool card so the +// "Confirm tool calls" gate works for the built-in tools (search, python, +// terminal, code, image) too -- not just the MCP tools that use the +// fallback renderer. The controls no-op unless the adapter registered a +// backend-gated pending call for this card, so non-gated tools are +// unaffected. Wrapped once at module scope to keep stable component +// identities (inline wrapping would remount the tool subtree each render). +const withToolConfirmation = ( + Component: ToolCallMessagePartComponent, +): ToolCallMessagePartComponent => { + const WithToolConfirmation: ToolCallMessagePartComponent = (props) => ( + <> + + + + ); + return WithToolConfirmation; +}; + +const WebSearchToolUIConfirmable = withToolConfirmation(WebSearchToolUI); +const PythonToolUIConfirmable = withToolConfirmation(PythonToolUI); +const TerminalToolUIConfirmable = withToolConfirmation(TerminalToolUI); +const CodeExecutionToolUIConfirmable = withToolConfirmation(CodeExecutionToolUI); +const ImageGenerationToolUIConfirmable = withToolConfirmation( + ImageGenerationToolUI, +); +const ToolFallbackConfirmable = withToolConfirmation(ToolFallback); + const AssistantMessage: FC = () => { return ( { ToolGroup: ToolGroup, tools: { by_name: { - web_search: WebSearchToolUI, - python: PythonToolUI, - terminal: TerminalToolUI, - code_execution: CodeExecutionToolUI, - image_generation: ImageGenerationToolUI, + web_search: WebSearchToolUIConfirmable, + python: PythonToolUIConfirmable, + terminal: TerminalToolUIConfirmable, + code_execution: CodeExecutionToolUIConfirmable, + image_generation: ImageGenerationToolUIConfirmable, }, - Fallback: ToolFallback, + Fallback: ToolFallbackConfirmable, }, }} /> diff --git a/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx b/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx new file mode 100644 index 0000000000..594d5db46e --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { Button } from "@/components/ui/button"; +import { resolveToolConfirmation } from "@/features/chat/api/chat-api"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; +import type { ToolCallMessagePartStatus } from "@assistant-ui/react"; +import { useCallback, useEffect, useState } from "react"; + +/** + * Allow / Always allow / Deny controls for a tool call paused awaiting the + * user's confirmation. Rendered alongside every tool card (built-in and + * MCP) so the gate works for all tools, not just the ones using the + * fallback renderer. + * + * A card is "awaiting" only when the adapter registered a backend-gated + * pending call for it (see `toolConfirmations` in the runtime store), so + * non-gated cards -- toggle off, or external-provider tools that already + * ran -- never show controls. + */ +export function ToolConfirmationControls({ + toolCallId, + toolName, + result, + status, +}: { + toolCallId?: string; + toolName: string; + result: unknown; + status?: ToolCallMessagePartStatus; +}) { + const confirmation = useChatRuntimeStore((s) => + toolCallId ? s.toolConfirmations[toolCallId] : undefined, + ); + const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways); + const clearToolConfirmation = useChatRuntimeStore( + (s) => s.clearToolConfirmation, + ); + const sessionId = confirmation?.sessionId ?? ""; + const autoAllowed = useChatRuntimeStore( + (s) => s.alwaysAllowToolsBySession.get(sessionId)?.has(toolName) ?? false, + ); + + const [decided, setDecided] = useState(false); + const [pending, setPending] = useState<"allow" | "deny" | null>(null); + const [failed, setFailed] = useState(false); + + // Still awaiting our decision: a gated pending entry exists, the tool has + // not produced a result, and the card is in its running state. + const awaiting = + confirmation !== undefined && + result === undefined && + status?.type === "running"; + const showControls = awaiting && !decided; + + const resolve = useCallback( + async (decision: "allow" | "deny") => { + if (!toolCallId || !confirmation) return; + setPending(decision); + setFailed(false); + try { + const ok = await resolveToolConfirmation( + confirmation.sessionId, + confirmation.approvalId, + decision, + ); + if (ok) { + // Only hide the controls once the backend confirms it matched the + // pending call -- otherwise the generation would stay blocked with + // no way to retry. + setDecided(true); + clearToolConfirmation(toolCallId); + } else { + setFailed(true); + } + } catch { + setFailed(true); + } finally { + setPending(null); + } + }, + [toolCallId, confirmation, clearToolConfirmation], + ); + + // Tools the user marked "Always allow" (this session) approve themselves. + useEffect(() => { + if (showControls && autoAllowed && pending === null && !failed) { + void resolve("allow"); + } + }, [showControls, autoAllowed, pending, failed, resolve]); + + if (!showControls) return null; + // Auto-approved tools resolve silently unless the post fails. + if (autoAllowed && !failed) return null; + + return ( +
+ + + + {failed ? ( + + Could not send your decision. Try again. + + ) : null} +
+ ); +} diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index 2ea87d63f3..b12855dfb8 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -319,6 +319,9 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ result, status, }) => { + // Allow/Deny confirmation controls are rendered uniformly for every tool + // card (built-in and fallback) by the `withToolConfirmation` wrapper in + // thread.tsx, so this renderer stays purely presentational. const isCancelled = status?.type === "incomplete" && status.reason === "cancelled"; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 666a304c3e..764468d4a0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1304,6 +1304,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { codeToolsEnabled, imageToolsEnabled, mcpEnabledForChat, + confirmToolCalls, webFetchToolsEnabled, } = runtime; const externalSelection = parseExternalModelId(params.checkpoint); @@ -2101,6 +2102,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(codeToolsEnabled ? ["python", "terminal"] : []), ], mcp_enabled: mcpEnabledForChat, + confirm_tool_calls: confirmToolCalls, auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, max_tool_calls_per_message: @@ -2202,6 +2204,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (toolEvent.type === "tool_start") { const id = (toolEvent.tool_call_id as string) || + (toolEvent.approval_id as string) || `${toolEvent.tool_name}_${Date.now()}`; const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"]; @@ -2212,11 +2215,28 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { argsText: JSON.stringify(toolArgs), args: toolArgs, }); + // Backend-gated tool calls pause for an allow/deny. Record + // the approval id + the session the generation runs under + // (the same id sent as session_id) so the tool card can + // resolve the exact pending call. Non-gated calls (toggle + // off, external providers) never set this, so their cards + // show no approval controls. + if (toolEvent.awaiting_confirmation === true) { + useChatRuntimeStore + .getState() + .setToolConfirmation( + id, + (toolEvent.approval_id as string) || "", + resolvedThreadId ?? "", + ); + } } else if (toolEvent.type === "tool_end") { const id = (toolEvent.tool_call_id as string) || toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; + // The call resolved; drop any pending confirmation entry. + useChatRuntimeStore.getState().clearToolConfirmation(id); const idx = toolCallParts.findIndex( (p) => p.toolCallId === id, ); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 81303d9311..2f00e42361 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -105,6 +105,31 @@ export async function unloadModel(payload: UnloadModelRequest): Promise { await parseJsonOrThrow(response); } +/** + * Allow or deny a tool call that is paused awaiting user confirmation + * (when the "Confirm tool calls" toggle is on). The call is identified by + * the backend ``approvalId`` echoed in the tool_start event; ``sessionId`` + * is a scope check. Resolves to ``true`` only when the backend matched a + * pending call, so the caller can surface a retry on a stale/failed post. + */ +export async function resolveToolConfirmation( + sessionId: string, + approvalId: string, + decision: "allow" | "deny", +): Promise { + const response = await authFetch("/api/inference/tool-confirm", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + session_id: sessionId, + approval_id: approvalId, + decision, + }), + }); + const parsed = await parseJsonOrThrow<{ resolved?: boolean }>(response); + return parsed.resolved === true; +} + export interface CachedGgufRepo { repo_id: string; size_bytes: number; diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 8c852b8189..4f19e95a88 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1338,6 +1338,7 @@ export function ChatSettingsPanel({
+
@@ -1512,6 +1513,30 @@ function AutoHealToolCallsToggle() { ); } +function ConfirmToolCallsToggle() { + const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); + const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); + + return ( +
+
+ + Confirm tool calls + + + When on, every tool call pauses for your approval in the chat before + it runs. + +
+ +
+ ); +} + function McpServersSection() { const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat); const setMcpEnabledForChat = useChatRuntimeStore( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 8f0767b7b8..9b85e56e9d 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -28,6 +28,7 @@ export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled"; +export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -301,6 +302,26 @@ type ChatRuntimeStore = { codeToolsEnabled: boolean; imageToolsEnabled: boolean; mcpEnabledForChat: boolean; + /** + * When on, every tool call pauses for an explicit allow/deny in the + * chat before it runs. + */ + confirmToolCalls: boolean; + /** + * Per-session set of tool names the user chose to auto-approve via + * "Always allow". Keyed by thread/session id so allowing a tool in one + * chat does not silently auto-approve it in another. Not persisted + * across reloads. + */ + alwaysAllowToolsBySession: Map>; + /** + * Tool calls currently paused awaiting the user's allow/deny decision, + * keyed by the frontend tool-call id. Each entry carries the backend + * ``approvalId`` to echo back and the ``sessionId`` the generation runs + * under, so the confirmation always resolves the exact pending call. + * Only backend-gated local tool calls are added here. + */ + toolConfirmations: Record; /** * Fetch pill state, independent of `toolsEnabled` (Search). Only * consulted when `providerSupportsBuiltinWebFetch` is true. @@ -369,6 +390,14 @@ type ChatRuntimeStore = { setCodeToolsEnabled: (enabled: boolean) => void; setImageToolsEnabled: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; + setConfirmToolCalls: (enabled: boolean) => void; + allowToolAlways: (sessionId: string, toolName: string) => void; + setToolConfirmation: ( + toolCallId: string, + approvalId: string, + sessionId: string, + ) => void; + clearToolConfirmation: (toolCallId: string) => void; setWebFetchToolsEnabled: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; @@ -620,6 +649,9 @@ export const useChatRuntimeStore = create((set, get) => ({ codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false), + confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false), + alwaysAllowToolsBySession: new Map>(), + toolConfirmations: {}, webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, @@ -901,6 +933,33 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_MCP_ENABLED_KEY, mcpEnabledForChat); return { mcpEnabledForChat }; }), + setConfirmToolCalls: (confirmToolCalls) => + set(() => { + saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls); + return { confirmToolCalls }; + }), + allowToolAlways: (sessionId, toolName) => + set((state) => { + const current = state.alwaysAllowToolsBySession.get(sessionId); + if (current?.has(toolName)) return state; + const next = new Map(state.alwaysAllowToolsBySession); + next.set(sessionId, new Set(current ?? []).add(toolName)); + return { alwaysAllowToolsBySession: next }; + }), + setToolConfirmation: (toolCallId, approvalId, sessionId) => + set((state) => ({ + toolConfirmations: { + ...state.toolConfirmations, + [toolCallId]: { approvalId, sessionId }, + }, + })), + clearToolConfirmation: (toolCallId) => + set((state) => { + if (!(toolCallId in state.toolConfirmations)) return state; + const next = { ...state.toolConfirmations }; + delete next[toolCallId]; + return { toolConfirmations: next }; + }), setWebFetchToolsEnabled: (webFetchToolsEnabled) => set(() => { saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled);