From 02f63cd4192cb05f42dcb7027231300a0f319475 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 14 Aug 2026 05:22:45 -0700 Subject: [PATCH] Carry a provider-run tool's result into the next turn (#8713) * Studio: carry a provider-run tool's result into the next turn Hosted and local tools coexist in one turn: Gemini can return a code-execution result while asking for a local web_search, and OpenAI can generate an image before requesting one. The hosted output reached the client as its own _toolEvent frame, but the loop rebuilds the assistant message from the text and tool calls it saw, so that output was absent from the conversation replayed on the follow-up request. The model answered from the local results alone, having lost what it had just produced. Recorded per call id, so a repeated end event cannot duplicate it, and replayed as text rather than as native items: the native shape differs per provider (Gemini codeExecutionResult, an OpenAI image call) while every provider can read its own prior turn's prose. A result that is missing, blank, not a string, or has no call id is ignored rather than trusted. Studio's own tool events are written with a top-level type and never appear as _toolEvent, so local results cannot be replayed twice by this. * Replay a hosted call's operation, strip its frontend sentinels, note its image Three gaps in the first cut, all found in review. The tool_end producers generally omit tool_name, and for Gemini code execution the code that ran is only ever in the tool_start arguments, so a result recorded on its own replayed as an unlabelled value the model could not interpret. Both halves of the call are recorded now and the operation labels its result. A hosted result can carry a full base64 data URI after the __IMAGES__ sentinel, which is there for the card rather than the model. Replaying it verbatim would have put megabytes into the next request for something the model cannot read, so hosted text now goes through the same strip_result_for_model the local results already use. image_generation reports an empty result and carries the picture in image_b64, so requiring non-empty text dropped every generated image, which is exactly the mixed hosted-and-local case this is meant to fix. That it happened is recorded without the bytes, for the same reason the sentinel is stripped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cap hosted output, keep it across a stall, and notice a stdout-less plot Three more from review, all in the replay added on this branch. Gemini code execution that produces a plot and nothing else returns a result that is only the __IMAGES__ sentinel, and does not set image_b64. Stripping it left an empty string with no image marker, so the entry looked empty and the follow-up was told nothing had been produced. The sentinel is noticed before the strip now. Hosted results went into the next prompt uncapped, while local execution has always held what the model sees to 16000 characters. One verbose hosted call could fill the follow-up context on its own, so the replayed copy is held to the same limit. A turn where a hosted tool completed but the model only said what it was about to do takes the stall reprompt, which returns to the provider from above the replay. The reprompted request was therefore told to continue from a search whose output it could no longer see. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the stalled replay in one turn, label the operation, honour the configured cap The stall reprompt appended its hosted block as a new assistant message, so a continued partial was split across two assistant turns. It now merges into the partial the same way the main replay does. The header rendered the whole arguments dict, which carries Gemini's native executableCode part with its thoughtSignature and OpenAI's paired reasoning item with its encrypted content. Those are replay plumbing for the provider, not text a model can read, and they filled the header's budget with truncated base64. The image generation prompt is also only on the end event, so the two halves are merged rather than read from the start alone. The hosted result now goes through the stripper with its own tool name, as local results do, so a fetched page ending in a well formed __FILES__ line keeps it, and the cap is read from the configured local one instead of a copy of its default. * Report a hosted call that ran and printed nothing, and say when its label was cut Gemini reports code that produced no stdout as an empty result, and the code it ran is only ever on the tool_start, so an entry skipped for having no result took the whole execution with it: the next turn saw no trace that anything had run. It now records that the call ended and replays it as (no output), the same answer the local tools and the other hosted paths give. A start with no end is still left out, and a non-string result is still a malformed frame rather than an empty outcome. The label was also cut at 2000 characters with no notice. Anthropic passes the model's whole tool input through as the arguments, so a created file lives there and answers with only Created, and a silent cut reads as a line that simply stops. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Carry the turn's thought signature into the stalled replay Gemini 3 stows a text part's thoughtSignature on the delta, the outbound translator pins it back onto the last text part from assistant.extra_content and nowhere else, and a turn replayed without it is rejected rather than answered. The main replay already carries it; the stall reprompt returns to the provider from above that and did not. * Tighten comments in hosted result replay * Keep an argument the model meant, and read the image sentinel as an envelope Dropping empty values took Anthropic's str_replace deletions with it: new_str of the empty string is the deletion, the schema allows it, and without the key the next turn cannot tell it from a value that was never captured. The guard is not needed anyway, since the provisional empty prompt OpenAI opens an image generation with is overwritten when the two halves merge. The image sentinel is now validated the way the local strippers validate theirs, so a fetched page that merely writes the marker is prose rather than a picture the turn never produced. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Assert on the snippet rather than the URL in the replay tests Checking that a bare URL appears in a string reads to the scanner as URL sanitization by substring, which is the one thing that check exists to flag. The fixture now carries a distinctive snippet and the assertions look for that, so they prove the same thing without the pattern. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- .../core/inference/studio_tool_loop.py | 199 +++ .../tests/test_hosted_result_replay.py | 1062 +++++++++++++++++ 2 files changed, 1261 insertions(+) create mode 100644 studio/backend/tests/test_hosted_result_replay.py diff --git a/studio/backend/core/inference/studio_tool_loop.py b/studio/backend/core/inference/studio_tool_loop.py index 073e5f2147..d84c2c8908 100644 --- a/studio/backend/core/inference/studio_tool_loop.py +++ b/studio/backend/core/inference/studio_tool_loop.py @@ -53,6 +53,7 @@ from dataclasses import dataclass, field from collections.abc import AsyncIterator from typing import Any, Protocol +from core.inference import tools as tools_module from core.inference.chat_template_helpers import append_assistant_turn from core.inference.passthrough_healing import StreamToolCallHealer, heal_gate from core.inference.sse_control_frames import sanitize_provider_sse_line @@ -66,6 +67,7 @@ from core.inference.tool_call_parser import ( from core.inference.tool_loop_controller import ( ToolLoopController, awaiting_approval_status, + strip_result_for_model, ) from core.inference.tool_stream_exec import ( TOOL_HEARTBEAT_INTERVAL_S, @@ -135,6 +137,67 @@ _USAGE_DETAIL_FIELDS = ( _STEP_DONE = object() + +def _truncate_for_model( + text: str, + limit: int | None = None, + *, + joiner: str = "\n", +) -> str: + """Hold a hosted result to the same cap a local result gets. + + Read off ``tools`` rather than copied, so an install that lowers + ``UNSLOTH_TOOL_RESULT_MAX_CHARS`` gets the lower cap here too. + """ + if limit is None: + limit = tools_module._MAX_OUTPUT_CHARS + if len(text) <= limit: + return text + return text[:limit] + f"{joiner}... [truncated, {len(text) - limit} more characters]" + + +# Cap on a call's label. Small next to the result cap: this is the query or the +# code, not the output. +_HOSTED_ARGUMENT_MAX_CHARS = 2000 + + +# Provider plumbing hung off ``arguments`` for the frontend and native-history +# replay, never for the model: Gemini's ``executableCode`` part plus an opaque +# ``thoughtSignature``, OpenAI's paired reasoning item with its multi-kilobyte +# ``encrypted_content``. As prose they are base64 cut off mid-token. +_HOSTED_ARGUMENT_PLUMBING_KEYS = frozenset({"google", "_server_tool"}) + + +def _carries_image_sentinel(result: str) -> bool: + """Whether a hosted result ends in the ``__IMAGES__`` envelope itself. + + Validated rather than matched on sight, the way the local strippers + validate theirs: a fetched page that merely writes the marker is prose, and + reading it as a picture would report an image the turn never made. + """ + _, sep, payload = result.rpartition("\n__IMAGES__:") + if not sep: + return False + try: + images = json.loads(payload) + except (ValueError, RecursionError): + return False + return ( + isinstance(images, list) and bool(images) and all(isinstance(i, str) and i for i in images) + ) + + +def _hosted_arguments_for_model(arguments: Any) -> dict[str, Any]: + """The part of a hosted tool's arguments worth showing the model.""" + if not isinstance(arguments, dict): + return {} + return { + key: value + for key, value in arguments.items() + if key not in _HOSTED_ARGUMENT_PLUMBING_KEYS and not key.startswith("openai_") + } + + # Consecutive turns that asked for a tool but ran none before the loop gives up. _MAX_FRUITLESS_TURNS = 2 @@ -295,6 +358,101 @@ class _Turn: text: list[str] = field(default_factory = list) reasoning_extra: dict[str, Any] | None = None finish_reason: str | None = None + # Results from tools the PROVIDER ran this turn, keyed by call id so a + # repeated end event cannot record the same result twice. + hosted_results: dict[str, dict[str, Any]] = field(default_factory = dict) + + def note_hosted_tool_event(self, event: Any) -> None: + """Record a provider-side tool call carried on ``_toolEvent``. + + These reach the client as their own frames but are not part of the + assistant message this loop replays, so the follow-up request would lose + whatever the provider just produced. Studio's own events carry a + top-level ``type``, so ``_toolEvent`` is unambiguously the provider's. + + Both halves matter: ``tool_end`` generally omits ``tool_name``, and for + Gemini code execution the code that ran is only in the ``tool_start`` + arguments, so a result recorded alone is unlabelled. + """ + if not isinstance(event, dict): + return + call_id = event.get("tool_call_id") + if not isinstance(call_id, str) or not call_id: + return + kind = event.get("type") + if kind not in ("tool_start", "tool_end"): + return + + entry = self.hosted_results.setdefault(call_id, {}) + name = event.get("tool_name") + if isinstance(name, str) and name: + entry["name"] = name + # The operation itself: Gemini's language and code, a search's query. + # Merged across both halves because OpenAI opens an image generation + # before it knows the prompt and only names it on the end event. + arguments = _hosted_arguments_for_model(event.get("arguments")) + if arguments: + merged = dict(entry.get("arguments_obj") or {}) + merged.update(arguments) + entry["arguments_obj"] = merged + # Truncated with the same notice a result gets: Anthropic hands the + # model's whole tool input through, so a file the code wrote lives + # here and nowhere else, and a silent cut reads as the whole thing. + entry["arguments"] = _truncate_for_model( + json.dumps(merged, separators = (",", ":")), + _HOSTED_ARGUMENT_MAX_CHARS, + joiner = " ", + ) + + if kind == "tool_start": + return + result = event.get("result") + if isinstance(result, str): + # The call finished, even if it produced nothing: Gemini reports + # code that printed nothing as an empty string. Recorded apart from + # the result so that stays distinguishable from a stream that died + # after the start. A non-string is a malformed frame, not an outcome. + entry["ended"] = True + if isinstance(result, str) and result.strip(): + if _carries_image_sentinel(result): + # A Gemini plot with no stdout is nothing BUT the sentinel, so + # stripping leaves an empty string and the entry looks empty. + entry["produced_image"] = True + # Same normalisation local results get: the frontend sentinels carry + # a full data URI, and replaying one sends megabytes of base64. The + # tool's name goes with it, as the local path passes it: only the + # sandbox tools emit __FILES__, so a fetched page ending in a well + # formed one keeps that line as the content it is. + stripped = strip_result_for_model(result, entry.get("name")) + if stripped.strip(): + entry["result"] = _truncate_for_model(stripped) + if event.get("image_b64"): + # image_generation reports an empty result and carries the picture + # apart. Record that it happened rather than the bytes. + entry["produced_image"] = True + + def hosted_replay_text(self) -> str: + """The provider-run calls of this turn, as prose for the next request.""" + blocks: list[str] = [] + for entry in self.hosted_results.values(): + result = entry.get("result", "") + produced_image = entry.get("produced_image") + if not result and not produced_image and not entry.get("ended"): + # A start with no outcome says only that something began. + continue + name = entry.get("name") or "tool" + header = f"[{name} result]" + arguments = entry.get("arguments") + if arguments: + header = f"[{name} {arguments}]" + # A call that ended with nothing to show still has to appear, as the + # same "(no output)" local tools report, so the model can tell the + # code ran from it never having run at all. + body = result or ("(produced an image)" if produced_image else "(no output)") + if result and produced_image: + body = f"{result}\n(produced an image)" + blocks.append(f"{header}\n{body}") + return "\n\n".join(blocks) def merge_structured(self, raw_calls: list[Any]) -> None: for raw_call in raw_calls: @@ -698,6 +856,7 @@ async def stream_with_studio_tools( extra = delta.get("extra_content") if isinstance(extra, dict): turn.reasoning_extra = extra + turn.note_hosted_tool_event(payload.get("_toolEvent")) if isinstance(choice.get("finish_reason"), str): turn.finish_reason = choice["finish_reason"] @@ -836,6 +995,33 @@ async def stream_with_studio_tools( ): reprompts += 1 last_reprompt_text = visible_answer + stalled_hosted = turn.hosted_replay_text() + if stalled_hosted: + # A hosted tool did run, the model just did not go on to ask + # for a local one. The replay below never happens on this + # path, so the reprompted request would be told to continue + # from output it can no longer see. + stalled_message: dict[str, Any] = { + "role": "assistant", + "content": ( + f"{visible_answer}\n\n{stalled_hosted}" + if visible_answer + else stalled_hosted + ), + } + if turn.reasoning_extra: + # Gemini 3 stows the text part's thoughtSignature here + # and its translator pins it back on from this field + # alone, so a turn replayed without it is rejected. + stalled_message["extra_content"] = turn.reasoning_extra + append_assistant_turn( + conversation, + stalled_message, + # A resumed partial is the same turn as what the model + # just added, so merge rather than append: appending + # puts a turn boundary mid-sentence. + continue_final_message = run.continue_final_message, + ) _append_user_turn(conversation, reprompt_to_act_message(tool_hint)) continue break @@ -1091,6 +1277,19 @@ async def stream_with_studio_tools( "".join(turn.text), final = True, enabled_tool_names = allowed_tool_names ), } + hosted_text = turn.hosted_replay_text() + if hosted_text: + # A tool the provider ran itself this turn. Its output went to the + # client as its own frame but is not otherwise part of this message, + # so the follow-up would answer from the local results alone. + # Replayed as text, not native items: the shape differs per provider + # (Gemini codeExecutionResult, an OpenAI image call), while every + # provider can read its own prior turn's prose. + assistant_message["content"] = ( + f"{assistant_message['content']}\n\n{hosted_text}" + if assistant_message["content"] + else hosted_text + ) if turn.reasoning_extra: assistant_message["extra_content"] = turn.reasoning_extra if assistant_tool_calls: diff --git a/studio/backend/tests/test_hosted_result_replay.py b/studio/backend/tests/test_hosted_result_replay.py new file mode 100644 index 0000000000..dd58fe535e --- /dev/null +++ b/studio/backend/tests/test_hosted_result_replay.py @@ -0,0 +1,1062 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A tool the provider ran must survive into the next turn's prompt. + +Hosted and local tools coexist in one turn: Gemini can return a code-execution +result while asking for a local ``web_search``. The hosted output reaches the +client as its own ``_toolEvent`` frame, but the loop rebuilds the assistant +message from the text and tool calls it saw, so that output was absent from the +replayed conversation and the model answered from the local results alone. + +Studio's own tool events carry a top-level ``type`` and never appear as +``_toolEvent``, so local results are not replayed twice. +""" + +from __future__ import annotations + +import asyncio +import json +import threading + +import pytest + +from core.inference import studio_tool_loop as loop_mod +from core.inference.studio_tool_loop import ( + ToolLoopPolicy, + ToolLoopRun, + stream_with_studio_tools, +) + + +_DONE = "data: [DONE]" + +WEB = { + "type": "function", + "function": { + "name": "web_search", + "description": "", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + + +def _hosted_event(payload: dict) -> str: + """A provider-side tool frame, shaped as external_provider emits it.""" + return "data: " + json.dumps( + { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + "_toolEvent": payload, + } + ) + + +def _call_line(call_id: str = "c1") -> str: + return "data: " + json.dumps( + { + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": { + "name": "web_search", + "arguments": '{"query": "x"}', + }, + } + ] + }, + } + ] + } + ) + + +def _text(content: str) -> str: + return "data: " + json.dumps({"choices": [{"index": 0, "delta": {"content": content}}]}) + + +def _finish(reason: str = "tool_calls") -> str: + return "data: " + json.dumps({"choices": [{"index": 0, "delta": {}, "finish_reason": reason}]}) + + +class FakeTransport: + heals_text_tool_calls = False + # The OAI-compat transport sanitizes at ingress, so the loop must not strip + # the provider's own _toolEvent frames. + sanitizes_provider_frames = True + + def __init__( + self, + turns, + *, + max_turns = 20, + ): + self.turns = [list(turn) for turn in turns] + self.requests: list[dict] = [] + self.max_turns = max_turns + + def stream(self, *, messages, tools, tool_choice, cancel_event): + self.requests.append({"messages": [dict(m) for m in messages]}) + assert len(self.requests) <= self.max_turns, "loop never terminated" + lines = self.turns.pop(0) if self.turns else [_DONE] + + async def _gen(): + for line in lines: + yield line + + return _gen() + + +@pytest.fixture +def executed(monkeypatch): + calls: list[str] = [] + + def _execute(name, arguments, **kwargs): + calls.append(name) + return f"LOCAL<{name}>" + + monkeypatch.setattr(loop_mod, "execute_tool", _execute) + monkeypatch.setattr(loop_mod, "build_rag_autoinject", lambda *a, **k: None) + monkeypatch.setattr(loop_mod, "is_high_risk_tool_call", lambda name, args: False) + return calls + + +def _run(transport): + async def _collect(): + out: list[str] = [] + agen = stream_with_studio_tools( + transport, + run = ToolLoopRun( + messages = [{"role": "user", "content": "hi"}], + session_id = "s1", + thread_id = "t1", + ), + policy = ToolLoopPolicy( + tools = [WEB], + max_calls = 25, + timeout = 300, + permission_mode = "off", + confirm_calls = False, + bypass_permissions = False, + rag_scope = None, + ), + cancel_event = threading.Event(), + ) + async for line in agen: + out.append(line) + return out + + return asyncio.run(asyncio.wait_for(_collect(), timeout = 30)) + + +def _replayed(transport) -> str: + """Everything the second request carried, as one searchable blob.""" + assert len(transport.requests) > 1, "the loop never made a follow-up request" + return json.dumps(transport.requests[1]["messages"]) + + +# ── the gap ────────────────────────────────────────────────────────── + + +def test_a_hosted_result_reaches_the_follow_up_request(executed): + """The provider searched, then asked us to run a tool. Both must survive.""" + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": "hosted-1", + "arguments": {}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "Title: Unsloth\nSnippet: gradient checkpointing lands", + } + ), + _text("Let me also compute that."), + _call_line(), + _finish(), + ], + [_text("done"), _finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "gradient checkpointing lands" in replayed, "the hosted result was dropped" + assert "LOCAL" in replayed, "the local result was dropped" + + +def test_the_hosted_result_still_reaches_the_client(executed): + """Capturing it for the replay must not stop it rendering.""" + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "hosted output", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + out = _run(transport) + assert any("hosted output" in line and "_toolEvent" in line for line in out) + + +def test_the_models_own_prose_is_kept_alongside(executed): + transport = FakeTransport( + [ + [ + _text("Searching now."), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "hosted output", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "Searching now." in replayed + assert "hosted output" in replayed + + +# ── what must not be replayed ──────────────────────────────────────── + + +def test_a_repeated_end_event_is_recorded_once(executed): + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "once only", + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "once only", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + assert _replayed(transport).count("once only") == 1 + + +def test_a_start_event_alone_adds_nothing(executed): + """A hosted tool that never reported a result has nothing to replay.""" + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": "hosted-1", + "arguments": {}, + } + ), + _text("hello"), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "hello" in replayed + assert "web_search result" not in replayed + + +@pytest.mark.parametrize("result", [None, 42, {"a": 1}]) +def test_a_malformed_result_is_ignored(executed, result): + """A non-string is a malformed frame rather than an outcome, so it records + nothing. An empty string IS an outcome, covered by + ``test_a_silent_hosted_execution_still_reaches_the_next_turn``. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": result, + } + ), + _text("hello"), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + assert "[web_search result]" not in _replayed(transport) + + +def test_an_event_without_a_call_id_is_ignored(executed): + transport = FakeTransport( + [ + [ + _hosted_event({"type": "tool_end", "tool_name": "web_search", "result": "orphan"}), + _text("hello"), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + assert "orphan" not in _replayed(transport) + + +def test_a_frontend_image_sentinel_is_not_replayed(executed): + """__IMAGES__ carries a full data URI for the card, not for the model. + + Replaying it verbatim puts megabytes of base64 into the next request. Local + results already go through the same stripper. + """ + huge = "data:image/png;base64," + ("A" * 20000) + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "code_execution", + "result": '4\n__IMAGES__:["' + huge + '"]', + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "4" in replayed + assert "__IMAGES__" not in replayed + assert "AAAA" not in replayed + + +def test_the_start_events_operation_labels_the_result(executed): + """tool_end generally omits tool_name, and for Gemini code execution the code + that ran is only in the start event, so a result recorded alone replays as an + unlabelled value. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "hosted-1", + "arguments": {"language": "python", "code": "print(2 + 2)"}, + } + ), + _hosted_event( + {"type": "tool_end", "tool_call_id": "hosted-1", "result": "4"}, + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "code_execution" in replayed, "the tool name came only from the start event" + assert "print(2 + 2)" in replayed, "the operation that produced the result was lost" + assert "4" in replayed + + +def test_a_generated_image_is_noted_without_its_bytes(executed): + """image_generation reports an empty result and carries the picture apart. + + Requiring non-empty text dropped it entirely, and replaying the base64 is + the sentinel mistake again, so record only that it happened. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": "hosted-1", + "arguments": {}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "result": "", + "image_b64": "B" * 5000, + "image_mime": "image/png", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "image_generation" in replayed + assert "produced an image" in replayed + assert "BBBB" not in replayed, "the base64 must not reach the next request" + + +def test_a_turn_with_no_hosted_tool_replays_exactly_as_before(executed): + """The common case must be untouched by any of this.""" + transport = FakeTransport( + [[_text("plain answer"), _call_line(), _finish()], [_finish("stop")], [_DONE]] + ) + _run(transport) + replayed = _replayed(transport) + assert "plain answer" in replayed + assert "result]" not in replayed + + +# ── image-only, oversized and stalled turns ────────────────────────── + + +def test_a_plot_with_no_stdout_is_still_reported(executed): + """Gemini code execution can return nothing but the image sentinel. + + Stripping it leaves an empty string and image_b64 is unset on that path, so + unnoticed the entry looks empty and the follow-up is told nothing was made. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "hosted-1", + "arguments": {"code": "plt.plot(x)"}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "result": '\n__IMAGES__:["data:image/png;base64,' + ("C" * 4000) + '"]', + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "produced an image" in replayed + assert "CCCC" not in replayed + + +def test_a_large_hosted_result_is_capped(executed): + """Local execution caps what the model sees; the hosted copy must too.""" + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "code_execution", + "result": "D" * 60000, + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "truncated" in replayed + assert len(replayed) < 40000, f"replayed {len(replayed)} chars" + + +def test_a_stalled_turn_keeps_its_hosted_result(executed): + """The model searched, then only said what it was about to do. + + That takes the stall reprompt, which returns to the provider from above the + main replay, so the request could no longer see the search output. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "Title: Unsloth\nSnippet: gradient checkpointing lands", + } + ), + _text("Let me check that."), + _finish("stop"), + ], + [_text("done"), _finish("stop")], + [_DONE], + ] + ) + _run(transport) + assert len(transport.requests) > 1, "the stall reprompt never happened" + assert "gradient checkpointing lands" in json.dumps(transport.requests[1]["messages"]) + + +def test_a_stalled_continuation_stays_one_assistant_turn(executed): + """The stalled turn's replay must merge into a resumed partial. + + The partial plus what the model just added are one turn. Appending leaves + two assistant messages in a row, splitting a sentence across a turn boundary + and getting rejected by a server that enforces role alternation. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "Title: Unsloth\nSnippet: gradient checkpointing lands", + } + ), + _text(" Let me check that."), + _finish("stop"), + ], + [_text("done"), _finish("stop")], + [_DONE], + ] + ) + + async def _collect(): + agen = stream_with_studio_tools( + transport, + run = ToolLoopRun( + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "The answer is"}, + ], + session_id = "s1", + thread_id = "t1", + continue_final_message = True, + ), + policy = ToolLoopPolicy( + tools = [WEB], + max_calls = 25, + timeout = 300, + permission_mode = "off", + confirm_calls = False, + bypass_permissions = False, + rag_scope = None, + ), + cancel_event = threading.Event(), + ) + async for _ in agen: + pass + + asyncio.run(asyncio.wait_for(_collect(), timeout = 30)) + + messages = transport.requests[1]["messages"] + roles = [m["role"] for m in messages] + assert not any( + roles[i] == "assistant" and roles[i + 1] == "assistant" for i in range(len(roles) - 1) + ), f"the resumed partial was split off its own turn: {roles}" + resumed = [m for m in messages if m["role"] == "assistant"] + assert len(resumed) == 1 + assert resumed[0]["content"].startswith("The answer is Let me check that.") + assert "gradient checkpointing lands" in resumed[0]["content"] + + +# ── the label is the operation, not the transport's plumbing ───────── + + +def test_gemini_code_execution_replays_the_code_not_its_thought_signature(executed): + """Gemini stows the native part on the same `arguments` the header renders. + + ``arguments.google.native_part`` replays Gemini's required history shape and + carries an opaque ``thoughtSignature``. As prose that is a kilobyte of base64 + cut mid-token, pushing the header to its cap on every hosted execution. + """ + signature = "SIG" + ("X" * 3000) + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "code_a", + "arguments": { + "kind": "code_execution", + "language": "python", + "code": "print(1 + 1)", + "_server_tool": True, + "google": { + "native_part": { + "parts": [ + { + "executableCode": { + "id": "code_a", + "language": "PYTHON", + "code": "print(1 + 1)", + }, + "thoughtSignature": signature, + } + ] + } + }, + }, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "code_a", + "result": "2\n", + "google": {"native_part": {"parts": [{"codeExecutionResult": {}}]}}, + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "print(1 + 1)" in replayed + assert "2" in replayed + assert "XXXX" not in replayed, "the thought signature reached the model" + assert "native_part" not in replayed + assert "_server_tool" not in replayed + + +def test_openai_image_generation_replays_the_prompt_it_actually_used(executed): + """The prompt is only on the end event; the start opens with an empty one. + + OpenAI emits ``image_generation_call`` before it knows the prompt, so reading + the start alone replayed ``"prompt": ""``. Those arguments also carry the + paired reasoning item, multi-kilobyte on a zero-data-retention org. + """ + plumbing = { + "openai_image_generation_call_id": "ig_abc", + "openai_response_id": "resp_123", + "openai_reasoning_item": { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "E" * 4000, + }, + } + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": "ig_abc", + "arguments": {"kind": "image", "prompt": "", **plumbing}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "ig_abc", + "result": "", + "arguments": { + "kind": "image", + "prompt": "A photorealistic ginger cat", + **plumbing, + }, + "image_b64": "B" * 5000, + "image_mime": "image/png", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "A photorealistic ginger cat" in replayed, "the prompt was dropped" + assert "produced an image" in replayed + assert "EEEE" not in replayed, "the encrypted reasoning reached the model" + assert "BBBB" not in replayed + + +# ── the cap and the envelope are the local ones, not copies ────────── + + +def test_the_hosted_cap_follows_the_configured_local_one(executed, monkeypatch): + """An install that lowers the local cap lowers the hosted one too. + + ``UNSLOTH_TOOL_RESULT_MAX_CHARS`` shrinks what a result may occupy on a + smaller context; a hosted copy held to its own hard-coded 16k would ignore + that on exactly those installs. + """ + monkeypatch.setattr(loop_mod.tools_module, "_MAX_OUTPUT_CHARS", 500) + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "code_execution", + "result": "D" * 4000, + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "truncated" in replayed + assert "D" * 600 not in replayed, "the configured cap was ignored" + + +def test_a_hosted_page_keeps_a_files_line_of_its_own(executed): + """Only the sandbox tools emit the ``__FILES__`` envelope. + + A fetched page ending in a well formed one is content, so the stripper is + given the tool's name, as the local path does. Otherwise the tail of the + document is dropped before the follow-up turn sees it. + """ + page = 'How the envelope looks\n__FILES__:[{"name": "plot.png", "size": 12}]' + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "web_fetch", + "tool_call_id": "hosted-1", + "arguments": {"url": "https://unsloth.ai/docs"}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "result": page, + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "How the envelope looks" in replayed + assert "plot.png" in replayed, "the page's own __FILES__ line was stripped" + + +# ── a call that ran is a call the next turn hears about ────────────── + + +def test_a_silent_hosted_execution_still_reaches_the_next_turn(executed): + """Gemini reports code that printed nothing as an empty result. + + ``codeExecutionResult.output`` is "" for a run that only wrote a file, and + the code is carried on the tool_start alone, never as assistant text, so + skipping an empty result drops the whole execution. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "code_a", + "arguments": {"language": "python", "code": "df.to_parquet('s.pq')"}, + } + ), + _hosted_event({"type": "tool_end", "tool_call_id": "code_a", "result": ""}), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "df.to_parquet" in replayed, "the execution vanished from the replay" + assert "(no output)" in replayed + + +def test_a_start_with_no_end_is_still_left_out(executed): + """The other half of the same rule: a call that never finished says nothing. + + A stream cut between the halves leaves a start alone, and reporting that + would tell the next turn an execution completed that never did. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "code_a", + "arguments": {"language": "python", "code": "df.to_parquet('s.pq')"}, + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "df.to_parquet" not in replayed + assert "no output" not in replayed + + +def test_a_long_hosted_argument_says_it_was_cut(executed): + """Anthropic passes the model's whole tool input through as arguments. + + A ``create`` carries the entire file body there and answers only "Created", + so the label is the sole record of what was written and a silent cut reads + as the whole of it. + """ + body = "# line\n" * 600 + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "code_a", + "arguments": { + "kind": "text_editor", + "command": "create", + "path": "/tmp/a.py", + "file_text": body, + }, + } + ), + _hosted_event({"type": "tool_end", "tool_call_id": "code_a", "result": "Created"}), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "truncated" in replayed, "the label was cut with no notice" + assert "Created" in replayed + + +def test_a_stalled_turn_keeps_its_thought_signature(executed): + """Gemini 3 will not take its own turn back without the signature. + + The outbound translator pins the ``thoughtSignature`` back on from + ``assistant.extra_content`` and nowhere else. The stall reprompt returns to + the provider from above the main replay, so it has to carry it too. + """ + signature = "SIGNATURE-abc123" + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "tool_name": "web_search", + "result": "Title: Unsloth\nSnippet: gradient checkpointing lands", + } + ), + "data: " + + json.dumps( + { + "choices": [ + { + "index": 0, + "delta": { + "content": "Let me check that.", + "extra_content": {"google": {"thought_signature": signature}}, + }, + } + ] + } + ), + _finish("stop"), + ], + [_text("done"), _finish("stop")], + [_DONE], + ] + ) + _run(transport) + reprompted = transport.requests[1]["messages"] + stalled = [ + m + for m in reprompted + if m["role"] == "assistant" and "gradient checkpointing lands" in str(m.get("content")) + ] + assert stalled, "the hosted result never reached the reprompt" + assert stalled[0].get("extra_content") == {"google": {"thought_signature": signature}} + + +def test_an_empty_argument_the_model_meant_survives(executed): + """Anthropic's text editor deletes by replacing with the empty string. + + ``str_replace`` with ``new_str: ""`` is an intentional deletion, and the + schema allows it, so dropping the key leaves the next turn unable to tell a + deletion from a replacement whose value was never captured. The provisional + empty prompt this once guarded against is overwritten by the end event + anyway, since the two halves merge in order. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "code_execution", + "tool_call_id": "edit_1", + "arguments": { + "kind": "text_editor", + "command": "str_replace", + "old_str": "debug = True", + "new_str": "", + }, + } + ), + _hosted_event({"type": "tool_end", "tool_call_id": "edit_1", "result": "Edited"}), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = "".join( + str(m.get("content")) for m in transport.requests[1]["messages"] if m["role"] == "assistant" + ) + assert '"new_str":""' in replayed, "the deletion looked like a missing value" + + +def test_a_page_that_writes_the_image_marker_is_not_an_image(executed): + """The sentinel is an envelope, not any line mentioning the marker. + + A fetched page documenting the output protocol contains the literal text. + Reading that as a picture reports an image the turn never produced. + """ + transport = FakeTransport( + [ + [ + _hosted_event( + { + "type": "tool_start", + "tool_name": "web_fetch", + "tool_call_id": "hosted-1", + "arguments": {"url": "https://unsloth.ai/docs/protocol"}, + } + ), + _hosted_event( + { + "type": "tool_end", + "tool_call_id": "hosted-1", + "result": "The card reads a line beginning __IMAGES__: and renders it.", + } + ), + _call_line(), + _finish(), + ], + [_finish("stop")], + [_DONE], + ] + ) + _run(transport) + replayed = _replayed(transport) + assert "produced an image" not in replayed