1 files changed, 126 insertions, 0 deletions
diff --git a/tests/test_llm.py b/tests/test_llm.py
index b4e68ad..edd9d61 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import base64
+from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from ollama import Message
@@ -105,3 +107,127 @@ def test_token_counts_extracts_from_object() -> None:
 def test_token_counts_returns_none_when_missing() -> None:
     assert _token_counts({}) is None
     assert _token_counts(None) is None
+
+
+def test_chat_with_tools_disables_tools_for_final_answer_after_limit(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from ollama_web import llm
+
+    executed: list[str] = []
+
+    def lookup(query: str) -> str:
+        """Return a test lookup result."""
+        executed.append(query)
+        return "lookup result"
+
+    registry = ToolRegistry()
+    registry.register(lookup)
+
+    class FakeClient:
+        def __init__(self) -> None:
+            self.calls: list[dict[str, Any]] = []
+
+        def chat(self, **kwargs: Any) -> dict[str, Any]:
+            self.calls.append(kwargs)
+            if len(self.calls) == 1:
+                return {
+                    "message": {
+                        "role": "assistant",
+                        "content": "",
+                        "tool_calls": [
+                            {"function": {"name": "lookup", "arguments": {"query": "test"}}}
+                        ],
+                    }
+                }
+            return {"message": {"role": "assistant", "content": "final answer"}}
+
+    client = FakeClient()
+    monkeypatch.setattr(llm, "get_client", lambda _host=None: client)
+    monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1)
+
+    events = list(
+        llm.chat_with_tools(
+            [Message(role="user", content="research this")],
+            model="test-model",
+            registry=registry,
+            capabilities={"tools"},
+            language="en",
+        )
+    )
+
+    assert executed == ["test"]
+    assert len(client.calls) == 2
+    assert client.calls[0]["tools"]
+    assert client.calls[1]["tools"] == []
+    assert any(event.get("type") == "status" for event in events)
+    assert any(event == {"type": "delta", "content": "final answer"} for event in events)
+    assert events[-1] == {"type": "done"}
+
+
+def test_stream_chat_with_tools_disables_mcp_tools_for_final_answer_after_limit(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from ollama_web import llm
+
+    tool_name = "mcp__test__lookup"
+    executed: list[tuple[str, dict[str, object] | str]] = []
+    registry = ToolRegistry()
+
+    def executor(full_name: str, arguments: dict[str, object] | str = "") -> str:
+        executed.append((full_name, arguments))
+        return "mcp lookup result"
+
+    registry.register_mcp_tool(
+        tool_name,
+        {"type": "function", "function": {"name": tool_name}},
+        executor,
+    )
+
+    class FakeClient:
+        def __init__(self) -> None:
+            self.calls: list[dict[str, Any]] = []
+
+        def chat(self, **kwargs: Any) -> list[SimpleNamespace]:
+            self.calls.append(kwargs)
+            if len(self.calls) == 1:
+                return [
+                    SimpleNamespace(
+                        message={
+                            "content": "",
+                            "tool_calls": [
+                                {
+                                    "function": {
+                                        "name": tool_name,
+                                        "arguments": {"query": "test"},
+                                    }
+                                }
+                            ],
+                        }
+                    )
+                ]
+            return [SimpleNamespace(message={"content": "final answer"})]
+
+    client = FakeClient()
+    monkeypatch.setattr(llm, "get_client", lambda _host=None: client)
+    monkeypatch.setattr(llm, "MAX_TOOL_ROUNDS", 1)
+
+    events = list(
+        llm.stream_chat_with_tools(
+            [Message(role="user", content="research this")],
+            model="test-model",
+            registry=registry,
+            capabilities={"tools"},
+            language="en",
+        )
+    )
+
+    assert executed == [(tool_name, {"query": "test"})]
+    assert len(client.calls) == 2
+    assert client.calls[0]["tools"]
+    assert client.calls[1]["tools"] == []
+    assert client.calls[0]["stream"] is True
+    assert client.calls[1]["stream"] is True
+    assert any(event.get("type") == "status" for event in events)
+    assert any(event == {"type": "delta", "content": "final answer"} for event in events)
+    assert events[-1] == {"type": "done"}