Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls (#5869)
* Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix race in tool-call confirmation gate * Studio: gate built-in tool calls and harden the confirmation handshake The Allow / Always allow / Deny controls only lived in the fallback tool card, but the built-in tools (web search, python, terminal, code execution, image generation) render with their own components and so never showed the buttons. Those calls paused after tool_start with no way to approve them, hanging until the 1 hour timeout. Only MCP tools, which use the fallback renderer, actually worked. Render the controls for every tool card by wrapping each registered tool component (and the fallback) in thread.tsx with a shared ToolConfirmationControls, so the gate applies uniformly. Also make the handshake robust: - The gate keys on a per-call approval_id minted by the backend and echoed in tool_start, instead of session_id alone, so a stale or concurrent confirmation can no longer resolve the wrong call. - The approval slot is registered before tool_start is yielded, closing the race where a fast click or an auto "Always allow" could reach the backend before the waiter existed. - The frontend resolves with the same session id the request was sent with (plus the approval_id), fixing the new-thread mismatch where the confirmation targeted a different session than the blocked stream. - The confirm endpoint returns {resolved}; the UI keeps the buttons and shows a retry hint until the backend confirms a match, instead of hiding them on a failed or mistargeted post. - The gate runs after the disabled-tool and duplicate-call checks, so a call that will not execute is not put up for approval. A denied call is still excluded from duplicate detection, so re-issuing and approving it works. - "Always allow" is scoped per session to match the backend gate. Add backend tests for the approval registry, the SSE no-deadlock handshake, and the loop integration (allow, deny, disabled, duplicate, re-issue after deny). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move "Confirm tool calls" to the Tools section * Studio: Keep tool group open while a tool call awaits confirmation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix tool confirmation session scope for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix confirmation follow-ups for PR #5869 * Apply pre-commit formatting for PR #5869 * Fix confirmation cleanup for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden confirmation lookups for PR #5869 * Studio: make the tool-call confirmation decision immutable resolve_tool_decision accepted a second confirmation for the same approval_id and overwrote slot["decision"] in the window before the waiter reads it and pops the slot, so a duplicate or out-of-order POST could flip an Allow to Deny (and returned a misleading resolved:true). Reject once the slot's event is already set so the first decision wins. Adds a regression test. * Fix/adjust tool confirmations for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: wasimysaid <wasimysdev@gmail.com>
This commit is contained in:
oobabooga
GitHub
parent
de0c5a2f09
commit
7f2986a413
22 changed files with 1694 additions and 29 deletions
|
|
@ -21,6 +21,8 @@ if _BACKEND_DIR not in sys.path:
|
|||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
from state import tool_approvals
|
||||
from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision
|
||||
|
||||
|
||||
def _sse(delta: dict) -> str:
|
||||
|
|
@ -70,6 +72,27 @@ def _tool_names(payload: dict) -> list[str]:
|
|||
]
|
||||
|
||||
|
||||
def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]:
|
||||
return [
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_name,
|
||||
"arguments": json.dumps(arguments),
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
_done(),
|
||||
]
|
||||
|
||||
|
||||
def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
|
||||
"""llama-server may emit content first and then native delta.tool_calls.
|
||||
|
||||
|
|
@ -1149,3 +1172,151 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
|
|||
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
|
||||
assert content_texts == ["I will use render_html now.", "Final note after tool."]
|
||||
assert len(payloads) == 3
|
||||
|
||||
|
||||
def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch):
|
||||
streams = [
|
||||
_structured_tool_call("python", {"code": "print(1)"}, "call_py"),
|
||||
[_sse({"content": "Done."}), _done()],
|
||||
]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
|
||||
calls: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
calls.append((name, arguments))
|
||||
return "OK"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1")
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_cpp.begin_tool_decision",
|
||||
lambda *_a, **_k: object(),
|
||||
)
|
||||
monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow")
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run python"}],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
max_tool_iterations = 1,
|
||||
confirm_tool_calls = True,
|
||||
session_id = "sess",
|
||||
)
|
||||
)
|
||||
|
||||
starts = [event for event in events if event.get("type") == "tool_start"]
|
||||
assert len(starts) == 1
|
||||
assert starts[0]["approval_id"]
|
||||
assert starts[0]["awaiting_confirmation"] is True
|
||||
assert calls == [("python", {"code": "print(1)"})]
|
||||
assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events)
|
||||
|
||||
|
||||
def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch):
|
||||
approval_id = "approval-close"
|
||||
streams = [_structured_tool_call("python", {"code": "print(1)"}, "call_py")]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"core.inference.tools.execute_tool",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("tool should not run")),
|
||||
)
|
||||
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: approval_id)
|
||||
|
||||
with tool_approvals._lock:
|
||||
tool_approvals._pending.clear()
|
||||
|
||||
gen = backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run python"}],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
max_tool_iterations = 1,
|
||||
confirm_tool_calls = True,
|
||||
session_id = "sess",
|
||||
)
|
||||
try:
|
||||
assert next(gen)["type"] == "status"
|
||||
start = next(gen)
|
||||
assert start["type"] == "tool_start"
|
||||
assert start["approval_id"] == approval_id
|
||||
with tool_approvals._lock:
|
||||
assert approval_id in tool_approvals._pending
|
||||
finally:
|
||||
gen.close()
|
||||
|
||||
with tool_approvals._lock:
|
||||
assert approval_id not in tool_approvals._pending
|
||||
assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False
|
||||
|
||||
|
||||
def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch):
|
||||
streams = [[_sse({"content": "Done."}), _done()]]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
|
||||
def fail_autoinject(*_args, **_kwargs):
|
||||
raise AssertionError("RAG autoinject must not run before approval")
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "use docs"}],
|
||||
tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}],
|
||||
max_tool_iterations = 1,
|
||||
confirm_tool_calls = True,
|
||||
session_id = "sess",
|
||||
rag_scope = {"thread_id": "t1"},
|
||||
)
|
||||
)
|
||||
|
||||
assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
|
||||
|
||||
|
||||
def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch):
|
||||
same_call = _structured_tool_call("python", {"code": "print(1)"}, "call_py")
|
||||
streams = [
|
||||
same_call,
|
||||
_structured_tool_call("python", {"code": "print(1)"}, "call_py_retry"),
|
||||
[_sse({"content": "Done."}), _done()],
|
||||
]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
|
||||
calls: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
calls.append((name, arguments))
|
||||
return "OK"
|
||||
|
||||
decisions = iter(["deny", "allow"])
|
||||
approvals = iter(["approval-1", "approval-2"])
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals))
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_cpp.begin_tool_decision",
|
||||
lambda *_a, **_k: object(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_cpp.wait_tool_decision",
|
||||
lambda *_a, **_k: next(decisions),
|
||||
)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run python"}],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
max_tool_iterations = 2,
|
||||
confirm_tool_calls = True,
|
||||
session_id = "sess",
|
||||
)
|
||||
)
|
||||
|
||||
starts = [event for event in events if event.get("type") == "tool_start"]
|
||||
ends = [event for event in events if event.get("type") == "tool_end"]
|
||||
assert len(starts) == 2
|
||||
assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"]
|
||||
assert calls == [("python", {"code": "print(1)"})]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue