From 282b8f9a9673de03ff3fe267855b9abf2372d851 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 29 May 2026 11:55:43 -0700 Subject: [PATCH 1/6] Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls --- studio/backend/core/inference/llama_cpp.py | 18 ++++- studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 16 +++- studio/backend/models/inference.py | 9 +++ studio/backend/routes/inference.py | 19 +++++ studio/backend/state/tool_approvals.py | 78 +++++++++++++++++++ .../components/assistant-ui/tool-fallback.tsx | 59 +++++++++++++- .../src/features/chat/api/chat-adapter.ts | 2 + .../src/features/chat/api/chat-api.ts | 17 ++++ .../src/features/chat/chat-settings-sheet.tsx | 18 +++++ .../chat/stores/chat-runtime-store.ts | 26 +++++++ 11 files changed, 259 insertions(+), 5 deletions(-) create mode 100644 studio/backend/state/tool_approvals.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2d95112d6d..8c9f67b188 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4452,6 +4452,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. @@ -4462,6 +4463,10 @@ 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, + request_tool_decision, + ) if not self.is_loaded: raise RuntimeError("llama-server is not loaded") @@ -5085,7 +5090,12 @@ class LlamaCppBackend: # 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]: + _denied = confirm_tool_calls and request_tool_decision( + session_id, cancel_event=cancel_event + ) == "deny" + if _denied: + result = TOOL_REJECTED_MESSAGE + elif _prev and _prev[0] == _tc_key and not _prev[1]: result = ( "You already made this exact call. " "Do not repeat the same tool call. " @@ -5144,7 +5154,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..157061e34e 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -24,6 +24,8 @@ from urllib.parse import urlparse from loggers import get_logger +from state.tool_approvals import TOOL_REJECTED_MESSAGE, request_tool_decision + from core.inference.tool_call_parser import ( BUDGET_EXHAUSTED_NUDGE, DUPLICATE_CALL_NUDGE, @@ -105,6 +107,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. @@ -317,7 +320,12 @@ def run_safetensors_tool_loop( } tc_key = tool_name + str(arguments) - if allowed_tool_names and tool_name not in allowed_tool_names: + denied = confirm_tool_calls and request_tool_decision( + session_id, cancel_event=cancel_event + ) == "deny" + if denied: + result = TOOL_REJECTED_MESSAGE + elif allowed_tool_names and tool_name not in allowed_tool_names: result = ( f"Error: tool '{tool_name}' is not enabled for this " "request. Use one of the enabled tools or provide a " @@ -355,7 +363,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..bebbee3fb0 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,11 @@ class ChatCompletionRequest(BaseModel): return self +class ToolConfirmRequest(BaseModel): + session_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..407099dfbc 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,22 @@ 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. + + Returns {"resolved": bool}. ``False`` means no matching call was + waiting (e.g. a stale or duplicate confirmation). + """ + from state.tool_approvals import resolve_tool_decision + + resolved = resolve_tool_decision(body.session_id, body.decision) + return {"resolved": resolved} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, @@ -2790,6 +2807,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 +3328,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..f187e630b9 --- /dev/null +++ b/studio/backend/state/tool_approvals.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Per-session 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. + +The agentic loop is strictly sequential, so at most one tool awaits a +decision per session at any moment -- the gate keys on ``session_id`` +alone and never needs to match individual tool-call ids (which some +models omit). +""" + +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() +# session_key -> {"event": threading.Event, "decision": "allow"|"deny"|None} +_pending: dict[str, dict] = {} + + +def _key(session_id: Optional[str]) -> str: + return session_id or "" + + +def request_tool_decision(session_id, cancel_event=None, timeout=_DECISION_TIMEOUT): + """Block until the user allows/denies the pending tool call. + + Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait + times out or generation is cancelled before the user decides. + """ + key = _key(session_id) + event = threading.Event() + with _lock: + _pending[key] = {"event": event, "decision": None} + try: + waited = 0.0 + while not 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" + with _lock: + slot = _pending.get(key) + return (slot or {}).get("decision") or "deny" + finally: + with _lock: + # Only remove our own entry — never a newer waiter's, in case the + # same session somehow re-registered while we were waiting. + if _pending.get(key, {}).get("event") is event: + _pending.pop(key, None) + + +def resolve_tool_decision(session_id, decision) -> 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). + """ + key = _key(session_id) + with _lock: + slot = _pending.get(key) + if not slot: + return False + slot["decision"] = decision + slot["event"].set() + return True diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index 2ea87d63f3..09ce87a97c 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -3,11 +3,14 @@ "use client"; +import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger, } from "@/components/ui/collapsible"; +import { resolveToolConfirmation } from "@/features/chat/api/chat-api"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock"; import { cn } from "@/lib/utils"; import { @@ -27,6 +30,7 @@ import { type ElementType, memo, useCallback, + useEffect, useRef, useState, } from "react"; @@ -319,11 +323,40 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ result, status, }) => { + const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); + const alwaysAllowTools = useChatRuntimeStore((s) => s.alwaysAllowTools); + const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways); + const [decided, setDecided] = useState(false); + + // A live tool call with no result yet, while confirmation is on, is one + // the backend has paused awaiting our decision (the loop is sequential, + // so only ever one at a time — the backend gates on the session alone). + const awaiting = + confirmToolCalls && result === undefined && status?.type === "running"; + const autoAllowed = alwaysAllowTools.has(toolName); + const showConfirm = awaiting && !decided; + + const resolve = useCallback((decision: "allow" | "deny") => { + setDecided(true); + // Falls back to "" so a thread that started before it had an id (matching + // the backend's empty-session gate key) still gets unblocked. + const sessionId = useChatRuntimeStore.getState().activeThreadId ?? ""; + void resolveToolConfirmation(sessionId, decision).catch(() => {}); + }, []); + + // Tools the user marked "Always allow" approve themselves this session. + useEffect(() => { + if (showConfirm && autoAllowed) resolve("allow"); + }, [showConfirm, autoAllowed, resolve]); + const isCancelled = status?.type === "incomplete" && status.reason === "cancelled"; return ( - + @@ -331,6 +364,30 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ argsText={argsText} className={cn(isCancelled && "opacity-60")} /> + {showConfirm && !autoAllowed && ( +
+ + + +
+ )} {!isCancelled && }
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 666a304c3e..368299aedc 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: diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 81303d9311..24364be138 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -105,6 +105,23 @@ 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 backend gates on the + * session id alone, since at most one call awaits a decision per thread. + */ +export async function resolveToolConfirmation( + sessionId: 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, decision }), + }); + await parseJsonOrThrow(response); +} + 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..8eda82fb02 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1517,6 +1517,8 @@ function McpServersSection() { const setMcpEnabledForChat = useChatRuntimeStore( (s) => s.setMcpEnabledForChat, ); + const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); + const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); const [enabledServerCount, setEnabledServerCount] = useState( null, ); @@ -1557,6 +1559,22 @@ function McpServersSection() { disabled={enabledServerCount === 0 && !mcpEnabledForChat} /> +
+
+ + Confirm tool calls + + + When on, every tool call pauses for your approval in the chat + before it runs. + +
+ +
{enabledServerCount === null 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..b677e999ed 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,16 @@ 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; + /** + * Tool names the user chose to auto-approve for the rest of this + * session via "Always allow". Not persisted across reloads. + */ + alwaysAllowTools: Set; /** * Fetch pill state, independent of `toolsEnabled` (Search). Only * consulted when `providerSupportsBuiltinWebFetch` is true. @@ -369,6 +380,8 @@ type ChatRuntimeStore = { setCodeToolsEnabled: (enabled: boolean) => void; setImageToolsEnabled: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; + setConfirmToolCalls: (enabled: boolean) => void; + allowToolAlways: (toolName: string) => void; setWebFetchToolsEnabled: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; @@ -620,6 +633,8 @@ 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), + alwaysAllowTools: new Set(), webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, @@ -901,6 +916,17 @@ 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: (toolName) => + set((state) => + state.alwaysAllowTools.has(toolName) + ? state + : { alwaysAllowTools: new Set(state.alwaysAllowTools).add(toolName) }, + ), setWebFetchToolsEnabled: (webFetchToolsEnabled) => set(() => { saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled); From f6f3b48a7c60c4c223fa2b002fdb0529c8383b85 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 29 May 2026 19:10:37 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 8 +++++--- studio/backend/core/inference/safetensors_agentic.py | 8 +++++--- studio/backend/state/tool_approvals.py | 4 ++-- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8c9f67b188..fb410f558c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -5090,9 +5090,11 @@ class LlamaCppBackend: # 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 - _denied = confirm_tool_calls and request_tool_decision( - session_id, cancel_event=cancel_event - ) == "deny" + _denied = ( + confirm_tool_calls + and request_tool_decision(session_id, cancel_event = cancel_event) + == "deny" + ) if _denied: result = TOOL_REJECTED_MESSAGE elif _prev and _prev[0] == _tc_key and not _prev[1]: diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 157061e34e..7f1a2273fc 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -320,9 +320,11 @@ def run_safetensors_tool_loop( } tc_key = tool_name + str(arguments) - denied = confirm_tool_calls and request_tool_decision( - session_id, cancel_event=cancel_event - ) == "deny" + denied = ( + confirm_tool_calls + and request_tool_decision(session_id, cancel_event = cancel_event) + == "deny" + ) if denied: result = TOOL_REJECTED_MESSAGE elif allowed_tool_names and tool_name not in allowed_tool_names: diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py index f187e630b9..cd376be24d 100644 --- a/studio/backend/state/tool_approvals.py +++ b/studio/backend/state/tool_approvals.py @@ -33,7 +33,7 @@ def _key(session_id: Optional[str]) -> str: return session_id or "" -def request_tool_decision(session_id, cancel_event=None, timeout=_DECISION_TIMEOUT): +def request_tool_decision(session_id, cancel_event = None, timeout = _DECISION_TIMEOUT): """Block until the user allows/denies the pending tool call. Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait @@ -45,7 +45,7 @@ def request_tool_decision(session_id, cancel_event=None, timeout=_DECISION_TIMEO _pending[key] = {"event": event, "decision": None} try: waited = 0.0 - while not event.wait(timeout=0.5): + while not event.wait(timeout = 0.5): if cancel_event is not None and cancel_event.is_set(): return "deny" waited += 0.5 From a758b4ac9ef79d75a566d0d9427510a10e4d79fb Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 29 May 2026 12:19:07 -0700 Subject: [PATCH 3/6] Fix race in tool-call confirmation gate --- studio/backend/state/tool_approvals.py | 33 ++++++++++++++++---------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py index cd376be24d..59f6264362 100644 --- a/studio/backend/state/tool_approvals.py +++ b/studio/backend/state/tool_approvals.py @@ -7,10 +7,12 @@ 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. -The agentic loop is strictly sequential, so at most one tool awaits a -decision per session at any moment -- the gate keys on ``session_id`` -alone and never needs to match individual tool-call ids (which some -models omit). +The agentic loop is sequential, so a session normally has a single tool +awaiting a decision -- the gate keys on ``session_id`` alone and never +needs to match individual tool-call ids (which some models omit). If a +second waiter ever registers for the same key (e.g. two id-less chats +both mapping to ""), the older one is unblocked as denied so it can't +hang. """ import threading @@ -40,25 +42,30 @@ def request_tool_decision(session_id, cancel_event = None, timeout = _DECISION_T times out or generation is cancelled before the user decides. """ key = _key(session_id) - event = threading.Event() + slot = {"event": threading.Event(), "decision": None} with _lock: - _pending[key] = {"event": event, "decision": None} + # If a waiter already holds this key (same session, or two id-less + # chats both mapping to ""), unblock it as denied so it can't hang + # once we orphan its event below. + old = _pending.get(key) + if old is not None: + old["decision"] = "deny" + old["event"].set() + _pending[key] = slot try: waited = 0.0 - while not event.wait(timeout = 0.5): + 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" - with _lock: - slot = _pending.get(key) - return (slot or {}).get("decision") or "deny" + # Read our own slot, not _pending[key], which a newer waiter may + # have replaced. + return slot["decision"] or "deny" finally: with _lock: - # Only remove our own entry — never a newer waiter's, in case the - # same session somehow re-registered while we were waiting. - if _pending.get(key, {}).get("event") is event: + if _pending.get(key) is slot: _pending.pop(key, None) From cfca72ce339e9060ce7a0eb6fe80b4572a108a7c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 05:32:01 +0000 Subject: [PATCH 4/6] Studio: gate built-in tool calls and harden the confirmation handshake The Allow / Always allow / Deny controls only lived in the fallback tool card, but the built-in tools (web search, python, terminal, code execution, image generation) render with their own components and so never showed the buttons. Those calls paused after tool_start with no way to approve them, hanging until the 1 hour timeout. Only MCP tools, which use the fallback renderer, actually worked. Render the controls for every tool card by wrapping each registered tool component (and the fallback) in thread.tsx with a shared ToolConfirmationControls, so the gate applies uniformly. Also make the handshake robust: - The gate keys on a per-call approval_id minted by the backend and echoed in tool_start, instead of session_id alone, so a stale or concurrent confirmation can no longer resolve the wrong call. - The approval slot is registered before tool_start is yielded, closing the race where a fast click or an auto "Always allow" could reach the backend before the waiter existed. - The frontend resolves with the same session id the request was sent with (plus the approval_id), fixing the new-thread mismatch where the confirmation targeted a different session than the blocked stream. - The confirm endpoint returns {resolved}; the UI keeps the buttons and shows a retry hint until the backend confirms a match, instead of hiding them on a failed or mistargeted post. - The gate runs after the disabled-tool and duplicate-call checks, so a call that will not execute is not put up for approval. A denied call is still excluded from duplicate detection, so re-issuing and approving it works. - "Always allow" is scoped per session to match the backend gate. Add backend tests for the approval registry, the SSE no-deadlock handshake, and the loop integration (allow, deny, disabled, duplicate, re-issue after deny). --- studio/backend/core/inference/llama_cpp.py | 92 ++++--- .../core/inference/safetensors_agentic.py | 50 ++-- studio/backend/models/inference.py | 1 + studio/backend/routes/inference.py | 10 +- studio/backend/state/tool_approvals.py | 90 ++++--- studio/backend/tests/test_tool_approvals.py | 224 +++++++++++++++++ .../backend/tests/test_tool_confirm_loop.py | 165 +++++++++++++ .../backend/tests/test_tool_confirm_stream.py | 225 ++++++++++++++++++ .../src/components/assistant-ui/thread.tsx | 47 +++- .../tool-confirmation-controls.tsx | 133 +++++++++++ .../components/assistant-ui/tool-fallback.tsx | 62 +---- .../src/features/chat/api/chat-adapter.ts | 18 ++ .../src/features/chat/api/chat-api.ts | 18 +- .../chat/stores/chat-runtime-store.ts | 55 ++++- 14 files changed, 1028 insertions(+), 162 deletions(-) create mode 100644 studio/backend/tests/test_tool_approvals.py create mode 100644 studio/backend/tests/test_tool_confirm_loop.py create mode 100644 studio/backend/tests/test_tool_confirm_stream.py create mode 100644 studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fb410f558c..80ead53960 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4465,7 +4465,9 @@ class LlamaCppBackend: from core.inference.tools import execute_tool from state.tool_approvals import ( TOOL_REJECTED_MESSAGE, - request_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, ) if not self.is_loaded: @@ -5077,27 +5079,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 - _denied = ( - confirm_tool_calls - and request_tool_decision(session_id, cancel_event = cancel_event) - == "deny" + _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 ) - if _denied: - result = TOOL_REJECTED_MESSAGE - elif _prev and _prev[0] == _tc_key and not _prev[1]: + _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. " @@ -5106,27 +5131,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, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 7f1a2273fc..8a536a4bfd 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -24,7 +24,12 @@ from urllib.parse import urlparse from loggers import get_logger -from state.tool_approvals import TOOL_REJECTED_MESSAGE, request_tool_decision +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, @@ -311,34 +316,51 @@ 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) - denied = ( - confirm_tool_calls - and request_tool_decision(session_id, cancel_event = cancel_event) - == "deny" - ) - if denied: - result = TOOL_REJECTED_MESSAGE - elif 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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index bebbee3fb0..d6f24b5159 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -951,6 +951,7 @@ class ChatCompletionRequest(BaseModel): class ToolConfirmRequest(BaseModel): session_id: Optional[str] = None + approval_id: Optional[str] = None decision: Literal["allow", "deny"] = "deny" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 407099dfbc..d87875ca01 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1268,12 +1268,16 @@ async def confirm_tool_call( ): """Allow or deny a tool call awaiting user confirmation. - Returns {"resolved": bool}. ``False`` means no matching call was - waiting (e.g. a stale or duplicate 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.session_id, body.decision) + resolved = resolve_tool_decision( + body.approval_id, body.decision, session_id = body.session_id + ) return {"resolved": resolved} diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py index 59f6264362..b3b9e212d6 100644 --- a/studio/backend/state/tool_approvals.py +++ b/studio/backend/state/tool_approvals.py @@ -1,20 +1,26 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Per-session tool-call confirmation gate. +"""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. -The agentic loop is sequential, so a session normally has a single tool -awaiting a decision -- the gate keys on ``session_id`` alone and never -needs to match individual tool-call ids (which some models omit). If a -second waiter ever registers for the same key (e.g. two id-less chats -both mapping to ""), the older one is unblocked as denied so it can't -hang. +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 @@ -27,31 +33,40 @@ _DECISION_TIMEOUT = 3600.0 TOOL_REJECTED_MESSAGE = "The user declined to run this tool call." _lock = threading.Lock() -# session_key -> {"event": threading.Event, "decision": "allow"|"deny"|None} +# approval_id -> {"event": threading.Event, "decision": str|None, "session": str} _pending: dict[str, dict] = {} -def _key(session_id: Optional[str]) -> str: - return session_id or "" +def new_approval_id() -> str: + """Mint an unguessable id for one pending tool-call confirmation.""" + return secrets.token_urlsafe(16) -def request_tool_decision(session_id, cancel_event = None, timeout = _DECISION_TIMEOUT): - """Block until the user allows/denies the pending tool call. +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. + times out or generation is cancelled before the user decides. Always + removes its own slot on exit. """ - key = _key(session_id) - slot = {"event": threading.Event(), "decision": None} - with _lock: - # If a waiter already holds this key (same session, or two id-less - # chats both mapping to ""), unblock it as denied so it can't hang - # once we orphan its event below. - old = _pending.get(key) - if old is not None: - old["decision"] = "deny" - old["event"].set() - _pending[key] = slot try: waited = 0.0 while not slot["event"].wait(timeout = 0.5): @@ -60,26 +75,37 @@ def request_tool_decision(session_id, cancel_event = None, timeout = _DECISION_T waited += 0.5 if waited >= timeout: return "deny" - # Read our own slot, not _pending[key], which a newer waiter may - # have replaced. return slot["decision"] or "deny" finally: with _lock: - if _pending.get(key) is slot: - _pending.pop(key, None) + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) -def resolve_tool_decision(session_id, decision) -> bool: +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). + stale or duplicate confirmation, or a session-scope mismatch). """ - key = _key(session_id) + if not approval_id: + return False with _lock: - slot = _pending.get(key) + 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..4855035435 --- /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..f3a8d7194f --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -0,0 +1,165 @@ +# 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 09ce87a97c..b12855dfb8 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -3,14 +3,11 @@ "use client"; -import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger, } from "@/components/ui/collapsible"; -import { resolveToolConfirmation } from "@/features/chat/api/chat-api"; -import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock"; import { cn } from "@/lib/utils"; import { @@ -30,7 +27,6 @@ import { type ElementType, memo, useCallback, - useEffect, useRef, useState, } from "react"; @@ -323,40 +319,14 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ result, status, }) => { - const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); - const alwaysAllowTools = useChatRuntimeStore((s) => s.alwaysAllowTools); - const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways); - const [decided, setDecided] = useState(false); - - // A live tool call with no result yet, while confirmation is on, is one - // the backend has paused awaiting our decision (the loop is sequential, - // so only ever one at a time — the backend gates on the session alone). - const awaiting = - confirmToolCalls && result === undefined && status?.type === "running"; - const autoAllowed = alwaysAllowTools.has(toolName); - const showConfirm = awaiting && !decided; - - const resolve = useCallback((decision: "allow" | "deny") => { - setDecided(true); - // Falls back to "" so a thread that started before it had an id (matching - // the backend's empty-session gate key) still gets unblocked. - const sessionId = useChatRuntimeStore.getState().activeThreadId ?? ""; - void resolveToolConfirmation(sessionId, decision).catch(() => {}); - }, []); - - // Tools the user marked "Always allow" approve themselves this session. - useEffect(() => { - if (showConfirm && autoAllowed) resolve("allow"); - }, [showConfirm, autoAllowed, resolve]); - + // 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"; return ( - + @@ -364,30 +334,6 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ argsText={argsText} className={cn(isCancelled && "opacity-60")} /> - {showConfirm && !autoAllowed && ( -
- - - -
- )} {!isCancelled && }
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 368299aedc..764468d4a0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2204,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"]; @@ -2214,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 24364be138..2f00e42361 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -107,19 +107,27 @@ export async function unloadModel(payload: UnloadModelRequest): Promise { /** * Allow or deny a tool call that is paused awaiting user confirmation - * (when the "Confirm tool calls" toggle is on). The backend gates on the - * session id alone, since at most one call awaits a decision per thread. + * (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 { +): Promise { const response = await authFetch("/api/inference/tool-confirm", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ session_id: sessionId, decision }), + body: JSON.stringify({ + session_id: sessionId, + approval_id: approvalId, + decision, + }), }); - await parseJsonOrThrow(response); + const parsed = await parseJsonOrThrow<{ resolved?: boolean }>(response); + return parsed.resolved === true; } export interface CachedGgufRepo { 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 b677e999ed..9b85e56e9d 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -308,10 +308,20 @@ type ChatRuntimeStore = { */ confirmToolCalls: boolean; /** - * Tool names the user chose to auto-approve for the rest of this - * session via "Always allow". Not persisted across reloads. + * 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. */ - alwaysAllowTools: Set; + 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. @@ -381,7 +391,13 @@ type ChatRuntimeStore = { setImageToolsEnabled: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; setConfirmToolCalls: (enabled: boolean) => void; - allowToolAlways: (toolName: string) => 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; @@ -634,7 +650,8 @@ export const useChatRuntimeStore = create((set, get) => ({ imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false), confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false), - alwaysAllowTools: new Set(), + alwaysAllowToolsBySession: new Map>(), + toolConfirmations: {}, webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, @@ -921,12 +938,28 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls); return { confirmToolCalls }; }), - allowToolAlways: (toolName) => - set((state) => - state.alwaysAllowTools.has(toolName) - ? state - : { alwaysAllowTools: new Set(state.alwaysAllowTools).add(toolName) }, - ), + 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); From f1451fcb71ea024b46fb23ac15c82443f33bb35e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 31 May 2026 05:32:41 +0000 Subject: [PATCH 5/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/safetensors_agentic.py | 8 ++++++-- studio/backend/state/tool_approvals.py | 4 +--- studio/backend/tests/test_tool_approvals.py | 8 ++++---- studio/backend/tests/test_tool_confirm_loop.py | 4 +++- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 8a536a4bfd..24b598a2a6 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -317,7 +317,9 @@ def run_safetensors_tool_loop( ) tc_key = tool_name + str(arguments) - is_disabled = bool(allowed_tool_names) and tool_name not in allowed_tool_names + 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 ) @@ -326,7 +328,9 @@ def run_safetensors_tool_loop( # 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 + 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 diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py index b3b9e212d6..1fc7981eb7 100644 --- a/studio/backend/state/tool_approvals.py +++ b/studio/backend/state/tool_approvals.py @@ -58,9 +58,7 @@ def begin_tool_decision(session_id, approval_id) -> dict: return slot -def wait_tool_decision( - slot, approval_id, cancel_event = None, timeout = _DECISION_TIMEOUT -): +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 diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py index 4855035435..10db93ea90 100644 --- a/studio/backend/tests/test_tool_approvals.py +++ b/studio/backend/tests/test_tool_approvals.py @@ -52,9 +52,7 @@ class _Waiter: 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 - ) + self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs) def start(self): self._thread.start() @@ -209,7 +207,9 @@ def test_concurrent_distinct_calls_route_their_own_decisions(): 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)} + 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(): diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py index f3a8d7194f..060d26cdef 100644 --- a/studio/backend/tests/test_tool_confirm_loop.py +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -36,7 +36,9 @@ class _FakeExecuteTool: def __init__(self): self.calls = [] - def __call__(self, name, arguments, *, cancel_event = None, timeout = None, session_id = None): + def __call__( + self, name, arguments, *, cancel_event = None, timeout = None, session_id = None + ): self.calls.append((name, arguments)) return f"RESULT[{name}]" From e5b2e9be64fa363590dcc1cfbc2017d66ea30a9e Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Sun, 31 May 2026 15:13:13 -0700 Subject: [PATCH 6/6] Move "Confirm tool calls" to the Tools section --- .../src/features/chat/chat-settings-sheet.tsx | 43 +++++++++++-------- 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 8eda82fb02..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,13 +1513,35 @@ 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( (s) => s.setMcpEnabledForChat, ); - const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); - const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); const [enabledServerCount, setEnabledServerCount] = useState( null, ); @@ -1559,22 +1582,6 @@ function McpServersSection() { disabled={enabledServerCount === 0 && !mcpEnabledForChat} />
-
-
- - Confirm tool calls - - - When on, every tool call pauses for your approval in the chat - before it runs. - -
- -
{enabledServerCount === null