tests/test_tools.py

"""Tool calls, on the way out and on the way back.

The fragments a provider streams for one tool call are the part most likely to
be got wrong: the name arrives in one chunk, the arguments a character at a time
across dozens more, and half of them carry no id. All of that is assembled here,
so all of it is asserted here.
"""

from __future__ import annotations

import json

from server.stream import (
    Accumulated, build_payload, normalise_messages, normalise_tools, split_thinking,
    translate,
)


def frames(lines: list[str]) -> tuple[list[tuple[str, str]], Accumulated]:
    acc = Accumulated()
    out: list[tuple[str, str]] = []
    for line in lines:
        for frame in translate(line, acc):
            name = "message"
            data: list[str] = []
            for part in frame.split("\n"):
                if part.startswith("event:"):
                    name = part[6:].strip()
                elif part.startswith("data:"):
                    data.append(part[5:].lstrip(" "))
            out.append((name, "\n".join(data)))
    return out, acc


def chunk(delta: dict, finish: str | None = None) -> str:
    return "data: " + json.dumps({"choices": [{"delta": delta, "finish_reason": finish}]})


class TestAssemblingCalls:
    STREAM = [
        chunk({"tool_calls": [{"index": 0, "id": "call_a", "function": {"name": "canvas_open"}}]}),
        chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"kind":'}}]}),
        chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"line"}'}}]}),
        chunk({}, finish="tool_calls"),
        "data: [DONE]",
    ]

    def test_emits_one_assembled_call(self):
        events, _ = frames(self.STREAM)
        tools = [json.loads(d) for name, d in events if name == "tools"]
        assert len(tools) == 1, "the calls are sent once, not per fragment"
        assert tools[0] == [{"id": "call_a", "name": "canvas_open",
                             "arguments": '{"kind":"line"}'}]

    def test_does_not_repeat_them_at_done(self):
        events, _ = frames(self.STREAM)
        assert [n for n, _ in events].count("tools") == 1

    def test_keeps_parallel_calls_in_index_order(self):
        events, _ = frames([
            chunk({"tool_calls": [{"index": 1, "id": "b", "function": {"name": "second", "arguments": "{}"}}]}),
            chunk({"tool_calls": [{"index": 0, "id": "a", "function": {"name": "first", "arguments": "{}"}}]}),
            "data: [DONE]",
        ])
        calls = json.loads(next(d for n, d in events if n == "tools"))
        assert [c["name"] for c in calls] == ["first", "second"]

    def test_invents_an_id_when_the_provider_omits_one(self):
        # The id only has to be unique within the turn: the model is handed back
        # whatever we send, so ours is as good as theirs.
        events, _ = frames([
            chunk({"tool_calls": [{"index": 0, "function": {"name": "compute", "arguments": "{}"}}]}),
            "data: [DONE]",
        ])
        calls = json.loads(next(d for n, d in events if n == "tools"))
        assert calls[0]["id"]

    def test_drops_a_fragment_that_never_named_a_tool(self):
        events, _ = frames([
            chunk({"tool_calls": [{"index": 0, "function": {"arguments": "{}"}}]}),
            "data: [DONE]",
        ])
        assert not [d for n, d in events if n == "tools"]

    def test_content_still_streams_alongside(self):
        events, acc = frames([
            chunk({"content": "one moment"}),
            chunk({"tool_calls": [{"index": 0, "id": "x", "function": {"name": "canvas_list", "arguments": "{}"}}]}),
            "data: [DONE]",
        ])
        assert acc.content == "one moment"
        assert [n for n, _ in events].count("delta") == 1


class TestNormaliseMessages:
    def test_translates_a_whole_tool_round_trip(self):
        out = normalise_messages([
            {"role": "user", "content": "draw it"},
            {"role": "assistant", "content": "",
             "toolCalls": [{"id": "c1", "name": "canvas_open", "arguments": '{"kind":"bar"}'}]},
            {"role": "tool", "content": '{"ok":true}', "toolCallId": "c1", "name": "canvas_open"},
        ])
        assert out[1]["tool_calls"] == [{
            "id": "c1", "type": "function",
            "function": {"name": "canvas_open", "arguments": '{"kind":"bar"}'},
        }]
        assert out[2] == {"role": "tool", "content": '{"ok":true}',
                          "tool_call_id": "c1", "name": "canvas_open"}

    def test_drops_a_tool_message_with_no_call_to_answer(self):
        # A provider rejects one outright, and it is never useful.
        assert normalise_messages([{"role": "tool", "content": "{}"}]) == []

    def test_drops_unknown_roles_and_non_objects(self):
        out = normalise_messages([{"role": "narrator", "content": "x"}, "hello", None])
        assert out == []

    def test_leaves_ordinary_messages_alone(self):
        assert normalise_messages([{"role": "user", "content": "hi"}]) == [
            {"role": "user", "content": "hi"},
        ]


class TestNormaliseTools:
    GOOD = {"type": "function", "function": {"name": "compute", "description": "d", "parameters": {"type": "object"}}}

    def test_passes_a_well_formed_tool(self):
        assert normalise_tools([self.GOOD])[0]["function"]["name"] == "compute"

    def test_rejects_anything_without_a_name(self):
        assert normalise_tools([{"type": "function", "function": {}}]) == []
        assert normalise_tools(["compute"]) == []
        assert normalise_tools(None) == []

    def test_caps_how_many_it_will_forward(self):
        assert len(normalise_tools([self.GOOD] * 100)) == 32

    def test_payload_carries_tools_and_an_auto_choice(self):
        p = build_payload(
            {"messages": [{"role": "user", "content": "hi"}], "tools": [self.GOOD]},
            "zai-org/GLM-5.3", reasoning=True,
        )
        assert p["tool_choice"] == "auto"
        assert len(p["tools"]) == 1

    def test_payload_omits_the_field_entirely_when_there_are_none(self):
        p = build_payload({"messages": [{"role": "user", "content": "hi"}]},
                          "zai-org/GLM-5.3", reasoning=True)
        assert "tools" not in p and "tool_choice" not in p


class TestInlineThinking:
    """Providers that drop the reasoning channel and inline <think> instead.

    The reader otherwise gets the model's scratchpad and a stray closing tag in
    the middle of the answer, which is what sent us looking for this.
    """

    def route(self, chunks: list[str]) -> tuple[str, str, Accumulated]:
        acc = Accumulated()
        answer, thinking = [], []
        for c in chunks:
            a, t = split_thinking(c, acc)
            answer.append(a)
            thinking.append(t)
        return "".join(answer), "".join(thinking), acc

    def test_splits_the_two_channels(self):
        answer, thinking, _ = self.route(["<think>let me see</think>The answer."])
        assert answer == "The answer."
        assert thinking == "let me see"

    def test_handles_a_tag_split_across_chunks(self):
        # The common case: tags arrive a few characters at a time.
        answer, thinking, acc = self.route(["<thi", "nk>hmm</thin", "k>done"])
        assert answer == "done"
        assert thinking == "hmm"
        assert acc.held == ""

    def test_holds_back_a_fragment_that_might_still_be_a_tag(self):
        answer, _, acc = self.route(["hello <"])
        assert answer == "hello "
        assert acc.held == "<"

    def test_releases_a_fragment_that_turned_out_to_be_text(self):
        answer, _, _ = self.route(["a <", "b c"])
        assert answer == "a <b c"

    def test_leaves_ordinary_content_alone(self):
        answer, thinking, _ = self.route(["just an answer"])
        assert answer == "just an answer"
        assert thinking == ""

    def test_an_unclosed_think_keeps_routing_to_reasoning(self):
        answer, thinking, acc = self.route(["<think>still going"])
        assert answer == ""
        assert thinking == "still going"
        assert acc.thinking is True

    def test_a_dangling_fragment_is_flushed_as_text_at_done(self):
        acc = Accumulated()
        list(translate('data: {"choices":[{"delta":{"content":"trailing <"}}]}', acc))
        frames = [f for f in translate("data: [DONE]", acc)]
        assert acc.content == "trailing <"
        assert any('"content": "<"' in f or '"content":"<"' in f for f in frames)

    def test_reaches_the_client_on_the_right_channels(self):
        events, acc = frames([
            chunk({"content": "<think>weighing it</think>Here."}),
            "data: [DONE]",
        ])
        deltas = [json.loads(d) for name, d in events if name == "delta"]
        assert deltas[0] == {"content": "Here.", "reasoning": "weighing it"}
        assert acc.content == "Here."
        assert acc.reasoning == "weighing it"


class TestArgumentsGoBackAsJson:
    """The router 400s the whole request over one call's malformed arguments.

    Which means a model that wrote a JavaScript object literal kills the
    conversation on the *next* turn, when its own call is echoed back — and the
    error names neither the call nor the turn it came from.
    """

    def send(self, arguments):
        out = normalise_messages([
            {"role": "assistant", "content": "",
             "toolCalls": [{"id": "c1", "name": "canvas_open", "arguments": arguments}]},
            {"role": "tool", "content": "{}", "toolCallId": "c1"},
        ])
        return out[0]["tool_calls"][0]["function"]["arguments"]

    def test_passes_valid_json_through_untouched(self):
        assert self.send('{"kind":"line"}') == '{"kind":"line"}'

    def test_replaces_a_javascript_object_literal(self):
        # Unquoted keys are the single most common thing a model writes.
        assert json.loads(self.send("{kind: 'line'}")) == {}

    def test_replaces_anything_that_is_not_an_object(self):
        # Valid JSON, still not an argument list.
        assert self.send("[1, 2]") == "{}"
        assert self.send('"line"') == "{}"

    def test_survives_a_missing_or_odd_arguments_field(self):
        assert self.send(None) == "{}"
        assert self.send("") == "{}"
        assert json.loads(self.send({"kind": "line"})) == {"kind": "line"}

    def test_every_call_in_a_request_is_guaranteed(self):
        out = normalise_messages([{
            "role": "assistant", "content": "",
            "toolCalls": [
                {"id": "a", "name": "one", "arguments": "{broken"},
                {"id": "b", "name": "two", "arguments": '{"ok":true}'},
            ],
        }])
        for call in out[0]["tool_calls"]:
            json.loads(call["function"]["arguments"])