diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6956d2d38b..40a10eecd3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -31,6 +31,7 @@ import uuid from dataclasses import dataclass, replace from pathlib import Path from typing import ( + Any, Callable, Collection, Generator, @@ -20230,6 +20231,9 @@ class LlamaCppBackend: in_thinking = False has_content_tokens = False tool_calls_acc = {} # Structured delta.tool_calls fragments + # accumulator key each index maps to now, (index, call_id) after a fork + tool_calls_open_key: dict[Any, Any] = {} + tool_calls_synthetic_ids: set[Any] = set() # slots still on a placeholder id has_structured_tc = False _iter_usage = None _iter_timings = None @@ -20332,31 +20336,53 @@ class LlamaCppBackend: yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) - if idx not in tool_calls_acc: - tool_calls_acc[idx] = { - "id": tc_d.get("id", f"call_{idx}"), + tc_id = tc_d.get("id") + # a second call at a reused index forks its own slot; bare fragments follow whichever call owns the index now + acc_key = tool_calls_open_key.get(idx, idx) + if isinstance(tc_id, str) and tc_id: + open_id = tool_calls_acc.get(acc_key, {}).get("id") + # a synthetic placeholder is not a rival call, so a first real id still lands on the slot that is open + if ( + open_id + and acc_key not in tool_calls_synthetic_ids + and open_id != tc_id + ): + first_id = tool_calls_acc.get(idx, {}).get("id") + acc_key = idx if first_id == tc_id else (idx, tc_id) + tool_calls_open_key[idx] = acc_key + if acc_key not in tool_calls_acc: + tool_calls_acc[acc_key] = { + "id": tc_id or f"call_{idx}", "type": "function", "function": { "name": "", "arguments": "", }, } - elif tc_d.get("id"): + if not tc_id: + tool_calls_synthetic_ids.add(acc_key) + elif tc_id: # Update ID if a real one # arrives on a later delta. - tool_calls_acc[idx]["id"] = tc_d["id"] + tool_calls_acc[acc_key]["id"] = tc_id + tool_calls_synthetic_ids.discard(acc_key) func = tc_d.get("function", {}) if func.get("name"): - tool_calls_acc[idx]["function"]["name"] += func["name"] - if func.get("arguments"): - tool_calls_acc[idx]["function"]["arguments"] += func[ - "arguments" + tool_calls_acc[acc_key]["function"]["name"] += func[ + "name" ] - current_name = tool_calls_acc[idx]["function"].get( + if func.get("arguments"): + tool_calls_acc[acc_key]["function"]["arguments"] += ( + func["arguments"] + ) + current_name = tool_calls_acc[acc_key]["function"].get( "name", "" ) fallback_id = f"call_{idx}" - current_id = tool_calls_acc[idx].get("id", fallback_id) + current_id = tool_calls_acc[acc_key].get("id", fallback_id) + is_first_accumulated_call = ( + next(iter(tool_calls_acc)) == acc_key + ) already_started = ( current_id in provisional_started_tool_calls ) @@ -20390,7 +20416,7 @@ class LlamaCppBackend: ) # Keep small-argument tools on the normal path. _args_len = len( - tool_calls_acc[idx]["function"].get("arguments", "") + tool_calls_acc[acc_key]["function"].get("arguments", "") ) _payload_is_large = ( current_name == "render_html" @@ -20398,7 +20424,10 @@ class LlamaCppBackend: ) if ( current_name - and (idx == 0 or not disable_parallel_tool_use) + and ( + is_first_accumulated_call + or not disable_parallel_tool_use + ) and has_real_id and not already_started and not _is_completed_one_shot @@ -20429,9 +20458,9 @@ class LlamaCppBackend: if current_id in provisional_started_tool_calls: if current_id not in arg_streamed_tool_call_ids: arg_streamed_tool_call_ids.add(current_id) - _args_backlog = tool_calls_acc[idx]["function"].get( - "arguments", "" - ) + _args_backlog = tool_calls_acc[acc_key][ + "function" + ].get("arguments", "") if _args_backlog: yield { "type": "tool_args", @@ -20940,10 +20969,11 @@ class LlamaCppBackend: if has_structured_tc: # Drop incomplete fragments (e.g. from max_tokens # truncation or disconnect). + # first-seen order: a call forked onto a reused index arrived after every call already open, so it must execute and replay last tool_calls = [ - tool_calls_acc[i] - for i in sorted(tool_calls_acc) - if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) + call + for call in tool_calls_acc.values() + if (call.get("function", {}).get("name", "").strip()) ] or None if not tool_calls: # Unconditional re-parse: we only reach DRAINING when the buffer looked like a diff --git a/studio/backend/core/inference/studio_tool_loop.py b/studio/backend/core/inference/studio_tool_loop.py index d84c2c8908..9acbefa8b2 100644 --- a/studio/backend/core/inference/studio_tool_loop.py +++ b/studio/backend/core/inference/studio_tool_loop.py @@ -352,6 +352,9 @@ class _Turn: by_index: dict[Any, dict[str, Any]] = field(default_factory = dict) order: list[Any] = field(default_factory = list) + # call key each delta index maps to: the index itself until a second call + # forks off it, then (index, call_id). + open_key_by_index: dict[int, Any] = field(default_factory = dict) last_index: int | None = None round: int = 0 healed: list[dict[str, Any]] = field(default_factory = list) @@ -466,15 +469,23 @@ class _Turn: # continue the call that is already open instead. index = self.last_index if self.last_index is not None else len(self.order) call_id = raw_call.get("id") - key: Any = index + # continue whichever call owns this index now: index restarts at 0 + # for every tool round, so after a fork the bare argument fragments + # belong to the newer call. + key: Any = self.open_key_by_index.get(index, index) if isinstance(call_id, str) and call_id: - open_id = self.by_index.get(index, {}).get("id") + open_id = self.by_index.get(key, {}).get("id") if open_id and open_id != call_id: # Two distinct calls reported at the same index. Merging them # concatenates their argument JSON into one unparseable blob # and loses an intent, so key the second on its own id. - key = (index, call_id) + # A fragment that names the call this index opened first goes + # back to it: an id beats the latest-index mapping, which only + # exists to place the fragments that carry no id. + first_id = self.by_index.get(index, {}).get("id") + key = index if first_id == call_id else (index, call_id) self.last_index = index + self.open_key_by_index[index] = key if key not in self.by_index: self.by_index[key] = { "id": "", diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 6dd270fbc7..9b83856b58 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -4939,3 +4939,283 @@ def test_provisional_text_card_closed_when_parse_fails(monkeypatch): assert executed == [] # nothing parsed, nothing ran assert ends, "provisional card left dangling (no tool_end)" assert ends[-1]["tool_call_id"] == "call_0" + + +def _tool_call_fragment( + index: int, + arguments: str, + call_id: str | None = None, +) -> str: + delta: dict = {"index": index, "function": {"arguments": arguments}} + if call_id is not None: + delta["id"] = call_id + return _sse({"tool_calls": [delta]}) + + +def _tool_call_opening(index: int, call_id: str, name: str, arguments: str) -> str: + return _sse( + { + "tool_calls": [ + { + "index": index, + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + ] + } + ) + + +def test_second_structured_call_at_one_index_keeps_its_own_fragments(monkeypatch): + """Two tool rounds in one llama-server response, both streamed at index 0. + + llama-server restarts ``delta.tool_calls[].index`` at 0 for every round + while giving each call its own id, and the continuation fragments carrying + the rest of the arguments arrive bare. Keying the accumulator on the index + alone appended round two's name and argument tail to round one, leaving a + single ``web_searchweb_search`` entry that matched no enabled tool, so + neither call ran. + """ + + stream = [ + _tool_call_opening(0, "call_a", "web_search", '{"query":'), + _tool_call_fragment(0, '"first"}'), + _tool_call_opening(0, "call_b", "web_search", '{"query":'), + _tool_call_fragment(0, '"second"}'), + _finish("tool_calls"), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append({"name": name, "arguments": arguments}) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search twice"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [ + {"name": "web_search", "arguments": {"query": "first"}}, + {"name": "web_search", "arguments": {"query": "second"}}, + ] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [ + "call_a", + "call_b", + ] + + # The replayed conversation must list both calls with their own arguments. + asst = next(m for m in payloads[1]["messages"] if m.get("tool_calls")) + assert [tc["id"] for tc in asst["tool_calls"]] == ["call_a", "call_b"] + assert [tc["function"]["arguments"] for tc in asst["tool_calls"]] == [ + '{"query":"first"}', + '{"query":"second"}', + ] + + +def test_structured_fragment_naming_its_call_goes_back_to_that_call(monkeypatch): + """Two calls at index 0 with the id repeated on every argument fragment. + + The latest-index mapping only exists to place fragments that carry no id, + so a fragment naming the call the index opened first has to go back to it + rather than fork a third slot. Forking left that call with truncated JSON + and dropped the fragment for having no function name. + """ + + stream = [ + _tool_call_opening(0, "call_a", "web_search", '{"query":'), + _tool_call_opening(0, "call_b", "web_search", '{"query":'), + _tool_call_fragment(0, '"first"}', call_id = "call_a"), + _tool_call_fragment(0, '"second"}', call_id = "call_b"), + _finish("tool_calls"), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append({"name": name, "arguments": arguments}) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search twice"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [ + {"name": "web_search", "arguments": {"query": "first"}}, + {"name": "web_search", "arguments": {"query": "second"}}, + ] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [ + "call_a", + "call_b", + ] + + +def test_structured_call_id_arriving_after_the_opening_delta_updates_that_call(monkeypatch): + """llama-server can open a call with no id and send the real one later. + + The opening slot holds the synthetic ``call_0`` until then, and treating + that placeholder as a rival id forked a second nameless slot: the fork was + dropped for having no function name and the original call ran on truncated + arguments. + """ + + stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "type": "function", + "function": {"name": "web_search", "arguments": '{"query":'}, + } + ] + } + ), + _tool_call_fragment(0, '"late"}', call_id = "call_late"), + _finish("tool_calls"), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append({"name": name, "arguments": arguments}) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [{"name": "web_search", "arguments": {"query": "late"}}] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == ["call_late"] + + +def test_structured_call_forked_onto_a_reused_index_executes_last(monkeypatch): + """A first round at indices 0 and 1, then a second round back at index 0. + + The fork belongs at the end: it arrived after both first-round calls, and + running it ahead of the index-1 call reorders side effects for stateful + tools. Grouping every fork next to the index it reused did exactly that. + """ + + stream = [ + _tool_call_opening(0, "call_a", "web_search", '{"query":"a"}'), + _tool_call_opening(1, "call_b", "web_search", '{"query":"b"}'), + _tool_call_opening(0, "call_c", "web_search", '{"query":"c"}'), + _finish("tool_calls"), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append(arguments) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search three times"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [{"query": "a"}, {"query": "b"}, {"query": "c"}] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [ + "call_a", + "call_b", + "call_c", + ] + asst = next(m for m in payloads[1]["messages"] if m.get("tool_calls")) + assert [tc["id"] for tc in asst["tool_calls"]] == ["call_a", "call_b", "call_c"] + + +def test_parallel_disabled_suppresses_provisional_for_reused_index(monkeypatch): + """A later call at reused index 0 is not the first accumulated call. + + ``parallel_tool_calls=false`` permits a provisional card only for the call + that can execute. Checking the raw index admitted every fork at index 0, + leaving a card for the truncated call that could only close empty. + """ + + 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)) + python_args = json.dumps({"code": big_code}) + terminal_args = json.dumps({"command": big_cmd}) + stream = [ + _tool_call_opening(0, "call_py", "python", python_args[:-1]), + _tool_call_fragment(0, python_args[-1]), + _tool_call_opening(0, "call_term", "terminal", terminal_args[:-1]), + _tool_call_fragment(0, terminal_args[-1]), + _finish("tool_calls"), + _done(), + ] + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [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 the first only"}], + tools = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "terminal"}}, + ], + max_tool_iterations = 1, + disable_parallel_tool_use = True, + permission_mode = "off", + ) + ) + + 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"] + assert calls == [("python", {"code": big_code})] + assert not [ + e + for e in events + if e.get("tool_call_id") == "call_term" + and e.get("type") in {"tool_start", "tool_args", "tool_end"} + ] diff --git a/studio/backend/tests/test_studio_tool_loop.py b/studio/backend/tests/test_studio_tool_loop.py index 982ef6b451..a32d3e5012 100644 --- a/studio/backend/tests/test_studio_tool_loop.py +++ b/studio/backend/tests/test_studio_tool_loop.py @@ -1036,3 +1036,71 @@ def test_a_skipped_duplicate_closes_the_card_the_provider_already_painted(execut # every tool_end has a matching tool_start. starts = _events(lines, "tool_start") assert [start["tool_call_id"] for start in starts] == ["call_a", "call_a"] + + +def test_a_second_call_at_one_index_keeps_its_own_argument_fragments(executed): + """Two tool rounds in one response, both streamed at index 0. + + Providers restart ``delta.tool_calls[].index`` at 0 for every round while + giving each call its own id, and the continuation fragments carrying the + rest of the arguments are sent bare. Routing those by index alone appended + round two's tail to round one, producing an unparseable blob and running + both tools on the wrong arguments. + """ + transport = FakeTransport( + [ + [ + _sse({"tool_calls": [_call_delta(0, "call_a", "web_search", '{"query":')]}), + _sse({"tool_calls": [{"index": 0, "function": {"arguments": '"first"}'}}]}), + _sse({"tool_calls": [_call_delta(0, "call_b", "web_search", '{"query":')]}), + _sse({"tool_calls": [{"index": 0, "function": {"arguments": '"second"}'}}]}), + _sse(finish = "tool_calls"), + _DONE, + ], + [_sse({"content": "done"}), _sse(finish = "stop"), _DONE], + ], + heals = False, + ) + _run(transport) + + assert [call["arguments"] for call in executed] == [{"query": "first"}, {"query": "second"}] + + +def test_a_fragment_naming_its_call_goes_back_to_that_call(executed): + """Two calls at index 0 with the id repeated on every argument fragment. + + The latest-index mapping only exists to place fragments that carry no id, so + a fragment that names the call the index opened first has to go back to it + rather than fork a third slot. Forking left that call with truncated JSON, + which reaches the tool as ``_raw``, and dropped the fragment for having no + function name. + """ + transport = FakeTransport( + [ + [ + _sse({"tool_calls": [_call_delta(0, "call_a", "web_search", '{"query":')]}), + _sse({"tool_calls": [_call_delta(0, "call_b", "web_search", '{"query":')]}), + _sse( + { + "tool_calls": [ + {"index": 0, "id": "call_a", "function": {"arguments": '"first"}'}} + ] + } + ), + _sse( + { + "tool_calls": [ + {"index": 0, "id": "call_b", "function": {"arguments": '"second"}'}} + ] + } + ), + _sse(finish = "tool_calls"), + _DONE, + ], + [_sse({"content": "done"}), _sse(finish = "stop"), _DONE], + ], + heals = False, + ) + _run(transport) + + assert [call["arguments"] for call in executed] == [{"query": "first"}, {"query": "second"}]