diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5b32103dc8..144aa1fd37 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -92,6 +92,7 @@ from utils.subprocess_compat import ( from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs from core.inference.tool_call_parser import ( MAX_ACT_REPROMPTS as _MAX_REPROMPTS, + NUDGE_TOOL_CALLS_STATUS as _NUDGE_TOOL_CALLS_STATUS, REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS, is_short_intent_without_action as _is_short_intent_without_action, reprompt_to_act_message as _reprompt_to_act_message, @@ -12419,7 +12420,10 @@ class LlamaCppBackend: _it_r = _iter_timings or {} _accumulated_predicted_ms += _it_r.get("predicted_ms", 0) _accumulated_predicted_n += _it_r.get("predicted_n", 0) + # Blank first (the route resets its text cursor only on an + # empty status), then the badge so the retry is not a hang. yield {"type": "status", "text": ""} + yield {"type": "status", "text": _NUDGE_TOOL_CALLS_STATUS} continue if _forced_tool_call_pending: diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index b593bc119b..3057f7c2ac 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -35,6 +35,7 @@ from core.inference.tool_call_parser import ( _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, MAX_ACT_REPROMPTS, + NUDGE_TOOL_CALLS_STATUS, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, @@ -1032,9 +1033,10 @@ def run_safetensors_tool_loop( "content": reprompt_to_act_message(tool_hint), } ) - # Empty status clears the badge and resets the route's - # per-turn text cursor before the re-prompted turn streams. + # Blank first: it clears the badge and resets the route's per-turn + # text cursor. The badge then shows the pause is a re-prompt, not a stall. yield {"type": "status", "text": ""} + yield {"type": "status", "text": NUDGE_TOOL_CALLS_STATUS} continue # Final answer. If a literal tool marker in prose was buffered but diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 9b6b0a7773..4c3fe234ae 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -183,6 +183,9 @@ INTENT_SIGNAL = re.compile( # times since #5620); safetensors and MLX inherit the same cap from here. MAX_ACT_REPROMPTS = 3 REPROMPT_MAX_CHARS = 2000 +# Composer badge while a hidden re-prompted turn regenerates, else the UI looks +# hung. Matched exactly by the frontend (utils/tool-status.ts); keep in sync. +NUDGE_TOOL_CALLS_STATUS = "Nudging tool calls" def is_short_intent_without_action(text: str) -> bool: diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index c629ff3be4..cbd1b07505 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -26,6 +26,7 @@ from core.inference.llama_cpp import ( _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend, ) +from core.inference.tool_call_parser import NUDGE_TOOL_CALLS_STATUS from state import tool_approvals from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision @@ -1841,6 +1842,140 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): assert len(payloads) == 3 +def _status_texts(events: list[dict]) -> list[str]: + return [event["text"] for event in events if event.get("type") == "status"] + + +_WEB_SEARCH_TOOL = { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web.", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + + +def _nudge_then_search_streams() -> list[list[str]]: + """Stall, then a re-prompted turn that finally searches, then the answer.""" + + return [ + [_sse({"content": "I will search the web now."}), _done()], + [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "red square"}), + }, + } + ] + } + ), + _done(), + ], + [_sse({"content": "Final answer: the square is red."}), _done()], + ] + + +def test_plan_without_action_nudge_is_announced_on_the_status_channel(monkeypatch): + """The re-prompted turn is hidden, so without a badge the UI looks frozen.""" + + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, _nudge_then_search_streams(), payloads) + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: "Search results: red is #f00.", + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "What colour is the square?"}], + tools = [_WEB_SEARCH_TOOL], + max_tool_iterations = 2, + ) + ) + + statuses = _status_texts(events) + assert NUDGE_TOOL_CALLS_STATUS in statuses + index = statuses.index(NUDGE_TOOL_CALLS_STATUS) + # Blank first: the route resets its text cursor only on an empty status. + # index > 0 matters: at 0, statuses[-1] wraps to the terminal clear. + assert index > 0 and statuses[index - 1] == "" + assert statuses[index + 1].startswith("Searching:") + assert statuses[-1] == "" + + +def test_plan_without_action_nudge_status_clears_when_the_retry_just_answers(monkeypatch): + streams = [ + [_sse({"content": "I will search the web now."}), _done()], + [_sse({"content": "No search needed. Final answer: the square is red."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "What colour is the square?"}], + tools = [_WEB_SEARCH_TOOL], + max_tool_iterations = 2, + ) + ) + + statuses = _status_texts(events) + assert NUDGE_TOOL_CALLS_STATUS in statuses + assert statuses[-1] == "" + + +def test_direct_answer_never_shows_the_nudge_status(monkeypatch): + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [[_sse({"content": "The square is red."}), _done()]], + payloads, + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "What colour is the square?"}], + tools = [_WEB_SEARCH_TOOL], + max_tool_iterations = 2, + ) + ) + + assert NUDGE_TOOL_CALLS_STATUS not in _status_texts(events) + + +def test_nudge_status_absent_when_nudging_is_disabled(monkeypatch): + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, _nudge_then_search_streams(), payloads) + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: "Search results: red is #f00.", + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "What colour is the square?"}], + tools = [_WEB_SEARCH_TOOL], + max_tool_iterations = 2, + nudge_tool_calls = False, + ) + ) + + assert NUDGE_TOOL_CALLS_STATUS not in _status_texts(events) + assert len(payloads) == 1 + + def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): streams = [ _structured_tool_call("python", {"code": "print(1)"}, "call_py"), diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 1043005f64..2e7e99fbba 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -24,6 +24,7 @@ from core.inference.safetensors_agentic import ( strip_tool_markup_streaming, ) from core.inference.tool_call_parser import ( + NUDGE_TOOL_CALLS_STATUS, RAG_MAX_SEARCHES_PER_TURN, has_tool_signal, parse_tool_calls_from_text, @@ -2231,6 +2232,24 @@ def test_reprompt_names_only_active_tools_not_hardcoded(): assert "python" not in reprompt["content"] +def test_reprompt_is_announced_on_the_status_channel(): + # The re-prompted turn is hidden, so the badge is the only sign of life. + # Blank still comes first: the route resets its text cursor only on that. + _captured, events = _reprompt_loop(auto_heal_tool_calls = True) + statuses = [e["text"] for e in events if e["type"] == "status"] + assert NUDGE_TOOL_CALLS_STATUS in statuses + index = statuses.index(NUDGE_TOOL_CALLS_STATUS) + # index > 0 matters: at 0, statuses[-1] wraps to the terminal clear. + assert index > 0 and statuses[index - 1] == "" + assert statuses[-1] == "" + + +def test_reprompt_status_absent_without_a_nudge(): + _captured, events = _reprompt_loop(auto_heal_tool_calls = False) + statuses = [e["text"] for e in events if e["type"] == "status"] + assert NUDGE_TOOL_CALLS_STATUS not in statuses + + def test_reprompt_suppressed_when_auto_heal_disabled(): # With Auto-Heal off the safetensors nudge must stay silent for backend parity # with the GGUF loop, so only the single initial generation runs. diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 6da8126421..9b3c7aa79e 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -91,6 +91,7 @@ import { useResearchRunStore, } from "@/features/chat/stores/research-run-store"; import { parseExternalModelId } from "@/features/chat/external-providers"; +import { toolStatusKind } from "@/features/chat/utils/tool-status"; import { McpComposerButton } from "@/features/chat/mcp-composer-button"; import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities"; import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled"; @@ -2847,15 +2848,28 @@ const ToolStatusDisplay: FC = () => { } // From the store's start time, so returning to the conversation resumes rather than restarting. const elapsed = Math.max(0, Math.floor((now - startedAt) / 1000)); - const isRunning = toolStatus.startsWith("Running"); - const StatusIcon = isRunning ? TerminalIcon : GlobeIcon; + const kind = toolStatusKind(toolStatus); + const isNudging = kind === "nudge"; + const StatusIcon = kind === "terminal" ? TerminalIcon : GlobeIcon; return (
-
- +
+ {isNudging ? ( + // label, not the default "Loading": the spinner is the badge's only + // role="status" region, so its name is what gets announced. + + ) : ( + + )} {toolStatus} {elapsed}s
diff --git a/studio/frontend/src/features/chat/utils/tool-status.ts b/studio/frontend/src/features/chat/utils/tool-status.ts new file mode 100644 index 0000000000..16c86bbd49 --- /dev/null +++ b/studio/frontend/src/features/chat/utils/tool-status.ts @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** Mirrors NUDGE_TOOL_CALLS_STATUS in backend core/inference/tool_call_parser.py; keep in sync. */ +export const NUDGE_TOOL_CALLS_STATUS = "Nudging tool calls"; + +export type ToolStatusKind = "nudge" | "terminal" | "web"; + +/** Which glyph the badge shows: exact match for the nudge, "Running" prefix for sandbox tools, globe otherwise. */ +export function toolStatusKind(status: string): ToolStatusKind { + if (status === NUDGE_TOOL_CALLS_STATUS) { + return "nudge"; + } + return status.startsWith("Running") ? "terminal" : "web"; +} diff --git a/studio/frontend/tests/tool-status.test.ts b/studio/frontend/tests/tool-status.test.ts new file mode 100644 index 0000000000..b20585bd08 --- /dev/null +++ b/studio/frontend/tests/tool-status.test.ts @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + NUDGE_TOOL_CALLS_STATUS, + toolStatusKind, +} from "../src/features/chat/utils/tool-status.ts"; + +test("the nudge status is the exact string the backend sends", () => { + // Mirrors tool_call_parser.py, so a reword on either side must break here. + assert.equal(NUDGE_TOOL_CALLS_STATUS, "Nudging tool calls"); + assert.equal(toolStatusKind(NUDGE_TOOL_CALLS_STATUS), "nudge"); +}); + +test("sandbox tools keep the terminal glyph", () => { + for (const status of [ + "Running Python: print(1)", + "Running Python...", + "Running: ls -la", + "Running command...", + ]) { + assert.equal(toolStatusKind(status), "terminal", status); + } +}); + +test("every other status keeps the globe", () => { + for (const status of [ + "Searching: red square", + "Reading: unsloth.ai", + "Reading page...", + "Searching documents: quarterly report", + "Calling: get_weather", + ]) { + assert.equal(toolStatusKind(status), "web", status); + } +}); + +test("a status that merely mentions nudging is not the nudge itself", () => { + // Exact match only: a tool named after the phrase must not steal the spinner. + assert.equal(toolStatusKind("Calling: Nudging tool calls"), "web"); + assert.equal(toolStatusKind("Nudging tool calls again"), "web"); +});