From 35d970c4a50ead559d90d200023184c440d42366 Mon Sep 17 00:00:00 2001
From: huanghuoguoguo <60681390+huanghuoguoguo@users.noreply.github.com>
Date: Mon, 29 Jun 2026 23:15:51 +0800
Subject: [PATCH] test(provider): support fake token counting

---
 tests/unit_tests/provider/conftest.py         | 52 +++++++++++++++++
 .../provider/test_fake_requester.py           | 57 +++++++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 tests/unit_tests/provider/test_fake_requester.py

diff --git a/tests/unit_tests/provider/conftest.py b/tests/unit_tests/provider/conftest.py
index 13b44fd14..db76aefe5 100644
--- a/tests/unit_tests/provider/conftest.py
+++ b/tests/unit_tests/provider/conftest.py
@@ -7,6 +7,8 @@ without calling real LLM APIs or network requests.
 
 from __future__ import annotations
 
+import json
+
 import pytest
 from unittest.mock import AsyncMock, Mock
 from types import SimpleNamespace
@@ -30,6 +32,25 @@ class FakeProviderAPIRequester(requester.ProviderAPIRequester):
         self._invoke_count = 0
         self._last_messages = None
         self._last_model = None
+        self._last_count_tokens_payload = None
+
+    @staticmethod
+    def _content_to_text(content) -> str:
+        if content is None:
+            return ''
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts: list[str] = []
+            for item in content:
+                if isinstance(item, dict):
+                    text = item.get('text')
+                else:
+                    text = getattr(item, 'text', None)
+                if text:
+                    parts.append(str(text))
+            return ''.join(parts)
+        return str(content)
 
     async def invoke_llm(
         self,
@@ -70,6 +91,37 @@ class FakeProviderAPIRequester(requester.ProviderAPIRequester):
             content=[provider_message.ContentElement(type='text', text='Fake stream chunk')],
         )
 
+    async def count_tokens(
+        self,
+        model: requester.RuntimeLLMModel,
+        messages: list,
+        funcs=None,
+        extra_args={},
+    ) -> int:
+        """Return deterministic token estimates for token-free integration tests."""
+        payload: list[dict] = []
+        for message in messages:
+            payload.append(
+                {
+                    'role': getattr(message, 'role', ''),
+                    'content': self._content_to_text(getattr(message, 'content', None)),
+                    'tool_calls': getattr(message, 'tool_calls', None),
+                }
+            )
+
+        for func in funcs or []:
+            payload.append(
+                {
+                    'name': getattr(func, 'name', ''),
+                    'description': getattr(func, 'description', ''),
+                    'parameters': getattr(func, 'parameters', {}),
+                }
+            )
+
+        self._last_count_tokens_payload = payload
+        text = json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str)
+        return max(1, (len(text) + 3) // 4)
+
     async def invoke_embedding(self, model, input_text: list, extra_args={}):
         """Return fake embedding vectors."""
         return [[0.1, 0.2, 0.3] for _ in input_text]
diff --git a/tests/unit_tests/provider/test_fake_requester.py b/tests/unit_tests/provider/test_fake_requester.py
new file mode 100644
index 000000000..9b0f85005
--- /dev/null
+++ b/tests/unit_tests/provider/test_fake_requester.py
@@ -0,0 +1,57 @@
+"""Tests for provider test doubles used by integration paths."""
+
+from __future__ import annotations
+
+import pytest
+
+from langbot.pkg.entity.persistence import model as persistence_model
+from langbot.pkg.provider.modelmgr import requester
+from langbot_plugin.api.entities.builtin.provider import message as provider_message
+from langbot_plugin.api.entities.builtin.resource import tool as resource_tool
+
+
+@pytest.mark.asyncio
+async def test_fake_requester_counts_messages_and_tools(runtime_provider):
+    """Fake requester should support token-free AgentRunner context budgeting."""
+    runtime_model = requester.RuntimeLLMModel(
+        model_entity=persistence_model.LLMModel(
+            uuid='fake-count-model',
+            name='fake-count-model',
+            provider_uuid=runtime_provider.provider_entity.uuid,
+            abilities=['func_call'],
+            extra_args={},
+        ),
+        provider=runtime_provider,
+    )
+
+    async def _placeholder_func(**kwargs):
+        return kwargs
+
+    messages = [
+        provider_message.Message(role='system', content='You are a test assistant.'),
+        provider_message.Message(
+            role='user',
+            content=[
+                provider_message.ContentElement(type='text', text='hello'),
+                provider_message.ContentElement(type='text', text=' world'),
+            ],
+        ),
+    ]
+    tools = [
+        resource_tool.LLMTool(
+            name='lookup',
+            human_desc='Lookup',
+            description='Lookup a value',
+            parameters={'type': 'object', 'properties': {'query': {'type': 'string'}}},
+            func=_placeholder_func,
+        )
+    ]
+
+    requester_inst = runtime_provider.requester
+    message_tokens = await requester_inst.count_tokens(runtime_model, messages, funcs=[])
+    message_and_tool_tokens = await requester_inst.count_tokens(runtime_model, messages, funcs=tools)
+
+    assert message_tokens > 0
+    assert message_and_tool_tokens > message_tokens
+    assert requester_inst._last_count_tokens_payload[-1]['name'] == 'lookup'
+    assert requester_inst._last_count_tokens_payload[1]['content'] == 'hello world'