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