kiln_ai.sandbox.test_code_tool_execution

   1"""Tests for the code-tool execution engine.
   2
   3Child/protocol tests spawn real child processes. Parent-side tests use
   4mock tool doubles. Shaped after the existing ``test_sandbox_worker.py`` suite.
   5"""
   6
   7import asyncio
   8import json
   9import textwrap
  10from unittest.mock import patch
  11
  12import pytest
  13
  14from kiln_ai.datamodel.code_tool import CodeTool
  15from kiln_ai.datamodel.project import Project
  16from kiln_ai.sandbox.spawn import _spawn_lock
  17from kiln_ai.tools.base_tool import (
  18    KilnToolInterface,
  19    ToolCallDefinition,
  20    ToolCallResult,
  21)
  22from kiln_ai.tools.code_tool import (
  23    PythonCodeTool,
  24    ToolCallLogEntry,
  25)
  26from kiln_ai.tools.sandbox_bridge import (
  27    CODE_SANDBOX_MAX_CONCURRENCY,
  28    NestedToolServer,
  29    _depth,
  30)
  31
  32# ---------------------------------------------------------------------------
  33# Helpers
  34# ---------------------------------------------------------------------------
  35
  36VALID_SCHEMA = {
  37    "type": "object",
  38    "properties": {"x": {"type": "string"}},
  39}
  40EMPTY_SCHEMA = {"type": "object", "properties": {}}
  41
  42
  43def _make_code_tool(code: str, **overrides) -> CodeTool:
  44    defaults = {
  45        "name": "Test Tool",
  46        "tool_function_name": "test_tool",
  47        "tool_description": "A test tool",
  48        "parameters_schema": VALID_SCHEMA,
  49        "code": code,
  50        "timeout_seconds": 10,
  51    }
  52    defaults.update(overrides)
  53    return CodeTool(**defaults)
  54
  55
  56def _make_project(tmp_path) -> Project:
  57    p = Project(name="test_project", path=tmp_path / "project")
  58    p.save_to_file()
  59    return p
  60
  61
  62def _make_python_code_tool(
  63    tmp_path,
  64    code: str,
  65    tool_allowlist=None,
  66    tool_call_recorder=None,
  67    **overrides,
  68) -> PythonCodeTool:
  69    project = _make_project(tmp_path)
  70    ct = _make_code_tool(
  71        code,
  72        tool_allowlist=tool_allowlist or [],
  73        **overrides,
  74    )
  75    ct.parent = project
  76    return PythonCodeTool(
  77        ct,
  78        project,
  79        tool_call_recorder=tool_call_recorder,
  80    )
  81
  82
  83class FakeTool(KilnToolInterface):
  84    """Minimal tool double for testing nested calls.
  85
  86    IMPORTANT: ``fn_name`` intentionally DIFFERS from the tool_id slug
  87    (e.g. ``fn_name="fake_add"`` for ``tool_id="kiln_tool::add_numbers"``).
  88    This ensures tests catch name-derivation bugs where the dispatch map
  89    would use the slug instead of ``tool.name()``.
  90    """
  91
  92    def __init__(
  93        self,
  94        tool_id: str,
  95        fn_name: str,
  96        fn_desc: str = "fake",
  97        params: dict | None = None,
  98        result: ToolCallResult | None = None,
  99        delay: float = 0,
 100    ):
 101        self._id = tool_id
 102        self._name = fn_name
 103        self._desc = fn_desc
 104        self._params = params or EMPTY_SCHEMA
 105        self._result = result or ToolCallResult(output="ok")
 106        self._delay = delay
 107
 108    async def id(self):
 109        return self._id
 110
 111    async def name(self):
 112        return self._name
 113
 114    async def description(self):
 115        return self._desc
 116
 117    async def toolcall_definition(self) -> ToolCallDefinition:
 118        return {
 119            "type": "function",
 120            "function": {
 121                "name": self._name,
 122                "description": self._desc,
 123                "parameters": self._params,
 124            },
 125        }
 126
 127    async def run(self, context=None, **kwargs) -> ToolCallResult:
 128        if self._delay > 0:
 129            await asyncio.sleep(self._delay)
 130        return self._result
 131
 132
 133# ---------------------------------------------------------------------------
 134# Child / protocol tests (real spawns)
 135# ---------------------------------------------------------------------------
 136
 137
 138class TestChildSyncRun:
 139    @pytest.mark.asyncio
 140    async def test_sync_run_returns_string(self, tmp_path):
 141        tool = _make_python_code_tool(
 142            tmp_path,
 143            'def run(x):\n    return "hello " + x\n',
 144        )
 145        result = await tool.run(None, x="world")
 146        assert not result.is_error
 147        assert result.output == "hello world"
 148
 149    @pytest.mark.asyncio
 150    async def test_sync_run_returns_dict(self, tmp_path):
 151        tool = _make_python_code_tool(
 152            tmp_path,
 153            'def run(x):\n    return {"value": x}\n',
 154        )
 155        result = await tool.run(None, x="test")
 156        assert not result.is_error
 157        assert json.loads(result.output) == {"value": "test"}
 158
 159    @pytest.mark.asyncio
 160    async def test_sync_run_returns_none(self, tmp_path):
 161        tool = _make_python_code_tool(
 162            tmp_path,
 163            "def run(x):\n    pass\n",
 164        )
 165        result = await tool.run(None, x="test")
 166        assert not result.is_error
 167        assert result.output == "null"
 168
 169
 170class TestChildAsyncRun:
 171    @pytest.mark.asyncio
 172    async def test_async_run_returns_string(self, tmp_path):
 173        tool = _make_python_code_tool(
 174            tmp_path,
 175            textwrap.dedent("""\
 176                import asyncio
 177                async def run(x):
 178                    async def greet(name):
 179                        return "hi " + name
 180                    results = await asyncio.gather(greet(x), greet(x + "!"))
 181                    return " ".join(results)
 182            """),
 183        )
 184        result = await tool.run(None, x="a")
 185        assert not result.is_error
 186        assert result.output == "hi a hi a!"
 187
 188    @pytest.mark.asyncio
 189    async def test_asyncio_run_inside_async_errors(self, tmp_path):
 190        tool = _make_python_code_tool(
 191            tmp_path,
 192            textwrap.dedent("""\
 193                import asyncio
 194                async def helper():
 195                    return 1
 196                async def run(x):
 197                    return asyncio.run(helper())
 198            """),
 199        )
 200        result = await tool.run(None, x="test")
 201        assert result.is_error
 202        assert (
 203            "cannot be called from a running event loop" in result.output.lower()
 204            or "cannot" in result.output.lower()
 205        )
 206
 207
 208class TestReturnSerialization:
 209    @pytest.mark.asyncio
 210    @pytest.mark.parametrize(
 211        "code,expected",
 212        [
 213            ('def run(x):\n    return "raw"\n', "raw"),
 214            ("def run(x):\n    return 42\n", "42"),
 215            ("def run(x):\n    return 3.14\n", "3.14"),
 216            ("def run(x):\n    return True\n", "true"),
 217            ("def run(x):\n    return False\n", "false"),
 218            ("def run(x):\n    return None\n", "null"),
 219            ("def run(x):\n    return [1, 2]\n", "[1, 2]"),
 220            ('def run(x):\n    return {"k": "v"}\n', '{"k": "v"}'),
 221        ],
 222        ids=[
 223            "str",
 224            "int",
 225            "float",
 226            "bool_true",
 227            "bool_false",
 228            "none",
 229            "list",
 230            "dict",
 231        ],
 232    )
 233    async def test_serialization(self, tmp_path, code, expected):
 234        tool = _make_python_code_tool(tmp_path, code)
 235        result = await tool.run(None, x="test")
 236        assert not result.is_error
 237        assert result.output == expected
 238
 239    @pytest.mark.asyncio
 240    async def test_non_serializable_type_errors(self, tmp_path):
 241        tool = _make_python_code_tool(
 242            tmp_path,
 243            "def run(x):\n    return object()\n",
 244        )
 245        result = await tool.run(None, x="test")
 246        assert result.is_error
 247        assert "must return str or JSON-serializable" in result.output
 248
 249    @pytest.mark.asyncio
 250    async def test_non_serializable_nested_value_errors(self, tmp_path):
 251        tool = _make_python_code_tool(
 252            tmp_path,
 253            "def run(x):\n    return {'fn': lambda: None}\n",
 254        )
 255        result = await tool.run(None, x="test")
 256        assert result.is_error
 257        assert "non-JSON-serializable" in result.output
 258
 259    @pytest.mark.asyncio
 260    async def test_string_passthrough_no_parsing(self, tmp_path):
 261        """JSON-shaped string returned by run() comes back as-is, not parsed."""
 262        tool = _make_python_code_tool(
 263            tmp_path,
 264            'def run(x):\n    return \'{"key": "value"}\'\n',
 265        )
 266        result = await tool.run(None, x="test")
 267        assert not result.is_error
 268        assert result.output == '{"key": "value"}'
 269
 270
 271class TestStdoutStderr:
 272    @pytest.mark.asyncio
 273    async def test_stdout_captured(self, tmp_path):
 274        project = _make_project(tmp_path)
 275        ct = _make_code_tool(
 276            'import sys\ndef run(x):\n    sys.stdout.write("debug")\n    return "ok"\n',
 277        )
 278        ct.parent = project
 279        pct = PythonCodeTool(ct, project)
 280        outcome = await pct._invoke(None, {"x": "test"})
 281        assert outcome.ok == "ok"
 282        assert "debug" in outcome.stdout
 283
 284    @pytest.mark.asyncio
 285    async def test_stdout_truncation(self, tmp_path):
 286        project = _make_project(tmp_path)
 287        ct = _make_code_tool(
 288            'import sys\ndef run(x):\n    sys.stdout.write("A" * 100000)\n    return "ok"\n',
 289        )
 290        ct.parent = project
 291        pct = PythonCodeTool(ct, project)
 292        outcome = await pct._invoke(None, {"x": "test"})
 293        assert outcome.ok == "ok"
 294        assert len(outcome.stdout) <= 64 * 1024 + 50
 295        assert "truncated" in outcome.stdout
 296
 297
 298class TestTraceback:
 299    @pytest.mark.asyncio
 300    async def test_traceback_shows_code_tool_lines(self, tmp_path):
 301        tool = _make_python_code_tool(
 302            tmp_path,
 303            textwrap.dedent("""\
 304                def helper():
 305                    raise ValueError("kaboom")
 306                def run(x):
 307                    helper()
 308            """),
 309        )
 310        result = await tool.run(None, x="test")
 311        assert result.is_error
 312        assert "kaboom" in result.output
 313        assert "<code_tool>" in result.output
 314        assert "worker.py" not in result.output
 315
 316
 317class TestMissingRun:
 318    @pytest.mark.asyncio
 319    async def test_missing_run_defense(self, tmp_path):
 320        """Even if save-time validation is bypassed, child handles missing run()."""
 321        project = _make_project(tmp_path)
 322        ct = CodeTool.__new__(CodeTool)
 323        object.__setattr__(
 324            ct,
 325            "__dict__",
 326            {
 327                "name": "bad",
 328                "tool_function_name": "bad",
 329                "tool_description": "bad",
 330                "parameters_schema": EMPTY_SCHEMA,
 331                "code": "x = 1\n",
 332                "timeout_seconds": 10,
 333                "tool_allowlist": [],
 334                "description": None,
 335                "is_archived": False,
 336                "id": "test123",
 337                "v": 1,
 338                "created_at": None,
 339                "created_by": None,
 340                "path": None,
 341            },
 342        )
 343        object.__setattr__(ct, "__pydantic_fields_set__", set())
 344        pct = PythonCodeTool(ct, project)
 345        result = await pct.run(None)
 346        assert result.is_error
 347        assert "run" in result.output.lower()
 348
 349
 350class TestImportForms:
 351    @pytest.mark.asyncio
 352    async def test_from_kiln_import_tools(self, tmp_path):
 353        tool = _make_python_code_tool(
 354            tmp_path,
 355            textwrap.dedent("""\
 356                from kiln import tools
 357                def run(x):
 358                    return type(tools).__name__
 359            """),
 360        )
 361        result = await tool.run(None, x="test")
 362        assert not result.is_error
 363
 364    @pytest.mark.asyncio
 365    async def test_import_kiln_tools(self, tmp_path):
 366        tool = _make_python_code_tool(
 367            tmp_path,
 368            textwrap.dedent("""\
 369                import kiln.tools
 370                def run(x):
 371                    return type(kiln.tools).__name__
 372            """),
 373        )
 374        result = await tool.run(None, x="test")
 375        assert not result.is_error
 376
 377    @pytest.mark.asyncio
 378    async def test_from_kiln_tools_import_exception(self, tmp_path):
 379        tool = _make_python_code_tool(
 380            tmp_path,
 381            textwrap.dedent("""\
 382                from kiln.tools import ToolCallError
 383                def run(x):
 384                    return ToolCallError.__name__
 385            """),
 386        )
 387        result = await tool.run(None, x="test")
 388        assert not result.is_error
 389        assert result.output == "ToolCallError"
 390
 391    @pytest.mark.asyncio
 392    async def test_from_kiln_import_async_tools(self, tmp_path):
 393        tool = _make_python_code_tool(
 394            tmp_path,
 395            textwrap.dedent("""\
 396                from kiln import async_tools
 397                def run(x):
 398                    return type(async_tools).__name__
 399            """),
 400        )
 401        result = await tool.run(None, x="test")
 402        assert not result.is_error
 403
 404    @pytest.mark.asyncio
 405    async def test_exception_classes_identical_across_modules(self, tmp_path):
 406        tool = _make_python_code_tool(
 407            tmp_path,
 408            textwrap.dedent("""\
 409                from kiln import tools, async_tools
 410                def run(x):
 411                    same_not_allowed = tools.ToolNotAllowed is async_tools.ToolNotAllowed
 412                    same_timeout = tools.ToolTimeout is async_tools.ToolTimeout
 413                    same_call_error = tools.ToolCallError is async_tools.ToolCallError
 414                    return str(same_not_allowed and same_timeout and same_call_error)
 415            """),
 416        )
 417        result = await tool.run(None, x="test")
 418        assert not result.is_error
 419        assert result.output == "True"
 420
 421
 422class TestJsonUnsafeKwargs:
 423    @pytest.mark.asyncio
 424    async def test_json_unsafe_kwargs_raise_in_frame(self, tmp_path):
 425        """Non-JSON-serializable tool kwargs raise ToolCallError inside child."""
 426        code = textwrap.dedent("""\
 427            from kiln.tools import ToolCallError
 428            from kiln import tools
 429            def run(x):
 430                try:
 431                    tools.some_tool(bad=object())
 432                except ToolCallError as e:
 433                    return f"caught: {e.tool}"
 434                return "no error"
 435        """)
 436        tool = _make_python_code_tool(tmp_path, code)
 437        result = await tool.run(None, x="test")
 438        assert not result.is_error
 439        assert "caught: some_tool" in result.output
 440
 441
 442# ---------------------------------------------------------------------------
 443# Parent-side tests (mock tools)
 444# ---------------------------------------------------------------------------
 445
 446
 447class TestHappyPath:
 448    @pytest.mark.asyncio
 449    async def test_simple_run(self, tmp_path):
 450        tool = _make_python_code_tool(
 451            tmp_path,
 452            'def run(x):\n    return "result_" + x\n',
 453        )
 454        result = await tool.run(None, x="abc")
 455        assert not result.is_error
 456        assert result.output == "result_abc"
 457
 458
 459class TestNestedToolCalls:
 460    @pytest.mark.asyncio
 461    async def test_nested_tool_success(self, tmp_path):
 462        fake = FakeTool(
 463            "kiln_tool::add_numbers",
 464            "fake_add",
 465            params=EMPTY_SCHEMA,
 466            result=ToolCallResult(output="42"),
 467        )
 468        code = textwrap.dedent("""\
 469            from kiln import tools
 470            def run(x):
 471                result = tools.fake_add()
 472                return "got: " + result
 473        """)
 474        tool = _make_python_code_tool(
 475            tmp_path,
 476            code,
 477            tool_allowlist=["kiln_tool::add_numbers"],
 478        )
 479        with patch(
 480            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 481            return_value=fake,
 482        ):
 483            result = await tool.run(None, x="test")
 484        assert not result.is_error
 485        assert result.output == "got: 42"
 486
 487    @pytest.mark.asyncio
 488    async def test_nested_tool_is_error(self, tmp_path):
 489        fake = FakeTool(
 490            "kiln_tool::add_numbers",
 491            "fake_add",
 492            params=EMPTY_SCHEMA,
 493            result=ToolCallResult(
 494                output="tool failed", is_error=True, error_message="tool failed"
 495            ),
 496        )
 497        code = textwrap.dedent("""\
 498            from kiln.tools import ToolCallError
 499            from kiln import tools
 500            def run(x):
 501                try:
 502                    tools.fake_add()
 503                except ToolCallError as e:
 504                    return f"error: {e.message}"
 505                return "no error"
 506        """)
 507        tool = _make_python_code_tool(
 508            tmp_path,
 509            code,
 510            tool_allowlist=["kiln_tool::add_numbers"],
 511        )
 512        with patch(
 513            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 514            return_value=fake,
 515        ):
 516            result = await tool.run(None, x="test")
 517        assert not result.is_error
 518        assert "error: tool failed" in result.output
 519
 520    @pytest.mark.asyncio
 521    async def test_nested_tool_not_allowed(self, tmp_path):
 522        fake = FakeTool(
 523            "kiln_tool::add_numbers",
 524            "fake_add",
 525            params=EMPTY_SCHEMA,
 526        )
 527        code = textwrap.dedent("""\
 528            from kiln.tools import ToolNotAllowed
 529            from kiln import tools
 530            def run(x):
 531                try:
 532                    tools.nonexistent_tool()
 533                except ToolNotAllowed as e:
 534                    return f"not allowed: {e.tool}"
 535                return "no error"
 536        """)
 537        tool = _make_python_code_tool(
 538            tmp_path,
 539            code,
 540            tool_allowlist=["kiln_tool::add_numbers"],
 541        )
 542        with patch(
 543            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 544            return_value=fake,
 545        ):
 546            result = await tool.run(None, x="test")
 547        assert not result.is_error
 548        assert "not allowed: nonexistent_tool" in result.output
 549
 550    @pytest.mark.asyncio
 551    async def test_nested_tool_ambiguous(self, tmp_path):
 552        fake1 = FakeTool("mcp::remote::server1::search", "search")
 553        fake2 = FakeTool("mcp::remote::server2::search", "search")
 554        fakes = {
 555            "mcp::remote::server1::search": fake1,
 556            "mcp::remote::server2::search": fake2,
 557        }
 558        code = textwrap.dedent("""\
 559            from kiln.tools import ToolCallError
 560            from kiln import tools
 561            def run(x):
 562                try:
 563                    tools.search()
 564                except ToolCallError as e:
 565                    return f"ambiguous: {e.message}"
 566                return "no error"
 567        """)
 568        tool = _make_python_code_tool(
 569            tmp_path,
 570            code,
 571            tool_allowlist=[
 572                "mcp::remote::server1::search",
 573                "mcp::remote::server2::search",
 574            ],
 575        )
 576        with patch(
 577            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 578            side_effect=lambda tid, **kw: fakes[tid],
 579        ):
 580            result = await tool.run(None, x="test")
 581        assert not result.is_error
 582        assert "ambiguous" in result.output.lower()
 583
 584    @pytest.mark.asyncio
 585    async def test_nested_tool_invalid_kwargs(self, tmp_path):
 586        fake = FakeTool(
 587            "kiln_tool::add_numbers",
 588            "fake_add",
 589            params={
 590                "type": "object",
 591                "properties": {"a": {"type": "integer"}},
 592                "required": ["a"],
 593            },
 594            result=ToolCallResult(output="42"),
 595        )
 596        code = textwrap.dedent("""\
 597            from kiln.tools import ToolCallError
 598            from kiln import tools
 599            def run(x):
 600                try:
 601                    tools.fake_add(a="not_an_int")
 602                except ToolCallError as e:
 603                    return f"invalid: {e.tool}"
 604                return "no error"
 605        """)
 606        tool = _make_python_code_tool(
 607            tmp_path,
 608            code,
 609            tool_allowlist=["kiln_tool::add_numbers"],
 610        )
 611        with patch(
 612            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 613            return_value=fake,
 614        ):
 615            result = await tool.run(None, x="test")
 616        assert not result.is_error
 617        assert "invalid: fake_add" in result.output
 618
 619
 620class TestListTools:
 621    @pytest.mark.asyncio
 622    async def test_list_tools_returns_content(self, tmp_path):
 623        fake = FakeTool(
 624            "kiln_tool::add_numbers",
 625            "fake_add",
 626            fn_desc="Add two numbers",
 627            params={
 628                "type": "object",
 629                "properties": {"a": {"type": "integer"}},
 630            },
 631        )
 632        code = textwrap.dedent("""\
 633            import json
 634            from kiln import tools
 635            def run(x):
 636                tool_list = tools.list_tools()
 637                return json.dumps(tool_list)
 638        """)
 639        tool = _make_python_code_tool(
 640            tmp_path,
 641            code,
 642            tool_allowlist=["kiln_tool::add_numbers"],
 643        )
 644        with patch(
 645            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 646            return_value=fake,
 647        ):
 648            result = await tool.run(None, x="test")
 649        assert not result.is_error
 650        tool_list = json.loads(result.output)
 651        assert len(tool_list) == 1
 652        assert tool_list[0]["name"] == "fake_add"
 653        assert tool_list[0]["description"] == "Add two numbers"
 654
 655
 656class TestTimeout:
 657    @pytest.mark.asyncio
 658    async def test_timeout_kills_child(self, tmp_path):
 659        tool = _make_python_code_tool(
 660            tmp_path,
 661            "import time\ndef run(x):\n    time.sleep(30)\n    return 'done'\n",
 662            timeout_seconds=1,
 663        )
 664        result = await tool.run(None, x="test")
 665        assert result.is_error
 666        assert "timed out" in result.output
 667
 668    @pytest.mark.asyncio
 669    async def test_timeout_during_nested_call(self, tmp_path):
 670        slow_fake = FakeTool(
 671            "kiln_tool::add_numbers",
 672            "fake_add",
 673            params=EMPTY_SCHEMA,
 674            result=ToolCallResult(output="42"),
 675            delay=30,
 676        )
 677        code = textwrap.dedent("""\
 678            from kiln import tools
 679            def run(x):
 680                return tools.fake_add()
 681        """)
 682        tool = _make_python_code_tool(
 683            tmp_path,
 684            code,
 685            tool_allowlist=["kiln_tool::add_numbers"],
 686            timeout_seconds=1,
 687        )
 688        with patch(
 689            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 690            return_value=slow_fake,
 691        ):
 692            result = await tool.run(None, x="test")
 693        assert result.is_error
 694        assert "timed out" in result.output
 695
 696
 697class TestCrash:
 698    @pytest.mark.asyncio
 699    async def test_crash_via_os_exit(self, tmp_path):
 700        tool = _make_python_code_tool(
 701            tmp_path,
 702            "import os\ndef run(x):\n    os._exit(3)\n",
 703        )
 704        result = await tool.run(None, x="test")
 705        assert result.is_error
 706        assert "crashed" in result.output
 707        assert "exit code" in result.output
 708
 709
 710class TestDepthCap:
 711    @pytest.mark.asyncio
 712    async def test_depth_cap_at_10(self, tmp_path):
 713        """Depth >= 10 returns an error without spawning."""
 714        tool = _make_python_code_tool(tmp_path, 'def run(x):\n    return "ok"\n')
 715        token = _depth.set(10)
 716        try:
 717            result = await tool.run(None, x="test")
 718        finally:
 719            _depth.reset(token)
 720        assert result.is_error
 721        assert "maximum nested code execution depth exceeded" in result.output
 722
 723
 724class TestSemaphore:
 725    @pytest.mark.asyncio
 726    async def test_semaphore_top_level_only_no_deadlock(self, tmp_path):
 727        """Regression: nested code-tool calls bypass the semaphore.
 728
 729        If nested calls counted against the semaphore, 8 parents each
 730        spawning a nested code-tool child would deadlock (parents hold
 731        all 8 slots, children wait forever).
 732
 733        This test sets MAX_CONCURRENCY parents running concurrently,
 734        each at depth 1 (simulating nested calls). All should complete
 735        without deadlock because nested calls bypass the semaphore.
 736        """
 737        code = 'def run(x):\n    return "nested_ok"\n'
 738        results = []
 739
 740        async def run_nested(i: int):
 741            tool = _make_python_code_tool(tmp_path, code)
 742            token = _depth.set(1)
 743            try:
 744                r = await tool.run(None, x=str(i))
 745                results.append(r)
 746            finally:
 747                _depth.reset(token)
 748
 749        await asyncio.gather(
 750            *(run_nested(i) for i in range(CODE_SANDBOX_MAX_CONCURRENCY))
 751        )
 752        assert len(results) == CODE_SANDBOX_MAX_CONCURRENCY
 753        assert all(not r.is_error for r in results)
 754
 755
 756class TestToolCallRecorder:
 757    @pytest.mark.asyncio
 758    async def test_recorder_gets_entries(self, tmp_path):
 759        fake = FakeTool(
 760            "kiln_tool::add_numbers",
 761            "fake_add",
 762            params=EMPTY_SCHEMA,
 763            result=ToolCallResult(output="42"),
 764        )
 765        log: list[ToolCallLogEntry] = []
 766        code = textwrap.dedent("""\
 767            from kiln import tools
 768            def run(x):
 769                r = tools.fake_add()
 770                return r
 771        """)
 772        tool = _make_python_code_tool(
 773            tmp_path,
 774            code,
 775            tool_allowlist=["kiln_tool::add_numbers"],
 776            tool_call_recorder=log.append,
 777        )
 778        with patch(
 779            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 780            return_value=fake,
 781        ):
 782            result = await tool.run(None, x="test")
 783        assert not result.is_error
 784        assert len(log) == 1
 785        assert log[0].tool_name == "fake_add"
 786        assert not log[0].is_error
 787        assert log[0].output_preview == "42"
 788        assert log[0].duration_ms >= 0
 789
 790
 791class TestAsyncToolsConcurrency:
 792    @pytest.mark.asyncio
 793    async def test_async_tools_gather_truly_concurrent(self, tmp_path):
 794        """async_tools + gather provides real parallelism via to_thread.
 795
 796        Two fake tools each take ~0.3s. If sequential, wall clock >= 0.6s.
 797        With true concurrency via gather + to_thread, wall clock < 0.6s.
 798        """
 799        slow_fake = FakeTool(
 800            "kiln_tool::add_numbers",
 801            "fake_add",
 802            params=EMPTY_SCHEMA,
 803            result=ToolCallResult(output="done"),
 804            delay=0.3,
 805        )
 806        code = textwrap.dedent("""\
 807            import asyncio, time
 808            from kiln import async_tools
 809            async def run(x):
 810                start = time.monotonic()
 811                a, b = await asyncio.gather(
 812                    async_tools.fake_add(),
 813                    async_tools.fake_add(),
 814                )
 815                elapsed = time.monotonic() - start
 816                return f"{elapsed:.2f}"
 817        """)
 818        tool = _make_python_code_tool(
 819            tmp_path,
 820            code,
 821            tool_allowlist=["kiln_tool::add_numbers"],
 822        )
 823        with patch(
 824            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 825            return_value=slow_fake,
 826        ):
 827            result = await tool.run(None, x="test")
 828        assert not result.is_error
 829        elapsed = float(result.output)
 830        assert elapsed < 0.6, f"Expected < 0.6s (concurrent), got {elapsed:.2f}s"
 831
 832
 833class _EchoFakeTool(KilnToolInterface):
 834    """Fake tool that echoes its kwargs back, proving per-call routing."""
 835
 836    def __init__(self, tool_id: str, fn_name: str, params: dict):
 837        self._id = tool_id
 838        self._name = fn_name
 839        self._params = params
 840
 841    async def id(self):
 842        return self._id
 843
 844    async def name(self):
 845        return self._name
 846
 847    async def description(self):
 848        return "echo"
 849
 850    async def toolcall_definition(self) -> ToolCallDefinition:
 851        return {
 852            "type": "function",
 853            "function": {
 854                "name": self._name,
 855                "description": "echo",
 856                "parameters": self._params,
 857            },
 858        }
 859
 860    async def run(self, context=None, **kwargs) -> ToolCallResult:
 861        return ToolCallResult(output=json.dumps(kwargs, sort_keys=True))
 862
 863
 864class TestCallIdRouting:
 865    @pytest.mark.asyncio
 866    async def test_concurrent_calls_routed_to_correct_caller(self, tmp_path):
 867        """4 threads each pass a unique idx kwarg; each gets its own value back.
 868
 869        The echo tool returns the kwargs it received. Each thread asserts it
 870        got back the idx it sent, proving call_id routing maps the right
 871        response to the right waiting caller under concurrency.
 872        """
 873        idx_schema = {
 874            "type": "object",
 875            "properties": {"idx": {"type": "string"}},
 876            "required": ["idx"],
 877        }
 878        code = textwrap.dedent("""\
 879            import json, threading
 880            from kiln import tools
 881            def run(x):
 882                results = [None] * 4
 883                errors = []
 884                def call_tool(i):
 885                    try:
 886                        raw = tools.fake_echo(idx=str(i))
 887                        results[i] = json.loads(raw)["idx"]
 888                    except Exception as e:
 889                        errors.append(f"thread {i}: {e}")
 890                threads = [threading.Thread(target=call_tool, args=(i,)) for i in range(4)]
 891                for t in threads:
 892                    t.start()
 893                for t in threads:
 894                    t.join()
 895                if errors:
 896                    return "errors: " + str(errors)
 897                return ",".join(results)
 898        """)
 899        fake = _EchoFakeTool(
 900            "kiln_tool::add_numbers",
 901            "fake_echo",
 902            params=idx_schema,
 903        )
 904        tool = _make_python_code_tool(
 905            tmp_path,
 906            code,
 907            tool_allowlist=["kiln_tool::add_numbers"],
 908        )
 909        with patch(
 910            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 911            return_value=fake,
 912        ):
 913            result = await tool.run(None, x="test")
 914        assert not result.is_error
 915        parts = result.output.split(",")
 916        assert len(parts) == 4
 917        assert parts == ["0", "1", "2", "3"]
 918
 919
 920class TestUnicodePassthrough:
 921    @pytest.mark.asyncio
 922    async def test_unicode_not_escaped(self, tmp_path):
 923        """ensure_ascii=False: non-ASCII chars pass through un-escaped."""
 924        tool = _make_python_code_tool(
 925            tmp_path,
 926            'def run(x):\n    return {"name": "\\u65e5\\u672c\\u8a9e"}\n',
 927        )
 928        result = await tool.run(None, x="test")
 929        assert not result.is_error
 930        parsed = json.loads(result.output)
 931        assert parsed["name"] == "日本語"
 932        assert "\\u" not in result.output
 933
 934
 935class TestSpawnLockIdentity:
 936    def test_spawn_lock_shared(self):
 937        """Code tools and code evals spawn through the same bridge / _spawn_lock."""
 938        from kiln_ai.tools import sandbox_bridge
 939
 940        assert (
 941            sandbox_bridge.start_process_with_light_main.__module__
 942            == "kiln_ai.sandbox.spawn"
 943        )
 944        from kiln_ai.sandbox.spawn import _spawn_lock as shared_lock
 945
 946        assert shared_lock is _spawn_lock
 947
 948
 949# ---------------------------------------------------------------------------
 950# End-to-end tests using REAL built-in tools (no mocking)
 951# ---------------------------------------------------------------------------
 952
 953
 954class TestRealBuiltInTools:
 955    """Tests that exercise real built-in tool dispatch without mocking
 956    ``tool_from_id_and_project``, ensuring the name-derivation path is
 957    exercised end-to-end.
 958    """
 959
 960    @pytest.mark.asyncio
 961    async def test_keyword_call_by_canonical_name_succeeds(self, tmp_path):
 962        """tools.add(a=1, b=2) returns '3' — the canonical name from list_tools."""
 963        code = textwrap.dedent("""\
 964            from kiln import tools
 965            def run(x):
 966                return tools.add(a=1, b=2)
 967        """)
 968        tool = _make_python_code_tool(
 969            tmp_path,
 970            code,
 971            tool_allowlist=["kiln_tool::add_numbers"],
 972        )
 973        result = await tool.run(None, x="test")
 974        assert not result.is_error, f"Expected success, got: {result.output}"
 975        assert result.output == "3"
 976
 977    @pytest.mark.asyncio
 978    async def test_list_tools_driven_call_succeeds(self, tmp_path):
 979        """Call using the name returned by list_tools() succeeds."""
 980        code = textwrap.dedent("""\
 981            from kiln import tools
 982            def run(x):
 983                tl = tools.list_tools()
 984                fn_name = tl[0]["name"]
 985                result = getattr(tools, fn_name)(a=5, b=3)
 986                return fn_name + ":" + result
 987        """)
 988        tool = _make_python_code_tool(
 989            tmp_path,
 990            code,
 991            tool_allowlist=["kiln_tool::add_numbers"],
 992        )
 993        result = await tool.run(None, x="test")
 994        assert not result.is_error, f"Expected success, got: {result.output}"
 995        assert result.output == "add:8"
 996
 997    @pytest.mark.asyncio
 998    async def test_friendly_name_not_allowed(self, tmp_path):
 999        """tools.Addition(a=1,b=2) raises ToolNotAllowed listing canonical names."""
1000        code = textwrap.dedent("""\
1001            from kiln.tools import ToolNotAllowed
1002            from kiln import tools
1003            def run(x):
1004                try:
1005                    tools.Addition(a=1, b=2)
1006                except ToolNotAllowed as e:
1007                    return e.message
1008                return "no error"
1009        """)
1010        tool = _make_python_code_tool(
1011            tmp_path,
1012            code,
1013            tool_allowlist=["kiln_tool::add_numbers"],
1014        )
1015        result = await tool.run(None, x="test")
1016        assert not result.is_error
1017        assert "not available" in result.output
1018        assert "'add'" in result.output
1019
1020    @pytest.mark.asyncio
1021    async def test_nonsense_name_not_allowed(self, tmp_path):
1022        """tools.bad_tool() raises ToolNotAllowed listing available names."""
1023        code = textwrap.dedent("""\
1024            from kiln.tools import ToolNotAllowed
1025            from kiln import tools
1026            def run(x):
1027                try:
1028                    tools.bad_tool(a=1)
1029                except ToolNotAllowed as e:
1030                    return e.message
1031                return "no error"
1032        """)
1033        tool = _make_python_code_tool(
1034            tmp_path,
1035            code,
1036            tool_allowlist=["kiln_tool::add_numbers"],
1037        )
1038        result = await tool.run(None, x="test")
1039        assert not result.is_error
1040        assert "not available" in result.output
1041        assert "'add'" in result.output
1042
1043    @pytest.mark.asyncio
1044    async def test_positional_args_error_message(self, tmp_path):
1045        """tools.add(1, 2) raises ToolCallError mentioning keyword args and params."""
1046        code = textwrap.dedent("""\
1047            from kiln.tools import ToolCallError
1048            from kiln import tools
1049            def run(x):
1050                try:
1051                    tools.add(1, 2)
1052                except ToolCallError as e:
1053                    return e.message
1054                return "no error"
1055        """)
1056        tool = _make_python_code_tool(
1057            tmp_path,
1058            code,
1059            tool_allowlist=["kiln_tool::add_numbers"],
1060        )
1061        result = await tool.run(None, x="test")
1062        assert not result.is_error
1063        assert "keyword arguments" in result.output
1064        assert "tools.add(" in result.output
1065        assert "a: number (required)" in result.output
1066        assert "b: number (required)" in result.output
1067
1068    @pytest.mark.asyncio
1069    async def test_wrong_kwargs_error_shows_schema(self, tmp_path):
1070        """tools.add(x=1) raises ToolCallError showing expected parameters."""
1071        code = textwrap.dedent("""\
1072            from kiln.tools import ToolCallError
1073            from kiln import tools
1074            def run(x):
1075                try:
1076                    tools.add(x=1)
1077                except ToolCallError as e:
1078                    return e.message
1079                return "no error"
1080        """)
1081        tool = _make_python_code_tool(
1082            tmp_path,
1083            code,
1084            tool_allowlist=["kiln_tool::add_numbers"],
1085        )
1086        result = await tool.run(None, x="test")
1087        assert not result.is_error
1088        assert "Expected parameters:" in result.output
1089        assert "a: number (required)" in result.output
1090
1091    @pytest.mark.asyncio
1092    async def test_name_consistency_across_all_builtins(self, tmp_path):
1093        """For every KilnBuiltInToolId the dispatch-map name matches tool.name()
1094        AND matches what list_tools reports."""
1095        from kiln_ai.datamodel.tool_id import KilnBuiltInToolId
1096        from kiln_ai.tools.tool_registry import tool_from_id_and_project
1097
1098        project = _make_project(tmp_path)
1099
1100        math_ids = [
1101            KilnBuiltInToolId.ADD_NUMBERS,
1102            KilnBuiltInToolId.SUBTRACT_NUMBERS,
1103            KilnBuiltInToolId.MULTIPLY_NUMBERS,
1104            KilnBuiltInToolId.DIVIDE_NUMBERS,
1105        ]
1106
1107        for builtin_id in math_ids:
1108            tool_id = builtin_id.value
1109            real_tool = tool_from_id_and_project(tool_id, project=project)
1110            real_name = await real_tool.name()
1111
1112            ct = _make_code_tool(
1113                'def run(x): return "ok"',
1114                tool_allowlist=[tool_id],
1115            )
1116            ct.parent = project
1117            server = NestedToolServer(
1118                allowlist=ct.tool_allowlist, project=project, task=None, context=None
1119            )
1120            dispatch_names = list((await server.name_map()).keys())
1121
1122            assert dispatch_names == [real_name], (
1123                f"For {builtin_id}: dispatch name {dispatch_names} != "
1124                f"tool.name() '{real_name}'"
1125            )
1126
1127    @pytest.mark.asyncio
1128    async def test_async_proxy_keyword_call(self, tmp_path):
1129        """async_tools.subtract(a=5, b=3) returns '2'."""
1130        code = textwrap.dedent("""\
1131            from kiln import async_tools
1132            async def run(x):
1133                return await async_tools.subtract(a=5, b=3)
1134        """)
1135        tool = _make_python_code_tool(
1136            tmp_path,
1137            code,
1138            tool_allowlist=["kiln_tool::subtract_numbers"],
1139        )
1140        result = await tool.run(None, x="test")
1141        assert not result.is_error, f"Expected success, got: {result.output}"
1142        assert result.output == "2"
1143
1144    @pytest.mark.asyncio
1145    async def test_async_proxy_positional_error(self, tmp_path):
1146        """async_tools.add(1, 2) raises ToolCallError with a helpful message."""
1147        code = textwrap.dedent("""\
1148            from kiln.tools import ToolCallError
1149            from kiln import async_tools
1150            async def run(x):
1151                try:
1152                    await async_tools.add(1, 2)
1153                except ToolCallError as e:
1154                    return e.message
1155                return "no error"
1156        """)
1157        tool = _make_python_code_tool(
1158            tmp_path,
1159            code,
1160            tool_allowlist=["kiln_tool::add_numbers"],
1161        )
1162        result = await tool.run(None, x="test")
1163        assert not result.is_error
1164        assert "keyword arguments" in result.output
1165
1166    @pytest.mark.asyncio
1167    async def test_positional_on_nonsense_name_still_not_allowed(self, tmp_path):
1168        """tools.bad_tool(1) raises ToolNotAllowed (not TypeError), regardless of args."""
1169        code = textwrap.dedent("""\
1170            from kiln.tools import ToolNotAllowed
1171            from kiln import tools
1172            def run(x):
1173                try:
1174                    tools.bad_tool(1, 2)
1175                except ToolNotAllowed as e:
1176                    return e.message
1177                return "no error"
1178        """)
1179        tool = _make_python_code_tool(
1180            tmp_path,
1181            code,
1182            tool_allowlist=["kiln_tool::add_numbers"],
1183        )
1184        result = await tool.run(None, x="test")
1185        assert not result.is_error
1186        assert "not available" in result.output
1187
1188
1189# ---------------------------------------------------------------------------
1190# UI example validation tests
1191#
1192# These tests execute the EXACT code strings shown in the "Code Tool Examples"
1193# modal (app/web_ui/src/lib/utils/code_tool_helpers.ts → generateExamples()).
1194# If you change those examples, you MUST update these tests to match.
1195# ---------------------------------------------------------------------------
1196
1197# The example code strings are duplicated here intentionally so any drift
1198# between the UI and these tests causes a test failure during review.
1199
1200EXAMPLE_PARALLEL_WITH_RETRIES = """\
1201import json
1202import time
1203from concurrent.futures import ThreadPoolExecutor, as_completed
1204from kiln import tools
1205
1206def run(urls: list[str], max_retries: int = 3) -> str:
1207    \"\"\"Fetch multiple URLs in parallel with retries.\"\"\"
1208    results = {}
1209
1210    def fetch_with_retry(url):
1211        for attempt in range(max_retries):
1212            try:
1213                result = tools.fetch_url(url=url)
1214                return url, json.loads(result)
1215            except Exception as e:
1216                if attempt == max_retries - 1:
1217                    return url, {"error": str(e)}
1218                time.sleep(0.5 * (attempt + 1))
1219
1220    with ThreadPoolExecutor(max_workers=5) as pool:
1221        futures = [pool.submit(fetch_with_retry, u) for u in urls]
1222        for future in as_completed(futures):
1223            url, data = future.result()
1224            results[url] = data
1225
1226    return json.dumps(results)
1227"""
1228
1229EXAMPLE_ASYNC_FAN_OUT = """\
1230import json
1231import asyncio
1232from kiln import async_tools
1233
1234async def run(user_ids: list[str]) -> str:
1235    \"\"\"Fetch user details concurrently using async_tools.\"\"\"
1236    async def fetch_user(uid):
1237        result = await async_tools.get_user(id=uid)
1238        return json.loads(result)
1239
1240    users = await asyncio.gather(*(fetch_user(uid) for uid in user_ids))
1241    return json.dumps(users)
1242"""
1243
1244EXAMPLE_FILTER_AND_TRANSFORM = """\
1245import json
1246from kiln import tools
1247
1248def run(query: str, max_results: int = 10) -> str:
1249    \"\"\"Search and filter results, returning only relevant fields.\"\"\"
1250    raw = tools.search(query=query)
1251    results = json.loads(raw)
1252
1253    filtered = [
1254        {"title": r["title"], "url": r["url"]}
1255        for r in results[:max_results]
1256        if "title" in r and "url" in r
1257    ]
1258
1259    return json.dumps(filtered)
1260"""
1261
1262
1263class TestUIExampleParallelWithRetries:
1264    """Validate the 'Parallel with Retries' example from the Code Tool Examples modal."""
1265
1266    @pytest.mark.asyncio
1267    async def test_parallel_with_retries_happy_path(self, tmp_path):
1268        fetch_url_responses = {
1269            "https://a.com": '{"status": "ok_a"}',
1270            "https://b.com": '{"status": "ok_b"}',
1271        }
1272        fake = FakeTool(
1273            "mcp::remote::test_server::fetch_url",
1274            "fetch_url",
1275            fn_desc="Fetch a URL",
1276            params={
1277                "type": "object",
1278                "properties": {"url": {"type": "string"}},
1279                "required": ["url"],
1280            },
1281        )
1282
1283        async def route_fetch(context=None, **kwargs):
1284            url = kwargs["url"]
1285            return ToolCallResult(output=fetch_url_responses[url])
1286
1287        fake.run = route_fetch  # type: ignore[assignment]
1288
1289        tool = _make_python_code_tool(
1290            tmp_path,
1291            EXAMPLE_PARALLEL_WITH_RETRIES,
1292            tool_allowlist=["mcp::remote::test_server::fetch_url"],
1293            parameters_schema={
1294                "type": "object",
1295                "properties": {
1296                    "urls": {"type": "array", "items": {"type": "string"}},
1297                    "max_retries": {"type": "integer"},
1298                },
1299                "required": ["urls"],
1300            },
1301        )
1302        with patch(
1303            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1304            return_value=fake,
1305        ):
1306            result = await tool.run(
1307                None, urls=["https://a.com", "https://b.com"], max_retries=1
1308            )
1309        assert not result.is_error, f"Expected success, got: {result.output}"
1310        parsed = json.loads(result.output)
1311        assert parsed["https://a.com"] == {"status": "ok_a"}
1312        assert parsed["https://b.com"] == {"status": "ok_b"}
1313
1314    @pytest.mark.asyncio
1315    async def test_parallel_with_retries_error_fallback(self, tmp_path):
1316        """When a tool call fails, the retry logic catches the exception and
1317        returns an error dict after exhausting retries."""
1318        fake = FakeTool(
1319            "mcp::remote::test_server::fetch_url",
1320            "fetch_url",
1321            fn_desc="Fetch a URL",
1322            params={
1323                "type": "object",
1324                "properties": {"url": {"type": "string"}},
1325                "required": ["url"],
1326            },
1327            result=ToolCallResult(
1328                output="connection refused",
1329                is_error=True,
1330                error_message="connection refused",
1331            ),
1332        )
1333        tool = _make_python_code_tool(
1334            tmp_path,
1335            EXAMPLE_PARALLEL_WITH_RETRIES,
1336            tool_allowlist=["mcp::remote::test_server::fetch_url"],
1337            parameters_schema={
1338                "type": "object",
1339                "properties": {
1340                    "urls": {"type": "array", "items": {"type": "string"}},
1341                    "max_retries": {"type": "integer"},
1342                },
1343                "required": ["urls"],
1344            },
1345        )
1346        with patch(
1347            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1348            return_value=fake,
1349        ):
1350            result = await tool.run(None, urls=["https://fail.com"], max_retries=1)
1351        assert not result.is_error, f"Expected success, got: {result.output}"
1352        parsed = json.loads(result.output)
1353        assert "error" in parsed["https://fail.com"]
1354
1355
1356class TestUIExampleAsyncFanOut:
1357    """Validate the 'Async Fan-Out' example from the Code Tool Examples modal."""
1358
1359    @pytest.mark.asyncio
1360    async def test_async_fan_out_happy_path(self, tmp_path):
1361        user_data = {
1362            "u1": '{"name": "Alice", "id": "u1"}',
1363            "u2": '{"name": "Bob", "id": "u2"}',
1364        }
1365        fake = FakeTool(
1366            "mcp::remote::test_server::get_user",
1367            "get_user",
1368            fn_desc="Get user details",
1369            params={
1370                "type": "object",
1371                "properties": {"id": {"type": "string"}},
1372                "required": ["id"],
1373            },
1374        )
1375
1376        async def route_user(context=None, **kwargs):
1377            uid = kwargs["id"]
1378            return ToolCallResult(output=user_data[uid])
1379
1380        fake.run = route_user  # type: ignore[assignment]
1381
1382        tool = _make_python_code_tool(
1383            tmp_path,
1384            EXAMPLE_ASYNC_FAN_OUT,
1385            tool_allowlist=["mcp::remote::test_server::get_user"],
1386            parameters_schema={
1387                "type": "object",
1388                "properties": {
1389                    "user_ids": {"type": "array", "items": {"type": "string"}},
1390                },
1391                "required": ["user_ids"],
1392            },
1393        )
1394        with patch(
1395            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1396            return_value=fake,
1397        ):
1398            result = await tool.run(None, user_ids=["u1", "u2"])
1399        assert not result.is_error, f"Expected success, got: {result.output}"
1400        parsed = json.loads(result.output)
1401        assert len(parsed) == 2
1402        assert parsed[0] == {"name": "Alice", "id": "u1"}
1403        assert parsed[1] == {"name": "Bob", "id": "u2"}
1404
1405
1406class TestUIExampleFilterAndTransform:
1407    """Validate the 'Filter & Transform' example from the Code Tool Examples modal."""
1408
1409    @pytest.mark.asyncio
1410    async def test_filter_and_transform_happy_path(self, tmp_path):
1411        search_results = json.dumps(
1412            [
1413                {"title": "Result 1", "url": "https://1.com", "score": 0.9},
1414                {"title": "Result 2", "url": "https://2.com", "score": 0.8},
1415                {"description": "no title or url"},
1416                {"title": "Result 3", "url": "https://3.com", "score": 0.7},
1417            ]
1418        )
1419        fake = FakeTool(
1420            "mcp::remote::test_server::search",
1421            "search",
1422            fn_desc="Search",
1423            params={
1424                "type": "object",
1425                "properties": {"query": {"type": "string"}},
1426                "required": ["query"],
1427            },
1428            result=ToolCallResult(output=search_results),
1429        )
1430        tool = _make_python_code_tool(
1431            tmp_path,
1432            EXAMPLE_FILTER_AND_TRANSFORM,
1433            tool_allowlist=["mcp::remote::test_server::search"],
1434            parameters_schema={
1435                "type": "object",
1436                "properties": {
1437                    "query": {"type": "string"},
1438                    "max_results": {"type": "integer"},
1439                },
1440                "required": ["query"],
1441            },
1442        )
1443        with patch(
1444            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1445            return_value=fake,
1446        ):
1447            result = await tool.run(None, query="test query")
1448        assert not result.is_error, f"Expected success, got: {result.output}"
1449        parsed = json.loads(result.output)
1450        assert len(parsed) == 3
1451        assert parsed[0] == {"title": "Result 1", "url": "https://1.com"}
1452        assert parsed[1] == {"title": "Result 2", "url": "https://2.com"}
1453        assert parsed[2] == {"title": "Result 3", "url": "https://3.com"}
1454
1455    @pytest.mark.asyncio
1456    async def test_filter_and_transform_respects_max_results(self, tmp_path):
1457        search_results = json.dumps(
1458            [{"title": f"R{i}", "url": f"https://{i}.com"} for i in range(20)]
1459        )
1460        fake = FakeTool(
1461            "mcp::remote::test_server::search",
1462            "search",
1463            fn_desc="Search",
1464            params={
1465                "type": "object",
1466                "properties": {"query": {"type": "string"}},
1467                "required": ["query"],
1468            },
1469            result=ToolCallResult(output=search_results),
1470        )
1471        tool = _make_python_code_tool(
1472            tmp_path,
1473            EXAMPLE_FILTER_AND_TRANSFORM,
1474            tool_allowlist=["mcp::remote::test_server::search"],
1475            parameters_schema={
1476                "type": "object",
1477                "properties": {
1478                    "query": {"type": "string"},
1479                    "max_results": {"type": "integer"},
1480                },
1481                "required": ["query"],
1482            },
1483        )
1484        with patch(
1485            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1486            return_value=fake,
1487        ):
1488            result = await tool.run(None, query="test", max_results=3)
1489        assert not result.is_error, f"Expected success, got: {result.output}"
1490        parsed = json.loads(result.output)
1491        assert len(parsed) == 3