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

Studio: show tool-call progress for large GGUF tool arguments (#6484)

Browse source 
* Studio: show tool-call progress for large GGUF tool arguments

The GGUF agentic tool loop only surfaced an early provisional tool card
for render_html, so any other tool (python, terminal, ...) was invisible
in the UI while its arguments streamed. For a large argument such as a
full HTML or code file this left the chat sitting on "Generating..." with
zero progress for tens of seconds while the model was clearly working.

Generalize the provisional tool_start to any enabled tool once its
streamed arguments grow past a threshold (render_html still surfaces
immediately, small-argument tools are unchanged). The provisional and the
real tool_start share the tool_call_id so the frontend reconciles them
into one card. Close the provisional on no-op, denial, parallel-drop,
post-loop, and on stream errors so a card can never spin forever, surface
each parallel call, and skip the early card while a human confirmation
gate is active. Apply the same confirmation-gate guard to the safetensors
agentic loop.

Additional hardening:
- Only emit a provisional card once a real, non-empty tool_call_id is
  known. llama.cpp can stream a tool call with an empty id, and a card
  keyed by "" cannot reconcile with the real tool_start (the frontend
  mints its own id per event), so it would dangle.
- On a connection drop or other mid-iteration failure, close the dangling
  provisional card with an error result instead of an empty success so the
  UI renders it as failed rather than completed.
- Mirror the provisional cleanup in the safetensors loop: close a
  provisional render_html card if the model generator raises mid-stream or
  the controller turns the call into an internal no-op.

Adds regression tests for the empty-id guard, the error-result on a
dropped connection, and the safetensors mid-stream exception cleanup.

* [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: wasimysaid <wasimysdev@gmail.com>
This commit is contained in:
Daniel Han 2026-06-22 05:50:10 -07:00 • committed by GitHub
commit e2e8e5ab46
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 602 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -226,6 +226,10 @@ _MAX_REPROMPTS = 1
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
# Only large streamed tool payloads get an early provisional card; render_html
# is exempt because it needs immediate artifact feedback.
_PROVISIONAL_ARGS_MIN_CHARS = 256
_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min
_REPROMPT_MAX_CHARS = 2000
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
@ -7900,7 +7904,9 @@ class LlamaCppBackend:
_iter_finish_reason = None
_stream_done = False
_last_emitted = ""
provisional_render_html_tool_call_ids = set()
# Provisional tool_start cards already shown, keyed by tool_call_id.
provisional_started_tool_calls: dict[str, str] = {}
resolved_provisional_tool_call_ids: set[str] = set()
_suppress_visible_output = _forced_tool_call_pending
with self._open_stream(url, payload, cancel_event) as (
@ -7966,11 +7972,8 @@ class LlamaCppBackend:
# ── Structured tool_calls ──
tc_deltas = delta.get("tool_calls")
if tc_deltas:
# llama-server can emit visible assistant
# preface content before native structured
# tool_calls. Preserve content_accum as
# the assistant pre-tool text and still
# drain/execute the structured call.
# Preserve any visible preface before draining
# the structured tool call.
has_structured_tc = True
detect_state = _S_DRAINING
for tc_d in tc_deltas:
@ -8001,27 +8004,54 @@ class LlamaCppBackend:
fallback_id = f"call_{idx}"
current_id = tool_calls_acc[idx].get("id", fallback_id)
already_started = (
current_id in provisional_render_html_tool_call_ids
current_id in provisional_started_tool_calls
)
has_real_id = current_id != fallback_id
if (
# Empty/synthetic ids cannot reconcile with real starts.
has_real_id = bool(current_id) and current_id != fallback_id
# Show one early card per eligible streamed tool call.
_is_completed_one_shot = (
current_name == "render_html"
and not _tool_succeeded("render_html")
and _tool_succeeded("render_html")
)
# render_html is one-shot.
_one_shot_already_provisional = (
current_name == "render_html"
and "render_html"
in provisional_started_tool_calls.values()
)
# Later parallel cards only reconcile when parallel use is enabled.
_confirm_gated = (
confirm_tool_calls and not bypass_permissions
)
# Keep small-argument tools on the normal path.
_args_len = len(
tool_calls_acc[idx]["function"].get("arguments", "")
)
_payload_is_large = (
current_name == "render_html"
or _args_len >= _PROVISIONAL_ARGS_MIN_CHARS
)
if (
current_name
and (idx == 0 or not disable_parallel_tool_use)
and has_real_id
and not already_started
and not _is_completed_one_shot
and not _one_shot_already_provisional
and not _confirm_gated
and _payload_is_large
and any(
(
(tool.get("function") or {}).get("name")
== "render_html"
)
(tool.get("function") or {}).get("name")
== current_name
for tool in active_tools
)
and not already_started
and not provisional_render_html_tool_call_ids
and has_real_id
):
provisional_render_html_tool_call_ids.add(current_id)
provisional_started_tool_calls[current_id] = (
current_name
)
yield {
"type": "tool_start",
"tool_name": "render_html",
"tool_name": current_name,
"tool_call_id": current_id,
"arguments": {},
"provenance": tool_event_provenance(
@ -8343,20 +8373,30 @@ class LlamaCppBackend:
for tc in tool_calls or []:
func = tc.get("function", {})
tool_name = func.get("name", "")
provisional_render_html_match = (
tool_name == "render_html"
and tc.get("id") in provisional_render_html_tool_call_ids
)
provisional_match = tc.get("id") in provisional_started_tool_calls
decision = tool_controller.prepare_call(
tc,
forced = _forced_tool_call_pending,
provisional = provisional_render_html_match,
provisional = provisional_match,
)
if not decision.should_execute:
if content_text and not assistant_appended:
conversation.append(assistant_msg)
assistant_appended = True
if provisional_match:
# A provisional tool card is already on screen for this
# id; close it so it never dangles when the controller
# turns the call into an internal no-op (duplicate /
# disabled / render_html_repeat).
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
"tool_call_id": decision.tool_call_id,
"result": "",
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
if _forced_tool_call_pending:
@ -8401,6 +8441,7 @@ class LlamaCppBackend:
== "deny"
):
decision_slot = None
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
@ -8444,12 +8485,25 @@ class LlamaCppBackend:
if decision.tool_name == "search_knowledge_base":
_kb_search_count += 1
completion = tool_controller.record_result(decision, result)
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
yield completion.tool_end_event()
conversation.append(completion.tool_message())
if _forced_tool_call_pending:
_forced_tool_call_pending = False
# Close provisional cards not resolved by execution/no-op handling.
for _pid, _pname in provisional_started_tool_calls.items():
if _pid not in resolved_provisional_tool_call_ids:
resolved_provisional_tool_call_ids.add(_pid)
yield {
"type": "tool_end",
"tool_name": _pname,
"tool_call_id": _pid,
"result": "",
"provenance": tool_event_provenance(provisional = True),
}
# Clear tool status badge before next generation/final pass.
yield {"type": "status", "text": ""}
if tool_controller.force_final_answer or not tool_controller.active_tools():
@ -8458,10 +8512,32 @@ class LlamaCppBackend:
continue
except httpx.ConnectError:
# Mark unresolved provisional cards as failed before raising.
for _pid, _pname in provisional_started_tool_calls.items():
if _pid not in resolved_provisional_tool_call_ids:
resolved_provisional_tool_call_ids.add(_pid)
yield {
"type": "tool_end",
"tool_name": _pname,
"tool_call_id": _pid,
"result": "Error: lost connection to llama-server before the tool call completed.",
"provenance": tool_event_provenance(provisional = True),
}
raise RuntimeError("Lost connection to llama-server")
except Exception as e:
if cancel_event is not None and cancel_event.is_set():
return
# Same cleanup for other mid-iteration failures.
for _pid, _pname in provisional_started_tool_calls.items():
if _pid not in resolved_provisional_tool_call_ids:
resolved_provisional_tool_call_ids.add(_pid)
yield {
"type": "tool_end",
"tool_name": _pname,
"tool_call_id": _pid,
"result": "Error: the tool call was interrupted before it completed.",
"provenance": tool_event_provenance(provisional = True),
}
raise
# ── Tool iteration cap reached -- synthesize final answer ──

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

@ -236,12 +236,39 @@ def run_safetensors_tool_loop(
cumulative_display = ""
last_emitted = ""
provisional_render_html_started = False
provisional_resolved = False
provisional_render_html_id = f"call_{next_call_id}"
# When a human confirmation gate is active the real tool_start is keyed
# by an approval id and carries awaiting_confirmation, so an early
# provisional card (keyed by tool_call_id, no approval) would show the
# tool as "running" before the user has approved it. Suppress the early
# card in that case and let the gated tool_start be the first signal.
_provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions
gen = _call_single_turn(single_turn, conversation, active_tools)
prev_cumulative = ""
for cumulative in gen:
_gen_iter = iter(gen)
while True:
try:
cumulative = next(_gen_iter)
except StopIteration:
break
except Exception:
# The model pipeline raised mid-stream. If a provisional
# render_html card was already surfaced, close it as errored so
# the UI never leaves a tool card spinning after the turn fails.
if provisional_render_html_started and not provisional_resolved:
provisional_resolved = True
yield {
"type": "tool_end",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"result": "Error: generation was interrupted before the tool call completed.",
"provenance": _tool_event_provenance(provisional = True),
}
raise
if cancel_event is not None and cancel_event.is_set():
return
@ -257,6 +284,7 @@ def run_safetensors_tool_loop(
if detect_state == _state_draining:
if (
not _tool_succeeded("render_html")
and not _provisional_confirm_gated
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
@ -295,6 +323,7 @@ def run_safetensors_tool_loop(
detect_state = _state_draining
if (
not _tool_succeeded("render_html")
and not _provisional_confirm_gated
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
@ -353,6 +382,7 @@ def run_safetensors_tool_loop(
detect_state = _state_draining
if (
not _tool_succeeded("render_html")
and not _provisional_confirm_gated
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
@ -460,7 +490,8 @@ def run_safetensors_tool_loop(
tool_protocol_active = False,
),
}
if provisional_render_html_started:
if provisional_render_html_started and not provisional_resolved:
provisional_resolved = True
yield {
"type": "tool_end",
"tool_name": "render_html",
@ -503,6 +534,18 @@ def run_safetensors_tool_loop(
if content_text and not assistant_appended:
conversation.append(assistant_msg)
assistant_appended = True
if provisional_match and not provisional_resolved:
# A provisional render_html card is already on screen for
# this id; close it so it never dangles when the controller
# turns the call into an internal no-op (duplicate / repeat).
provisional_resolved = True
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
"tool_call_id": decision.tool_call_id,
"result": "",
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
logger.info(
@ -541,6 +584,8 @@ def run_safetensors_tool_loop(
== "deny"
):
decision_slot = None
if provisional_match:
provisional_resolved = True
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
@ -587,6 +632,8 @@ def run_safetensors_tool_loop(
kb_search_count += 1
completion = tool_controller.record_result(decision, result)
if provisional_match:
provisional_resolved = True
yield completion.tool_end_event()
conversation.append(completion.tool_message())

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

@ -20,7 +20,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_cpp import _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend
from state import tool_approvals
from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision
@ -1325,3 +1325,323 @@ def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypat
assert len(starts) == 2
assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"]
assert calls == [("python", {"code": "print(1)"})]
def _streamed_structured_tool_call(
tool_name: str,
arguments: dict,
call_id: str,
frag: int = 24,
) -> list[str]:
"""A structured tool call whose arguments arrive token-by-token across many
deltas (id + name on the first delta), mirroring how llama-server streams a
large tool-call argument such as a full HTML/code file."""
args_json = json.dumps(arguments)
fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""]
chunks = [
_sse(
{
"tool_calls": [
{
"index": 0,
"id": call_id,
"type": "function",
"function": {"name": tool_name, "arguments": fragments[0]},
}
]
}
)
]
for fragment in fragments[1:]:
chunks.append(_sse({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}))
chunks.append(_done())
return chunks
def test_large_python_tool_call_emits_early_provisional_start(monkeypatch):
"""Regression: a large streamed tool-call argument surfaces a provisional
tool card BEFORE the full arguments finish, so the UI shows progress during
generation instead of a frozen 'Generating...'. (The bug: only render_html
surfaced early; python/terminal/etc. were silent until the call completed.)"""
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
args_json = json.dumps({"code": big_code})
assert len(args_json) > _PROVISIONAL_ARGS_MIN_CHARS
first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_big")
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], 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)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "write code"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
tool_starts = [e for e in events if e.get("type") == "tool_start"]
provisional = [e for e in tool_starts if not e.get("arguments")]
real = [e for e in tool_starts if e.get("arguments", {}).get("code")]
# Exactly one provisional (empty args) and one real (full args), same id so
# the frontend reconciles them into a single card.
assert len(provisional) == 1, tool_starts
assert provisional[0]["tool_name"] == "python"
assert provisional[0]["tool_call_id"] == "call_py_big"
assert provisional[0]["provenance"].get("provisional") is True
assert len(real) == 1
assert real[0]["tool_call_id"] == "call_py_big"
# The provisional card appears before the real (completed) tool_start.
assert events.index(provisional[0]) < events.index(real[0])
assert calls == [("python", {"code": big_code})]
assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events)
def test_small_python_tool_call_has_no_provisional_start(monkeypatch):
"""A small tool-call argument finishes streaming instantly, so it keeps the
existing behavior of a single (real) tool_start with no provisional card."""
first_stream = _structured_tool_call("python", {"code": "print(1)"}, "call_py_small")
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK")
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "x"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
tool_starts = [e for e in events if e.get("type") == "tool_start"]
assert [e for e in tool_starts if not e.get("arguments")] == []
assert len([e for e in tool_starts if e.get("arguments", {}).get("code")]) == 1
def _streamed_parallel_tool_calls(specs, frag: int = 24) -> list[str]:
"""Two or more structured tool calls, each streamed token-by-token across
deltas, one index fully before the next, mirroring how llama-server streams
several parallel tool calls whose arguments are large."""
chunks: list[str] = []
for index, (tool_name, arguments, call_id) in enumerate(specs):
args_json = json.dumps(arguments)
fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""]
chunks.append(
_sse(
{
"tool_calls": [
{
"index": index,
"id": call_id,
"type": "function",
"function": {"name": tool_name, "arguments": fragments[0]},
}
]
}
)
)
for fragment in fragments[1:]:
chunks.append(
_sse({"tool_calls": [{"index": index, "function": {"arguments": fragment}}]})
)
chunks.append(_done())
return chunks
def test_parallel_large_tool_calls_each_emit_provisional_start(monkeypatch):
"""With parallel tool use enabled (the default), every streamed large tool
call surfaces its own provisional card, not just the first one, so the UI
shows progress for each call as its arguments stream."""
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
big_cmd = "echo start\n" + "\n".join(f"echo line {i}" for i in range(60))
assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS
assert len(json.dumps({"command": big_cmd})) > _PROVISIONAL_ARGS_MIN_CHARS
first_stream = _streamed_parallel_tool_calls(
[
("python", {"code": big_code}, "call_py"),
("terminal", {"command": big_cmd}, "call_term"),
]
)
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], 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)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "do both"}],
tools = [
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "terminal"}},
],
max_tool_iterations = 1,
)
)
provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")]
assert sorted(e["tool_call_id"] for e in provisional) == ["call_py", "call_term"]
assert all(e["provenance"].get("provisional") is True for e in provisional)
# Both calls actually executed (parallel tool use is enabled by default).
assert sorted(name for name, _ in calls) == ["python", "terminal"]
def test_parallel_disabled_suppresses_provisional_for_later_calls(monkeypatch):
"""When parallel tool use is disabled the downstream truncates to the first
call, so only the first streamed call may surface a provisional; a later
call must not get a card that could never reconcile or be closed."""
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
big_cmd = "echo start\n" + "\n".join(f"echo line {i}" for i in range(60))
first_stream = _streamed_parallel_tool_calls(
[
("python", {"code": big_code}, "call_py"),
("terminal", {"command": big_cmd}, "call_term"),
]
)
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], 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)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "do both"}],
tools = [
{"type": "function", "function": {"name": "python"}},
{"type": "function", "function": {"name": "terminal"}},
],
max_tool_iterations = 1,
disable_parallel_tool_use = True,
)
)
provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")]
assert [e["tool_call_id"] for e in provisional] == ["call_py"]
# Only the first call executes when parallel use is disabled.
assert calls == [("python", {"code": big_code})]
# The lone provisional is closed exactly once (no dangling card).
closing = [
e for e in events if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py"
]
assert len(closing) == 1
def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
"""If llama-server drops mid tool-call after a provisional card is shown, the
loop must close that card before surfacing the error so the UI never leaves a
tool spinning forever."""
import httpx
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
fragments = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_err")
# Drop the trailing [DONE]; raise a connection error after the fragments
# stream (and after the provisional card has been emitted).
fragments = fragments[:-1]
def raising_stream():
for chunk in fragments:
yield chunk
raise httpx.ConnectError("connection lost mid stream")
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [raising_stream()], payloads)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK")
collected: list[dict] = []
raised = False
gen = backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "write code"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
try:
for event in gen:
collected.append(event)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
provisional = [e for e in collected if e.get("type") == "tool_start" and not e.get("arguments")]
assert len(provisional) == 1
assert provisional[0]["tool_call_id"] == "call_py_err"
# The provisional card is closed before the error propagates.
closing = [
e
for e in collected
if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py_err"
]
assert len(closing) == 1
# The closing card is marked as an error, not an empty success, so the UI
# renders it as failed.
assert "Error" in (closing[0].get("result") or "")
def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):
"""llama.cpp can stream a tool call whose id is an empty string. A provisional
card keyed by "" cannot reconcile with the real tool_start (the frontend mints
its own id per event), so it must not be emitted -- otherwise the empty card
would dangle. The real call must still execute normally."""
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS
# Same large streamed call as the provisional test, but with an empty id.
first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "")
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], 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)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "write code"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
# No provisional card (empty-args tool_start) was surfaced for the empty id.
provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")]
assert provisional == []
# The real call still executes despite the missing id.
assert calls == [("python", {"code": big_code})]

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

@ -393,6 +393,138 @@ class TestLoopBasic:
assert exec_fn.calls[0][0] == "render_html"
assert "<!doctype html>" in exec_fn.calls[0][1]["code"]
def test_render_html_confirmation_gate_suppresses_early_provisional(self, monkeypatch):
"""When a human confirmation gate is active, render_html must not surface
an early provisional tool_start: that card (keyed by tool_call_id, no
approval) would show the tool 'running' before the user approves. The
gated real tool_start is the first signal the UI receives instead."""
monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "approval-rh")
monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object())
monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "allow")
exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
turn_iter = iter(
[
[
"<function=render_html>",
"<parameter=code><!doctype html><html>",
"<body>Hi</body></html></parameter></function>",
],
["Done."],
]
)
def _gen(_messages):
chunks = next(turn_iter)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "make html"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
execute_tool = exec_fn,
confirm_tool_calls = True,
session_id = "sess",
max_tool_iterations = 3,
)
events = _collect_events(loop)
tool_starts = [e for e in events if e["type"] == "tool_start"]
# No early provisional (empty-args) card while confirmation is pending.
assert [e for e in tool_starts if e.get("arguments") == {}] == []
# The real, gated tool_start still surfaces with the full arguments.
real = [e for e in tool_starts if e.get("arguments", {}).get("code")]
assert len(real) == 1
assert real[0].get("awaiting_confirmation") is True
assert "<!doctype html>" in real[0]["arguments"]["code"]
assert exec_fn.calls[0][0] == "render_html"
def test_render_html_bypass_permissions_keeps_early_provisional(self, monkeypatch):
"""bypass_permissions wins over the confirm gate, so the early provisional
card is preserved (no human approval is required)."""
exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
turn_iter = iter(
[
[
"<function=render_html>",
"<parameter=code><!doctype html><html>",
"<body>Hi</body></html></parameter></function>",
],
["Done."],
]
)
def _gen(_messages):
chunks = next(turn_iter)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "make html"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
execute_tool = exec_fn,
confirm_tool_calls = True,
bypass_permissions = True,
session_id = "sess",
max_tool_iterations = 3,
)
events = _collect_events(loop)
tool_starts = [e for e in events if e["type"] == "tool_start"]
assert len(tool_starts) == 2
assert tool_starts[0]["arguments"] == {}
assert "<!doctype html>" in tool_starts[1]["arguments"]["code"]
def test_render_html_provisional_card_closed_on_generator_exception(self):
"""If the model generator raises mid-stream after a provisional render_html
card was surfaced, the loop must close that card as errored before the
exception propagates, so the UI never leaves a tool spinning forever."""
exec_fn = FakeExecuteTool([])
def _gen(_messages):
acc = ""
for chunk in ["<function=render_html>", "<parameter=code><!doctype html><html>"]:
acc += chunk
yield acc
raise RuntimeError("model pipeline exploded")
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "make html"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
execute_tool = exec_fn,
)
collected: list[dict] = []
raised = False
try:
for event in loop:
collected.append(event)
except RuntimeError as exc:
raised = True
assert "exploded" in str(exc)
assert raised
provisional = [
e for e in collected if e["type"] == "tool_start" and e.get("arguments") == {}
]
assert len(provisional) == 1
# The provisional card is closed (as an error) before the exception
# propagates, so it never dangles.
closing = [
e
for e in collected
if e["type"] == "tool_end" and e.get("tool_call_id") == provisional[0]["tool_call_id"]
]
assert len(closing) == 1
assert "Error" in (closing[0].get("result") or "")
def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(self):
loop, exec_fn = _make_loop(
turns = [