ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls (#5869)

Browse source 
* Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix race in tool-call confirmation gate

* 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).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move "Confirm tool calls" to the Tools section

* Studio: Keep tool group open while a tool call awaits confirmation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tool confirmation session scope for PR #5869

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix confirmation follow-ups for PR #5869

* Apply pre-commit formatting for PR #5869

* Fix confirmation cleanup for PR #5869

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden confirmation lookups for PR #5869

* Studio: make the tool-call confirmation decision immutable

resolve_tool_decision accepted a second confirmation for the same approval_id
and overwrote slot["decision"] in the window before the waiter reads it and
pops the slot, so a duplicate or out-of-order POST could flip an Allow to Deny
(and returned a misleading resolved:true). Reject once the slot's event is
already set so the first decision wins. Adds a regression test.

* Fix/adjust tool confirmations for PR #5869

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
This commit is contained in:
oobabooga 2026-06-12 05:55:26 -03:00 • committed by GitHub
commit 7f2986a413
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 1694 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

71
studio/backend/core/inference/llama_cpp.py
View file

@ -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.

2
studio/backend/core/inference/orchestrator.py
View file

@ -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(

53
studio/backend/core/inference/safetensors_agentic.py
View file

@ -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.

10
studio/backend/models/inference.py
View file

@ -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 ─────────────────────

78
studio/backend/routes/inference.py
View file

@ -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(

139
studio/backend/state/tool_approvals.py Normal file
View file

@ -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

13
studio/backend/tests/test_anthropic_messages.py
View file

@ -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(

171
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -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)"})]

101
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -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

49
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -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 </parameter> must not truncate: the
# parser uses end-of-body as the only boundary for single-param calls.
text = (
"<function=python><parameter=code>html = '<a></a>'\n"
"print('hi')</parameter></function>"
"<function=python><parameter=code>html = '<a></a>'\nprint('hi')</parameter></function>"
)
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 = [['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>']],
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(
[

261
studio/backend/tests/test_tool_approvals.py Normal file
View file

@ -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()

170
studio/backend/tests/test_tool_confirm_loop.py Normal file
View file

@ -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'<tool_call>{{"name": "{name}", "arguments": {args_json}}}</tool_call>'
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]"

219
studio/backend/tests/test_tool_confirm_stream.py Normal file
View file

@ -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"]

30
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -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 (
<MessagePrimitive.Root
@ -2538,15 +2552,15 @@ const AssistantMessage: FC = () => {
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,
},
}}
/>

157
studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx Normal file
View file

@ -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 (
<div className="flex flex-wrap items-center gap-2 pt-1">
<Button
size="xs"
disabled={pending !== null}
onClick={() => void resolve("allow")}
>
Allow
</Button>
<Button
size="xs"
variant="outline"
disabled={pending !== null}
onClick={() => {
if (autoAllowKey) allowToolAlways(autoAllowKey, toolName);
void resolve("allow");
}}
>
Always allow
</Button>
<Button
size="xs"
variant="destructive"
disabled={pending !== null}
onClick={() => void resolve("deny")}
>
Deny
</Button>
{failed ? (
<span className="text-xs text-destructive">
Could not send your decision. Try again.
</span>
) : null}
</div>
);
}
export function withToolConfirmation(
Component: ToolCallMessagePartComponent,
): ToolCallMessagePartComponent {
const WithToolConfirmation: ToolCallMessagePartComponent = (props) => (
<>
<Component {...props} />
<ToolConfirmationControls
toolCallId={props.toolCallId}
toolName={props.toolName}
result={props.result}
status={props.status}
/>
</>
);
return WithToolConfirmation;
}

3
studio/frontend/src/components/assistant-ui/tool-fallback.tsx
View file

@ -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";

28
studio/frontend/src/components/assistant-ui/tool-group.tsx
View file

@ -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 (
<ToolGroupRoot>
<ToolGroupRoot open={forceOpen ? true : undefined}>
<ToolGroupTrigger count={toolCount} />
<ToolGroupContent>{children}</ToolGroupContent>
</ToolGroupRoot>

46
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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<string, string>();
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);

25
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -110,6 +110,31 @@ export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
await parseJsonOrThrow<unknown>(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<boolean> {
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;

25
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -1496,6 +1496,7 @@ export function ChatSettingsPanel({
<CollapsibleSection label="Tools">
<div className="flex flex-col gap-5 pt-1">
<AutoHealToolCallsToggle />
<ConfirmToolCallsToggle />
<MaxToolCallsSlider />
<ToolCallTimeoutSlider />
</div>
@ -1682,6 +1683,30 @@ function AutoHealToolCallsToggle() {
);
}
function ConfirmToolCallsToggle() {
const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls);
const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls);
return (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Confirm tool calls
</span>
<InfoHint>
When on, local Studio tool calls pause for your approval before they
run. Provider-hosted tools are not gated here.
</InfoHint>
</div>
<Switch
className="panel-switch"
checked={confirmToolCalls}
onCheckedChange={setConfirmToolCalls}
/>
</div>
);
}
function ChatTemplateFields() {
const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate);
const override = useChatRuntimeStore((s) => s.chatTemplateOverride);

70
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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<string, Set<string>>;
/**
* 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<ChatRuntimeStore>((set, get) => ({
false,
),
mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false),
confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false),
alwaysAllowToolsBySession: new Map<string, Set<string>>(),
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<ChatRuntimeStore>((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);

2
studio/frontend/src/features/chat/types/api.ts
View file

@ -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;