Enhance AIClient and MCPServer to support tool registration with source tracking. Added logging for tool calls and improved error handling. Introduced methods for embedding token limit extraction and budget application in OpenAIModel. Added tests for MCP tool registration and execution.

tests/test_openai_model_compat.py

@@ -43,6 +43,18 @@ class _FakeEmbeddings:
         return SimpleNamespace(data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3])])
 
 
+class _LengthLimitedEmbeddings:
+    def __init__(self):
+        self.inputs = []
+
+    async def create(self, **kwargs):
+        text = kwargs.get("input", "")
+        self.inputs.append(text)
+        if len(text) > 512:
+            raise RuntimeError("input must be less than 512 tokens")
+        return SimpleNamespace(data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3])])
+
+
 class _ModernCompletions:
     def __init__(self):
         self.last_params = None
@@ -197,6 +209,13 @@ class _RuntimeRejectToolsAsyncOpenAI:
         self.embeddings = _FakeEmbeddings()
 
 
+class _LengthLimitedEmbedAsyncOpenAI:
+    def __init__(self, **kwargs):
+        self.completions = _ModernCompletions()
+        self.chat = SimpleNamespace(completions=self.completions)
+        self.embeddings = _LengthLimitedEmbeddings()
+
+
 def test_openai_model_uses_tools_when_supported(monkeypatch):
     monkeypatch.setattr(openai_model_module, "AsyncOpenAI", _ModernAsyncOpenAI)
 
@@ -276,3 +295,28 @@ def test_openai_model_retries_with_functions_when_tools_rejected(monkeypatch):
     assert model._supports_tools is False
     assert result.tool_calls is not None
     assert result.tool_calls[0]["function"]["name"] == "demo_tool"
+
+
+def test_openai_model_learns_embedding_limit_and_pretruncates(monkeypatch):
+    monkeypatch.setattr(
+        openai_model_module, "AsyncOpenAI", _LengthLimitedEmbedAsyncOpenAI
+    )
+
+    model = OpenAIModel(_model_config())
+    long_text = "你" * 600
+
+    first_embedding = asyncio.run(model.embed(long_text))
+    assert first_embedding == [0.1, 0.2, 0.3]
+    assert model._embedding_token_limit == 512
+
+    inputs_after_first = list(model.client.embeddings.inputs)
+    assert len(inputs_after_first) == 2
+    assert len(inputs_after_first[0]) == 600
+    assert len(inputs_after_first[1]) < len(inputs_after_first[0])
+
+    second_embedding = asyncio.run(model.embed(long_text))
+    assert second_embedding == [0.1, 0.2, 0.3]
+
+    inputs_after_second = list(model.client.embeddings.inputs)
+    assert len(inputs_after_second) == 3
+    assert len(inputs_after_second[-1]) < 512