diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 80b4862aa8..8cf37ed9ec 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -60,6 +60,13 @@ from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) @@ -2192,8 +2199,9 @@ class LlamaCppBackend: else None ), ( - f"{general['general.organization']}/" - f"{general['general.basename']}".replace(" ", "-") + f"{general['general.organization']}/{general['general.basename']}".replace( + " ", "-" + ) if general.get("general.organization") and general.get("general.basename") else None ), @@ -3748,7 +3756,7 @@ class LlamaCppBackend: ) logger.info( - f"llama-server ready on port {self._port} " f"for model '{model_identifier}'" + f"llama-server ready on port {self._port} for model '{model_identifier}'" ) # Probe outside _lock (interruptible by /unload); init inside. @@ -4373,7 +4381,7 @@ class LlamaCppBackend: proc.kill() logger.info( - f"Killed orphaned llama-server process " f"(pid={proc.info['pid']})" + f"Killed orphaned llama-server process (pid={proc.info['pid']})" ) except ( psutil.NoSuchProcess, @@ -4910,6 +4918,7 @@ class LlamaCppBackend: rag_scope: Optional[dict] = None, seed: Optional[int] = None, disable_parallel_tool_use: bool = False, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4928,7 +4937,7 @@ class LlamaCppBackend: # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits # the same tool card + citations a real call would. - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -5077,7 +5086,7 @@ class LlamaCppBackend: if response.status_code != 200: error_body = response.read().decode() raise RuntimeError( - f"llama-server returned {response.status_code}: " f"{error_body}" + f"llama-server returned {response.status_code}: {error_body}" ) raw_buf = "" @@ -5488,8 +5497,7 @@ class LlamaCppBackend: force = True, ) logger.info( - f"Safety net: parsed {len(tool_calls)} tool call(s) " - f"from streamed content" + f"Safety net: parsed {len(tool_calls)} tool call(s) from streamed content" ) else: # ── DRAINING path: assemble tool_calls ── @@ -5609,8 +5617,51 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = ( + begin_tool_decision(session_id, approval_id) if needs_confirm else None + ) + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + if _forced_tool_call_pending: + _forced_tool_call_pending = False + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4e4788ace9..4ccac2912e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -861,6 +861,7 @@ class InferenceOrchestrator: tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -922,6 +923,7 @@ class InferenceOrchestrator: tool_call_timeout = tool_call_timeout, session_id = session_id, rag_scope = rag_scope, + 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 7942edb6d7..3b6a393f3d 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -35,6 +35,13 @@ from core.inference.tool_loop_controller import ( status_for_tool, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) @@ -146,6 +153,7 @@ def run_safetensors_tool_loop( tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -174,7 +182,7 @@ def run_safetensors_tool_loop( # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. from core.inference.tools import build_rag_autoinject - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -509,8 +517,47 @@ def run_safetensors_tool_loop( else: assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b70202d6ff..039a5d75dd 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -690,6 +690,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.", @@ -926,6 +930,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 856d9bacf8..9ff30cd6db 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -655,6 +655,7 @@ from models.inference import ( ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ToolConfirmRequest, ChatMessage, ChunkChoice, ChoiceDelta, @@ -702,6 +703,7 @@ from core.inference.anthropic_compat import ( AnthropicPassthroughEmitter, ) from auth.authentication import get_current_subject +from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key from core.inference.providers import get_provider_info, get_base_url @@ -1495,8 +1497,7 @@ async def load_model( # Shouldn't happen on already-validated args; degrade to # no-extras rather than 400 if managed flags changed. logger.warning( - "Stored llama_extra_args failed revalidation; " - "loading without them: %s", + "Stored llama_extra_args failed revalidation; loading without them: %s", stripped, ) extra_llama_args = [] @@ -1938,6 +1939,20 @@ async def cancel_inference(request: Request, current_subject: str = Depends(get_ return {"cancelled": n} +@studio_router.post("/tool-confirm") +async def confirm_tool_call( + request: ToolConfirmRequest, current_subject: str = Depends(get_current_subject) +): + matched = resolve_tool_decision( + request.approval_id, + request.decision, + session_id = request.session_id, + ) + if not matched: + raise HTTPException(status_code = 404, detail = "No pending tool call confirmation") + return {"resolved": True} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, current_subject: str = Depends(get_current_subject) @@ -3196,6 +3211,22 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + if payload.confirm_tool_calls and ( + payload.enable_tools is True + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls is only supported for local streaming tools.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("external provider chat completions") return await _proxy_to_external_provider(payload, request) @@ -3567,6 +3598,16 @@ async def openai_chat_completions( use_tools = False if use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── @@ -3637,6 +3678,7 @@ async def openai_chat_completions( session_id = payload.session_id, rag_scope = payload.rag_scope, disable_parallel_tool_use = payload.parallel_tool_calls is False, + confirm_tool_calls = bool(payload.confirm_tool_calls), ) _tool_sentinel = object() @@ -3646,6 +3688,7 @@ async def openai_chat_completions( _tracker.__enter__() async def gguf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -3768,6 +3811,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _tracker.__exit__(None, None, None) return StreamingResponse( @@ -4091,6 +4139,16 @@ async def openai_chat_completions( _sf_use_tools = False if _sf_use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) _sf_nudge = _build_tool_action_nudge( tools = _sf_tools_to_use, model_name = model_name, @@ -4167,6 +4225,7 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, rag_scope = payload.rag_scope, + confirm_tool_calls = bool(payload.confirm_tool_calls), use_adapter = payload.use_adapter, stats_holder = _sf_stats_holder, ) @@ -4177,6 +4236,7 @@ async def openai_chat_completions( _sf_tracker.__enter__() async def sf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -4293,6 +4353,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _sf_tracker.__exit__(None, None, None) if payload.stream: @@ -5934,6 +5999,15 @@ async def anthropic_messages( ) if server_tools: + if bool(getattr(payload, "confirm_tool_calls", False)): + raise HTTPException( + status_code = 400, + detail = anthropic_error_body( + "confirm_tool_calls is not supported for Anthropic Messages server tools.", + status = 400, + err_type = "invalid_request_error", + ), + ) from core.inference.tools import ALL_TOOLS openai_tools = _select_anthropic_server_tools( diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py new file mode 100644 index 0000000000..f66226b61d --- /dev/null +++ b/studio/backend/state/tool_approvals.py @@ -0,0 +1,139 @@ +# 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 abort_tool_decision(slot, approval_id) -> None: + """Remove a slot that was announced but never entered ``wait_tool_decision``. + + Streaming wrappers may stop after ``tool_start`` is yielded and before + the loop resumes into ``wait_tool_decision``. In that case there is no + waiter to run the normal cleanup path, so the generator close path calls + this explicitly. + """ + 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). + + The first decision wins: once a slot's event is set, a later (duplicate or + out-of-order) confirmation for the same id is rejected without mutating the + recorded decision, so an Allow can never be flipped to Deny in the window + before the waiter reads ``slot["decision"]`` and pops the slot. + """ + 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 + if slot["event"].is_set(): + return False + slot["decision"] = decision + slot["event"].set() + return True diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index e1230ae113..f8d3f44d4b 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1543,6 +1543,19 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "tools" + def test_confirm_tool_calls_rejected_for_server_tools(self, monkeypatch): + backend = _mock_backend(monkeypatch) + payload = _basic_payload( + confirm_tool_calls = True, + tools = [{"type": "web_search_20250305", "name": "web_search"}], + ) + + with pytest.raises(HTTPException) as exc: + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert exc.value.status_code == 400 + assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"] + assert backend.calls == [] + def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): backend = _mock_backend(monkeypatch) payload = _basic_payload( diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fa583ef53d..3c121c281d 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -21,6 +21,8 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) from core.inference.llama_cpp import LlamaCppBackend +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision def _sse(delta: dict) -> str: @@ -70,6 +72,27 @@ def _tool_names(payload: dict) -> list[str]: ] +def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]: + return [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": json.dumps(arguments), + }, + } + ] + } + ), + _done(), + ] + + def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): """llama-server may emit content first and then native delta.tool_calls. @@ -1149,3 +1172,151 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now.", "Final note after tool."] assert len(payloads) == 3 + + +def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): + streams = [ + _structured_tool_call("python", {"code": "print(1)"}, "call_py"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1") + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow") + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + assert len(starts) == 1 + assert starts[0]["approval_id"] + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events) + + +def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch): + approval_id = "approval-close" + streams = [_structured_tool_call("python", {"code": "print(1)"}, "call_py")] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("tool should not run")), + ) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: approval_id) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + gen = backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + try: + assert next(gen)["type"] == "status" + start = next(gen) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + gen.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + + +def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch): + streams = [[_sse({"content": "Done."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "use docs"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + rag_scope = {"thread_id": "t1"}, + ) + ) + + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + + +def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch): + same_call = _structured_tool_call("python", {"code": "print(1)"}, "call_py") + streams = [ + same_call, + _structured_tool_call("python", {"code": "print(1)"}, "call_py_retry"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + decisions = iter(["deny", "allow"]) + approvals = iter(["approval-1", "approval-2"]) + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals)) + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr( + "core.inference.llama_cpp.wait_tool_decision", + lambda *_a, **_k: next(decisions), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 2, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + ends = [event for event in events if event.get("type") == "tool_end"] + assert len(starts) == 2 + assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"] + assert calls == [("python", {"code": "print(1)"})] diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1d994b46c0..79e8c9ae49 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -401,6 +401,28 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_rejected_for_provider_tools(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "external_model": "gpt-4.1", + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "only supported for local streaming tools" in body["error"]["message"] + def test_logprobs_rejected_until_supported(self, monkeypatch): class _UnusedBackend: is_loaded = False @@ -480,6 +502,7 @@ class TestChatCompletionRequestToolFields: def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False + supports_tools = False class _InferenceBackend: active_model_name = "test-model" @@ -495,6 +518,45 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_requires_streaming_for_safetensors_tools(self, monkeypatch): + import routes.inference as inference_route + + class _NoGGUFBackend: + is_loaded = False + supports_tools = False + + class _InferenceBackend: + active_model_name = "test-model" + models = {"test-model": {"chat_template_info": {"template": "chatml"}}} + + def generate_chat_completion_with_tools(self, **kwargs): + raise AssertionError("tool loop should be rejected before starting") + + def generate_chat_completion(self, **kwargs): + raise AssertionError("plain path should not be used") + + monkeypatch.setattr( + inference_route, + "_detect_safetensors_features", + lambda backend, chat_template: {"supports_tools": True}, + ) + client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + "stream": False, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "requires stream=true" in body["error"]["message"] + def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( messages = [ @@ -1206,6 +1268,45 @@ class TestGgufVisionToolRouting: assert captured["kwargs"]["disable_parallel_tool_use"] is True + def test_confirm_tool_calls_requires_streaming_for_gguf_tools(self, monkeypatch): + import routes.inference as inf_mod + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + raise AssertionError("tool loop should be rejected before starting") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + model_identifier = "test-gguf", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + confirm_tool_calls = True, + stream = False, + messages = [{"role": "user", "content": "search once"}], + ) + + with pytest.raises(HTTPException) as exc: + self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert exc.value.status_code == 400 + assert "requires stream=true" in exc.value.detail["error"]["message"] + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 8aa6e5df4e..12731783a0 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -28,6 +28,8 @@ from core.inference.tool_call_parser import ( parse_tool_calls_from_text, strip_tool_markup, ) +from state import tool_approvals +from state.tool_approvals import resolve_tool_decision from utils.datasets import is_gpt_oss_model_name @@ -84,8 +86,7 @@ class TestParser: # A code parameter with a literal must not truncate: the # parser uses end-of-body as the only boundary for single-param calls. text = ( - "html = ''\n" - "print('hi')" + "html = ''\nprint('hi')" ) result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -1033,6 +1034,50 @@ class TestGuardrails: _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] + def test_confirm_tool_calls_close_after_prompt_cleans_slot(self, monkeypatch): + approval_id = "approval-close-sf" + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: approval_id) + + loop, exec_fn = _make_loop( + turns = [['{"name":"python","arguments":{"code":"print(1)"}}']], + exec_results = ["OK"], + confirm_tool_calls = True, + session_id = "sess", + max_tool_iterations = 1, + ) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + try: + assert next(loop)["type"] == "status" + start = next(loop) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + loop.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + assert exec_fn.calls == [] + + def test_confirm_tool_calls_skips_rag_autoinject(self, monkeypatch): + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + loop, exec_fn = _make_loop( + turns = [["plain answer"]], + confirm_tool_calls = True, + rag_scope = {"thread_id": "t1"}, + ) + events = _collect_events(loop) + assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events) + assert exec_fn.calls == [] + def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): turns = iter( [ diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py new file mode 100644 index 0000000000..af792e652c --- /dev/null +++ b/studio/backend/tests/test_tool_approvals.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Concurrency tests for the per-call tool-call confirmation gate. + +``state.tool_approvals`` coordinates two threads: the agentic loop thread +blocked in ``wait_tool_decision`` and the request thread that delivers the +user's choice through ``resolve_tool_decision``. Each gated call carries a +unique ``approval_id`` so a stale or concurrent confirmation can never +resolve the wrong call. These tests exercise that handshake directly -- +no model, no server -- so the race windows are fast and deterministic. +""" + +import threading +import time + +import pytest + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + request_tool_decision, + resolve_tool_decision, + wait_tool_decision, +) + + +@pytest.fixture(autouse = True) +def _clear_pending(): + """Each test starts and ends with an empty ``_pending`` map.""" + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _Waiter: + """Run ``request_tool_decision`` in a thread and capture its result.""" + + def __init__( + self, + session_id, + approval_id, + cancel_event = None, + timeout = None, + ): + self.session_id = session_id + self.approval_id = approval_id + self.cancel_event = cancel_event + self.timeout = timeout + self.result = None + self._thread = threading.Thread(target = self._run, daemon = True) + + def _run(self): + kwargs = {"cancel_event": self.cancel_event} + if self.timeout is not None: + kwargs["timeout"] = self.timeout + self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs) + + def start(self): + self._thread.start() + _wait_until(lambda: _has_pending(self.approval_id)) + return self + + def join(self, timeout = 5.0): + self._thread.join(timeout = timeout) + assert not self._thread.is_alive(), "waiter thread did not finish" + return self.result + + +def _has_pending(approval_id) -> bool: + with tool_approvals._lock: + return approval_id in tool_approvals._pending + + +def _wait_until( + pred, + timeout = 2.0, + interval = 0.005, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(interval) + return False + + +# ── Basic allow / deny ─────────────────────────────────────────────── + + +def test_allow_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + assert w.join() == "allow" + + +def test_deny_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "deny", session_id = "sess") is True + assert w.join() == "deny" + + +def test_slot_cleaned_up_after_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + resolve_tool_decision(aid, "allow") + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_abort_tool_decision_removes_unwaited_slot(): + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + abort_tool_decision(slot, aid) + assert not _has_pending(aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is False + + +def test_approval_ids_are_unique(): + ids = {new_approval_id() for _ in range(1000)} + assert len(ids) == 1000 + + +# ── Pre-registration race (begin before wait) ──────────────────────── + + +def test_resolve_before_wait_is_not_lost(): + """A decision delivered after ``begin`` but before ``wait`` survives. + + The loop registers the slot before it yields ``tool_start``, so even a + confirmation that races ahead of the blocking ``wait`` is recorded on + the slot and returned -- never dropped. + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # wait() is only entered now, after the decision already landed. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Resolver edge cases ────────────────────────────────────────────── + + +def test_resolve_unknown_approval_returns_false(): + assert resolve_tool_decision(new_approval_id(), "allow") is False + + +def test_resolve_empty_approval_returns_false(): + assert resolve_tool_decision("", "allow") is False + assert resolve_tool_decision(None, "allow") is False + + +def test_resolve_wrong_session_scope_returns_false(): + aid = new_approval_id() + w = _Waiter("sess-a", aid).start() + # Correct approval_id but the wrong session must not resolve it. + assert resolve_tool_decision(aid, "allow", session_id = "sess-b") is False + assert _has_pending(aid) + # The right session still works. + assert resolve_tool_decision(aid, "allow", session_id = "sess-a") is True + assert w.join() == "allow" + + +def test_duplicate_resolve_after_completion_returns_false(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow") is True + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + assert resolve_tool_decision(aid, "deny") is False + + +def test_first_decision_is_immutable(): + """A second confirmation cannot flip an already-recorded decision. + + The waiter reads ``slot["decision"]`` outside the lock and then cleans up, + so a duplicate or out-of-order POST that lands in that window must be + rejected and must not overwrite the first decision -- an Allow can never + become a Deny. Distinct from the after-completion case above: here the slot + is still pending (no waiter has consumed it yet). + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # Second decision, same id, before any waiter consumes/cleans the slot. + assert resolve_tool_decision(aid, "deny", session_id = "sess") is False + assert slot["decision"] == "allow" + # The waiter still observes the first (immutable) decision. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Cancellation and timeout ───────────────────────────────────────── + + +def test_cancel_event_breaks_wait_as_deny(): + cancel = threading.Event() + aid = new_approval_id() + w = _Waiter("sess", aid, cancel_event = cancel).start() + cancel.set() + assert w.join(timeout = 3.0) == "deny" + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_timeout_returns_deny(): + aid = new_approval_id() + start = time.monotonic() + result = request_tool_decision("sess", aid, timeout = 0.1) + assert result == "deny" + assert time.monotonic() - start < 2.0 + assert not _has_pending(aid) + + +# ── Independence across concurrent calls ───────────────────────────── + + +def test_two_pending_calls_same_session_are_independent(): + """Keying on approval_id, not session, keeps concurrent calls distinct. + + Resolving the first call's id must not unblock or alter the second + call pending in the same session. + """ + a1, a2 = new_approval_id(), new_approval_id() + w1 = _Waiter("sess", a1).start() + w2 = _Waiter("sess", a2).start() + + assert resolve_tool_decision(a1, "deny", session_id = "sess") is True + assert w1.join() == "deny" + # w2 is still waiting on its own id. + assert _has_pending(a2) + assert resolve_tool_decision(a2, "allow", session_id = "sess") is True + assert w2.join() == "allow" + + +def test_concurrent_distinct_calls_route_their_own_decisions(): + n = 25 + waiters = {} + for i in range(n): + aid = new_approval_id() + waiters[aid] = _Waiter(f"s{i}", aid).start() + expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)} + for aid, decision in expected.items(): + assert resolve_tool_decision(aid, decision) is True + for aid, w in waiters.items(): + assert w.join() == expected[aid] + + +# ── Constants ──────────────────────────────────────────────────────── + + +def test_rejected_message_is_user_facing_text(): + assert isinstance(TOOL_REJECTED_MESSAGE, str) + assert TOOL_REJECTED_MESSAGE.strip() 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..ce7852c95f --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -0,0 +1,170 @@ +# 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.safetensors_agentic import run_safetensors_tool_loop +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, + rag_scope = 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(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + [], + tools = [{"type": "function", "function": {"name": "web_search"}}], + ) + assert _tool_starts(events) == [] + assert _tool_ends(events) == [] + assert calls == [] + + +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) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert len(_tool_ends(events)) == 1 + + +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..b8e0472e12 --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -0,0 +1,219 @@ +# 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 84e1654079..100143ed33 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -19,6 +19,7 @@ import { thinkEffortAriaLabel, thinkToggleAriaLabel, } from "@/components/assistant-ui/think-aria-label"; +import { withToolConfirmation } 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"; @@ -2520,6 +2521,19 @@ const CancelledIndicator: FC = () => { ); }; +const WebSearchToolUIConfirmable = withToolConfirmation(WebSearchToolUI); +const KnowledgeBaseToolUIConfirmable = + withToolConfirmation(KnowledgeBaseToolUI); +const PythonToolUIConfirmable = withToolConfirmation(PythonToolUI); +const TerminalToolUIConfirmable = withToolConfirmation(TerminalToolUI); +const CodeExecutionToolUIConfirmable = + withToolConfirmation(CodeExecutionToolUI); +const ImageGenerationToolUIConfirmable = withToolConfirmation( + ImageGenerationToolUI, +); +const RenderHtmlToolUIConfirmable = withToolConfirmation(RenderHtmlToolUI); +const ToolFallbackConfirmable = withToolConfirmation(ToolFallback); + const AssistantMessage: FC = () => { return ( { ToolGroup: ToolGroup, tools: { by_name: { - web_search: WebSearchToolUI, - search_knowledge_base: KnowledgeBaseToolUI, - python: PythonToolUI, - terminal: TerminalToolUI, - code_execution: CodeExecutionToolUI, - image_generation: ImageGenerationToolUI, - render_html: RenderHtmlToolUI, + web_search: WebSearchToolUIConfirmable, + search_knowledge_base: KnowledgeBaseToolUIConfirmable, + python: PythonToolUIConfirmable, + terminal: TerminalToolUIConfirmable, + code_execution: CodeExecutionToolUIConfirmable, + image_generation: ImageGenerationToolUIConfirmable, + render_html: RenderHtmlToolUIConfirmable, }, - 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..89a15c61ab --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx @@ -0,0 +1,157 @@ +// 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 { + ToolCallMessagePartComponent, + 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 && + Object.prototype.hasOwnProperty.call(s.toolConfirmations, toolCallId) + ? s.toolConfirmations[toolCallId] + : undefined, + ); + const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways); + const clearToolConfirmation = useChatRuntimeStore( + (s) => s.clearToolConfirmation, + ); + const autoAllowKey = confirmation?.autoAllowKey ?? ""; + const autoAllowed = useChatRuntimeStore( + (s) => + s.alwaysAllowToolsBySession.get(autoAllowKey)?.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} +
+ ); +} + +export function withToolConfirmation( + Component: ToolCallMessagePartComponent, +): ToolCallMessagePartComponent { + const WithToolConfirmation: ToolCallMessagePartComponent = (props) => ( + <> + + + + ); + return WithToolConfirmation; +} diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index eeb4c2059f..8930b6386f 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -325,6 +325,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/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx index 569e673da3..3ce65627e2 100644 --- a/studio/frontend/src/components/assistant-ui/tool-group.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx @@ -9,6 +9,7 @@ import { type PropsWithChildren, } from "react"; import { useAuiState } from "@assistant-ui/react"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { ChevronDownIcon } from "lucide-react"; import { Wrench01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -216,6 +217,31 @@ const ToolGroupImpl: FC< (part) => part.type === "tool-call" && part.toolName === "render_html", ), ); + // A blocking allow/deny prompt must never be hidden inside a collapsed + // group, so force the group open while any of its calls awaits confirmation. + const toolConfirmations = useChatRuntimeStore((s) => s.toolConfirmations); + const hasPendingConfirmation = useAuiState(({ message }) => + message.parts + .slice(startIndex, endIndex + 1) + .some( + (part) => + part.type === "tool-call" && + Object.prototype.hasOwnProperty.call( + toolConfirmations, + part.toolCallId, + ), + ), + ); + const messageRunning = useAuiState( + ({ message }) => message.status?.type === "running", + ); + // Keep the group open once a confirmation forced it open, so answering an + // allow/deny doesn't snap it shut between sequential tool calls. It reverts + // to the default collapsed state once the turn finishes. + const forcedOpenRef = useRef(false); + if (hasPendingConfirmation) forcedOpenRef.current = true; + const forceOpen = + hasPendingConfirmation || (forcedOpenRef.current && messageRunning); // Render single tool calls and artifacts directly so cards never hide in a // collapsed group. @@ -224,7 +250,7 @@ const ToolGroupImpl: FC< } return ( - + {children} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 983d4f5aa2..ce9f1c6544 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1440,6 +1440,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const resolvedThreadId = (unstable_threadId ?? runtime.activeThreadId) || undefined; const sandboxSessionId = await resolveSandboxSessionId(resolvedThreadId); + const toolConfirmationScopeId = resolvedThreadId + ? `${sandboxSessionId || "_default"}:${resolvedThreadId}` + : sandboxSessionId || "_default"; + const toolConfirmationIdsByBackendId = new Map(); const resolvedThreadKey = resolvedThreadId ?? null; const pendingImageEditReferenceForRun = runtime.pendingImageEditReference; const selectedImageEditReference = @@ -1513,6 +1517,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { imageToolsEnabled, artifactsEnabled, mcpEnabledForChat, + confirmToolCalls, webFetchToolsEnabled, ragEnabled, ragSource, @@ -2435,6 +2440,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : []), ], mcp_enabled: mcpEnabledForChat, + confirm_tool_calls: confirmToolCalls, // Scope: thread_id = this thread's docs, kb_id = a KB. ...(ragEnabled ? { @@ -2560,9 +2566,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toolEvent.provenance, ); if (toolEvent.type === "tool_start") { + const backendToolCallId = + (toolEvent.tool_call_id as string) || ""; + const approvalId = (toolEvent.approval_id as string) || ""; + const awaitingConfirmation = + toolEvent.awaiting_confirmation === true; const id = - (toolEvent.tool_call_id as string) || - `${toolEvent.tool_name}_${Date.now()}`; + awaitingConfirmation && approvalId + ? `${toolConfirmationScopeId}:${approvalId}` + : backendToolCallId || + approvalId || + `${toolEvent.tool_name}_${Date.now()}`; + if (awaitingConfirmation && backendToolCallId) { + toolConfirmationIdsByBackendId.set(backendToolCallId, id); + } const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"]; const idx = toolCallParts.findIndex( @@ -2593,11 +2610,30 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(toolProvenance ? { provenance: toolProvenance } : {}), } as PositionedToolCallPart); } + if (awaitingConfirmation) { + useChatRuntimeStore + .getState() + .setToolConfirmation( + id, + approvalId, + sandboxSessionId ?? "", + toolConfirmationScopeId, + ); + } } else if (toolEvent.type === "tool_end") { + const backendToolCallId = + (toolEvent.tool_call_id as string) || ""; const id = - (toolEvent.tool_call_id as string) || + (backendToolCallId + ? toolConfirmationIdsByBackendId.get(backendToolCallId) + : undefined) || + backendToolCallId || toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; + if (backendToolCallId) { + toolConfirmationIdsByBackendId.delete(backendToolCallId); + } + useChatRuntimeStore.getState().clearToolConfirmation(id); const idx = toolCallParts.findIndex( (p) => p.toolCallId === id, ); @@ -3149,6 +3185,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw err; } finally { abortSignal.removeEventListener("abort", onAbortCancel); + const confirmStore = useChatRuntimeStore.getState(); + for (const part of toolCallParts) { + confirmStore.clearToolConfirmation(part.toolCallId); + } runtime.setGeneratingStatus(null); runtime.setToolStatus(null); clearTimeout(warmupTimer); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 9d0939c37d..110d573571 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -110,6 +110,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 5751a1a426..59ac2ae6b4 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1496,6 +1496,7 @@ export function ChatSettingsPanel({
+
@@ -1682,6 +1683,30 @@ function AutoHealToolCallsToggle() { ); } +function ConfirmToolCallsToggle() { + const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); + const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); + + return ( +
+
+ + Confirm tool calls + + + When on, local Studio tool calls pause for your approval before they + run. Provider-hosted tools are not gated here. + +
+ +
+ ); +} + function ChatTemplateFields() { const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate); const override = useChatRuntimeStore((s) => s.chatTemplateOverride); 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 3ef63d2945..2f0e3027a8 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -34,6 +34,7 @@ export const CHAT_COLLAPSE_HTML_ARTIFACTS_KEY = export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY = "unsloth_chat_allow_artifact_network_access"; 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"; export const CHAT_RAG_SOURCE_KEY = "unsloth_chat_rag_source"; @@ -402,6 +403,29 @@ type ChatRuntimeStore = { // autoInject = forced first-pass retrieval before answering. ragAutoInject: RagAutoInject; ragAutoInjectMinScore: number; + /** + * When on, local Studio tool calls pause for an explicit allow/deny in the + * chat before they run. + */ + confirmToolCalls: boolean; + /** + * Per-chat set of tool names the user chose to auto-approve via "Always + * allow". Keyed by UI confirmation scope, not necessarily the backend + * sandbox session id. Not persisted across reloads. + */ + alwaysAllowToolsBySession: Map>; + /** + * Tool calls currently paused awaiting the user's allow/deny decision, + * keyed by the scoped 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. The + * ``autoAllowKey`` scopes the UI-only "Always allow" bucket per chat. + * Only backend-gated local tool calls are added here. + */ + toolConfirmations: Record< + string, + { approvalId: string; sessionId: string; autoAllowKey: string } + >; /** * Fetch pill state, independent of `toolsEnabled` (Search). Only * consulted when `providerSupportsBuiltinWebFetch` is true. @@ -483,6 +507,15 @@ type ChatRuntimeStore = { setCollapseHtmlArtifacts: (enabled: boolean) => void; setAllowArtifactNetworkAccess: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; + setConfirmToolCalls: (enabled: boolean) => void; + allowToolAlways: (sessionId: string, toolName: string) => void; + setToolConfirmation: ( + toolCallId: string, + approvalId: string, + sessionId: string, + autoAllowKey: string, + ) => void; + clearToolConfirmation: (toolCallId: string) => void; setWebFetchToolsEnabled: (enabled: boolean) => void; setRagEnabled: (enabled: boolean) => void; setRagSource: (source: RagSource) => void; @@ -749,6 +782,9 @@ export const useChatRuntimeStore = create((set, get) => ({ 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), // RAG is opt-in per session: always starts off, never restored from storage. ragEnabled: false, @@ -1074,6 +1110,40 @@ 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, autoAllowKey) => + set((state) => ({ + toolConfirmations: { + ...state.toolConfirmations, + [toolCallId]: { approvalId, sessionId, autoAllowKey }, + }, + })), + clearToolConfirmation: (toolCallId) => + set((state) => { + if ( + !Object.prototype.hasOwnProperty.call( + state.toolConfirmations, + toolCallId, + ) + ) { + return state; + } + const next = { ...state.toolConfirmations }; + delete next[toolCallId]; + return { toolConfirmations: next }; + }), setWebFetchToolsEnabled: (webFetchToolsEnabled) => set(() => { saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index a0e5d5355e..96caa38bd4 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -282,6 +282,8 @@ export interface OpenAIChatCompletionsRequest { enabled_tools?: string[]; /** Local models + enable_tools only. */ mcp_enabled?: boolean; + /** Local models + enable_tools only. */ + confirm_tool_calls?: boolean; /** Exactly one of `kb_id` (a KB) or `thread_id` (thread docs). */ rag_scope?: { kb_id?: string;