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 TestBrokenAllowlistEntry:
 657    """One unresolvable allowlist entry (e.g. a deleted RAG config) must not
 658    take down the whole nested-tool surface."""
 659
 660    BROKEN_ID = "kiln_tool::rag::missing"
 661
 662    def _patch_registry(self, healthy: FakeTool):
 663        def resolver(tool_id, project=None, task=None):
 664            if tool_id == self.BROKEN_ID:
 665                raise ValueError(f"RAG config not found: {tool_id}")
 666            return healthy
 667
 668        return patch(
 669            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 670            side_effect=resolver,
 671        )
 672
 673    @pytest.mark.asyncio
 674    async def test_list_tools_shows_healthy_and_unavailable(self, tmp_path):
 675        fake = FakeTool("kiln_tool::add_numbers", "fake_add", params=EMPTY_SCHEMA)
 676        code = textwrap.dedent("""\
 677            import json
 678            from kiln import tools
 679            def run(x):
 680                return json.dumps(tools.list_tools())
 681        """)
 682        tool = _make_python_code_tool(
 683            tmp_path,
 684            code,
 685            tool_allowlist=["kiln_tool::add_numbers", self.BROKEN_ID],
 686        )
 687        with self._patch_registry(fake):
 688            result = await tool.run(None, x="test")
 689        assert not result.is_error
 690        tool_list = json.loads(result.output)
 691        assert len(tool_list) == 2
 692        by_name = {t["name"]: t for t in tool_list}
 693        assert by_name["fake_add"]["description"] == "fake"
 694        broken = by_name[self.BROKEN_ID]
 695        assert broken["description"].startswith("(unavailable:")
 696        assert "RAG config not found" in broken["description"]
 697
 698    @pytest.mark.asyncio
 699    async def test_healthy_tool_still_callable(self, tmp_path):
 700        fake = FakeTool(
 701            "kiln_tool::add_numbers",
 702            "fake_add",
 703            params=EMPTY_SCHEMA,
 704            result=ToolCallResult(output="42"),
 705        )
 706        code = textwrap.dedent("""\
 707            from kiln import tools
 708            def run(x):
 709                return tools.fake_add()
 710        """)
 711        tool = _make_python_code_tool(
 712            tmp_path,
 713            code,
 714            tool_allowlist=["kiln_tool::add_numbers", self.BROKEN_ID],
 715        )
 716        with self._patch_registry(fake):
 717            result = await tool.run(None, x="test")
 718        assert not result.is_error
 719        assert result.output == "42"
 720
 721    @pytest.mark.asyncio
 722    async def test_calling_broken_tool_reports_unavailable(self, tmp_path):
 723        fake = FakeTool("kiln_tool::add_numbers", "fake_add", params=EMPTY_SCHEMA)
 724        code = textwrap.dedent("""\
 725            from kiln.tools import ToolNotAllowed
 726            from kiln import tools
 727            def run(x):
 728                try:
 729                    tools.missing_rag()
 730                except ToolNotAllowed as e:
 731                    return f"unavailable: {e}"
 732                return "no error"
 733        """)
 734        tool = _make_python_code_tool(
 735            tmp_path,
 736            code,
 737            tool_allowlist=["kiln_tool::add_numbers", self.BROKEN_ID],
 738        )
 739        with self._patch_registry(fake):
 740            result = await tool.run(None, x="test")
 741        assert not result.is_error
 742        assert "unavailable:" in result.output
 743        assert "not available" in result.output
 744        assert "fake_add" in result.output
 745
 746
 747class TestTimeout:
 748    @pytest.mark.asyncio
 749    async def test_timeout_kills_child(self, tmp_path):
 750        tool = _make_python_code_tool(
 751            tmp_path,
 752            "import time\ndef run(x):\n    time.sleep(30)\n    return 'done'\n",
 753            timeout_seconds=1,
 754        )
 755        result = await tool.run(None, x="test")
 756        assert result.is_error
 757        assert "timed out" in result.output
 758
 759    @pytest.mark.asyncio
 760    async def test_timeout_during_nested_call(self, tmp_path):
 761        slow_fake = FakeTool(
 762            "kiln_tool::add_numbers",
 763            "fake_add",
 764            params=EMPTY_SCHEMA,
 765            result=ToolCallResult(output="42"),
 766            delay=30,
 767        )
 768        code = textwrap.dedent("""\
 769            from kiln import tools
 770            def run(x):
 771                return tools.fake_add()
 772        """)
 773        tool = _make_python_code_tool(
 774            tmp_path,
 775            code,
 776            tool_allowlist=["kiln_tool::add_numbers"],
 777            timeout_seconds=1,
 778        )
 779        with patch(
 780            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 781            return_value=slow_fake,
 782        ):
 783            result = await tool.run(None, x="test")
 784        assert result.is_error
 785        assert "timed out" in result.output
 786
 787
 788class TestNestedTimeoutKind:
 789    """The dispatcher classifies a nested code tool's timeout from the typed
 790    ``timed_out`` flag on its result, never from the error text."""
 791
 792    def _nested_code_tool(self, project, code: str, timeout_seconds: int = 10):
 793        ct = _make_code_tool(
 794            code,
 795            name="Nested Tool",
 796            tool_function_name="nested_tool",
 797            parameters_schema=EMPTY_SCHEMA,
 798            tool_allowlist=[],
 799            timeout_seconds=timeout_seconds,
 800        )
 801        ct.parent = project
 802        return PythonCodeTool(ct, project)
 803
 804    OUTER_CODE = textwrap.dedent("""\
 805        from kiln.tools import ToolCallError, ToolTimeout
 806        from kiln import tools
 807        def run(x):
 808            try:
 809                tools.nested_tool()
 810            except ToolTimeout:
 811                return "timeout"
 812            except ToolCallError:
 813                return "call_error"
 814            return "no error"
 815    """)
 816
 817    @pytest.mark.asyncio
 818    async def test_real_nested_timeout_raises_tool_timeout(self, tmp_path):
 819        project = _make_project(tmp_path)
 820        nested = self._nested_code_tool(
 821            project,
 822            "import time\ndef run():\n    time.sleep(30)\n",
 823            timeout_seconds=1,
 824        )
 825        outer_ct = _make_code_tool(
 826            self.OUTER_CODE, tool_allowlist=["kiln_tool::add_numbers"]
 827        )
 828        outer_ct.parent = project
 829        outer = PythonCodeTool(outer_ct, project)
 830        with patch(
 831            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 832            return_value=nested,
 833        ):
 834            result = await outer.run(None, x="test")
 835        assert not result.is_error
 836        assert result.output == "timeout"
 837
 838    @pytest.mark.asyncio
 839    async def test_failure_text_mentioning_timeout_raises_call_error(self, tmp_path):
 840        # An ordinary failure whose message merely contains "timed out" must
 841        # not spoof the timeout kind (which callers treat as retryable).
 842        project = _make_project(tmp_path)
 843        nested = self._nested_code_tool(
 844            project,
 845            'def run():\n    raise Exception("upstream request timed out after 3 retries")\n',
 846        )
 847        outer_ct = _make_code_tool(
 848            self.OUTER_CODE, tool_allowlist=["kiln_tool::add_numbers"]
 849        )
 850        outer_ct.parent = project
 851        outer = PythonCodeTool(outer_ct, project)
 852        with patch(
 853            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 854            return_value=nested,
 855        ):
 856            result = await outer.run(None, x="test")
 857        assert not result.is_error
 858        assert result.output == "call_error"
 859
 860
 861class TestCrash:
 862    @pytest.mark.asyncio
 863    async def test_crash_via_os_exit(self, tmp_path):
 864        tool = _make_python_code_tool(
 865            tmp_path,
 866            "import os\ndef run(x):\n    os._exit(3)\n",
 867        )
 868        result = await tool.run(None, x="test")
 869        assert result.is_error
 870        assert "crashed" in result.output
 871        assert "exit code" in result.output
 872
 873
 874class TestDepthCap:
 875    @pytest.mark.asyncio
 876    async def test_depth_cap_at_10(self, tmp_path):
 877        """Depth >= 10 returns an error without spawning."""
 878        tool = _make_python_code_tool(tmp_path, 'def run(x):\n    return "ok"\n')
 879        token = _depth.set(10)
 880        try:
 881            result = await tool.run(None, x="test")
 882        finally:
 883            _depth.reset(token)
 884        assert result.is_error
 885        assert "maximum nested code execution depth exceeded" in result.output
 886
 887
 888class TestSemaphore:
 889    @pytest.mark.asyncio
 890    async def test_semaphore_top_level_only_no_deadlock(self, tmp_path):
 891        """Regression: nested code-tool calls bypass the semaphore.
 892
 893        If nested calls counted against the semaphore, 8 parents each
 894        spawning a nested code-tool child would deadlock (parents hold
 895        all 8 slots, children wait forever).
 896
 897        This test sets MAX_CONCURRENCY parents running concurrently,
 898        each at depth 1 (simulating nested calls). All should complete
 899        without deadlock because nested calls bypass the semaphore.
 900        """
 901        code = 'def run(x):\n    return "nested_ok"\n'
 902        results = []
 903
 904        async def run_nested(i: int):
 905            tool = _make_python_code_tool(tmp_path, code)
 906            token = _depth.set(1)
 907            try:
 908                r = await tool.run(None, x=str(i))
 909                results.append(r)
 910            finally:
 911                _depth.reset(token)
 912
 913        await asyncio.gather(
 914            *(run_nested(i) for i in range(CODE_SANDBOX_MAX_CONCURRENCY))
 915        )
 916        assert len(results) == CODE_SANDBOX_MAX_CONCURRENCY
 917        assert all(not r.is_error for r in results)
 918
 919
 920class TestToolCallRecorder:
 921    @pytest.mark.asyncio
 922    async def test_recorder_gets_entries(self, tmp_path):
 923        fake = FakeTool(
 924            "kiln_tool::add_numbers",
 925            "fake_add",
 926            params=EMPTY_SCHEMA,
 927            result=ToolCallResult(output="42"),
 928        )
 929        log: list[ToolCallLogEntry] = []
 930        code = textwrap.dedent("""\
 931            from kiln import tools
 932            def run(x):
 933                r = tools.fake_add()
 934                return r
 935        """)
 936        tool = _make_python_code_tool(
 937            tmp_path,
 938            code,
 939            tool_allowlist=["kiln_tool::add_numbers"],
 940            tool_call_recorder=log.append,
 941        )
 942        with patch(
 943            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 944            return_value=fake,
 945        ):
 946            result = await tool.run(None, x="test")
 947        assert not result.is_error
 948        assert len(log) == 1
 949        assert log[0].tool_name == "fake_add"
 950        assert not log[0].is_error
 951        assert log[0].output_preview == "42"
 952        assert log[0].duration_ms >= 0
 953
 954
 955class TestAsyncToolsConcurrency:
 956    @pytest.mark.asyncio
 957    async def test_async_tools_gather_truly_concurrent(self, tmp_path):
 958        """async_tools + gather provides real parallelism via to_thread.
 959
 960        Two fake tools each take ~0.3s. If sequential, wall clock >= 0.6s.
 961        With true concurrency via gather + to_thread, wall clock < 0.6s.
 962        """
 963        slow_fake = FakeTool(
 964            "kiln_tool::add_numbers",
 965            "fake_add",
 966            params=EMPTY_SCHEMA,
 967            result=ToolCallResult(output="done"),
 968            delay=0.3,
 969        )
 970        code = textwrap.dedent("""\
 971            import asyncio, time
 972            from kiln import async_tools
 973            async def run(x):
 974                start = time.monotonic()
 975                a, b = await asyncio.gather(
 976                    async_tools.fake_add(),
 977                    async_tools.fake_add(),
 978                )
 979                elapsed = time.monotonic() - start
 980                return f"{elapsed:.2f}"
 981        """)
 982        tool = _make_python_code_tool(
 983            tmp_path,
 984            code,
 985            tool_allowlist=["kiln_tool::add_numbers"],
 986        )
 987        with patch(
 988            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
 989            return_value=slow_fake,
 990        ):
 991            result = await tool.run(None, x="test")
 992        assert not result.is_error
 993        elapsed = float(result.output)
 994        assert elapsed < 0.6, f"Expected < 0.6s (concurrent), got {elapsed:.2f}s"
 995
 996
 997class _EchoFakeTool(KilnToolInterface):
 998    """Fake tool that echoes its kwargs back, proving per-call routing."""
 999
1000    def __init__(self, tool_id: str, fn_name: str, params: dict):
1001        self._id = tool_id
1002        self._name = fn_name
1003        self._params = params
1004
1005    async def id(self):
1006        return self._id
1007
1008    async def name(self):
1009        return self._name
1010
1011    async def description(self):
1012        return "echo"
1013
1014    async def toolcall_definition(self) -> ToolCallDefinition:
1015        return {
1016            "type": "function",
1017            "function": {
1018                "name": self._name,
1019                "description": "echo",
1020                "parameters": self._params,
1021            },
1022        }
1023
1024    async def run(self, context=None, **kwargs) -> ToolCallResult:
1025        return ToolCallResult(output=json.dumps(kwargs, sort_keys=True))
1026
1027
1028class TestCallIdRouting:
1029    @pytest.mark.asyncio
1030    async def test_concurrent_calls_routed_to_correct_caller(self, tmp_path):
1031        """4 threads each pass a unique idx kwarg; each gets its own value back.
1032
1033        The echo tool returns the kwargs it received. Each thread asserts it
1034        got back the idx it sent, proving call_id routing maps the right
1035        response to the right waiting caller under concurrency.
1036        """
1037        idx_schema = {
1038            "type": "object",
1039            "properties": {"idx": {"type": "string"}},
1040            "required": ["idx"],
1041        }
1042        code = textwrap.dedent("""\
1043            import json, threading
1044            from kiln import tools
1045            def run(x):
1046                results = [None] * 4
1047                errors = []
1048                def call_tool(i):
1049                    try:
1050                        raw = tools.fake_echo(idx=str(i))
1051                        results[i] = json.loads(raw)["idx"]
1052                    except Exception as e:
1053                        errors.append(f"thread {i}: {e}")
1054                threads = [threading.Thread(target=call_tool, args=(i,)) for i in range(4)]
1055                for t in threads:
1056                    t.start()
1057                for t in threads:
1058                    t.join()
1059                if errors:
1060                    return "errors: " + str(errors)
1061                return ",".join(results)
1062        """)
1063        fake = _EchoFakeTool(
1064            "kiln_tool::add_numbers",
1065            "fake_echo",
1066            params=idx_schema,
1067        )
1068        tool = _make_python_code_tool(
1069            tmp_path,
1070            code,
1071            tool_allowlist=["kiln_tool::add_numbers"],
1072        )
1073        with patch(
1074            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1075            return_value=fake,
1076        ):
1077            result = await tool.run(None, x="test")
1078        assert not result.is_error
1079        parts = result.output.split(",")
1080        assert len(parts) == 4
1081        assert parts == ["0", "1", "2", "3"]
1082
1083
1084class TestUnicodePassthrough:
1085    @pytest.mark.asyncio
1086    async def test_unicode_not_escaped(self, tmp_path):
1087        """ensure_ascii=False: non-ASCII chars pass through un-escaped."""
1088        tool = _make_python_code_tool(
1089            tmp_path,
1090            'def run(x):\n    return {"name": "\\u65e5\\u672c\\u8a9e"}\n',
1091        )
1092        result = await tool.run(None, x="test")
1093        assert not result.is_error
1094        parsed = json.loads(result.output)
1095        assert parsed["name"] == "日本語"
1096        assert "\\u" not in result.output
1097
1098
1099class TestSpawnLockIdentity:
1100    def test_spawn_lock_shared(self):
1101        """Code tools and code evals spawn through the same bridge / _spawn_lock."""
1102        from kiln_ai.tools import sandbox_bridge
1103
1104        assert (
1105            sandbox_bridge.start_process_with_light_main.__module__
1106            == "kiln_ai.sandbox.spawn"
1107        )
1108        from kiln_ai.sandbox.spawn import _spawn_lock as shared_lock
1109
1110        assert shared_lock is _spawn_lock
1111
1112
1113# ---------------------------------------------------------------------------
1114# End-to-end tests using REAL built-in tools (no mocking)
1115# ---------------------------------------------------------------------------
1116
1117
1118class TestRealBuiltInTools:
1119    """Tests that exercise real built-in tool dispatch without mocking
1120    ``tool_from_id_and_project``, ensuring the name-derivation path is
1121    exercised end-to-end.
1122    """
1123
1124    @pytest.mark.asyncio
1125    async def test_keyword_call_by_canonical_name_succeeds(self, tmp_path):
1126        """tools.add(a=1, b=2) returns '3' — the canonical name from list_tools."""
1127        code = textwrap.dedent("""\
1128            from kiln import tools
1129            def run(x):
1130                return tools.add(a=1, b=2)
1131        """)
1132        tool = _make_python_code_tool(
1133            tmp_path,
1134            code,
1135            tool_allowlist=["kiln_tool::add_numbers"],
1136        )
1137        result = await tool.run(None, x="test")
1138        assert not result.is_error, f"Expected success, got: {result.output}"
1139        assert result.output == "3"
1140
1141    @pytest.mark.asyncio
1142    async def test_list_tools_driven_call_succeeds(self, tmp_path):
1143        """Call using the name returned by list_tools() succeeds."""
1144        code = textwrap.dedent("""\
1145            from kiln import tools
1146            def run(x):
1147                tl = tools.list_tools()
1148                fn_name = tl[0]["name"]
1149                result = getattr(tools, fn_name)(a=5, b=3)
1150                return fn_name + ":" + result
1151        """)
1152        tool = _make_python_code_tool(
1153            tmp_path,
1154            code,
1155            tool_allowlist=["kiln_tool::add_numbers"],
1156        )
1157        result = await tool.run(None, x="test")
1158        assert not result.is_error, f"Expected success, got: {result.output}"
1159        assert result.output == "add:8"
1160
1161    @pytest.mark.asyncio
1162    async def test_friendly_name_not_allowed(self, tmp_path):
1163        """tools.Addition(a=1,b=2) raises ToolNotAllowed listing canonical names."""
1164        code = textwrap.dedent("""\
1165            from kiln.tools import ToolNotAllowed
1166            from kiln import tools
1167            def run(x):
1168                try:
1169                    tools.Addition(a=1, b=2)
1170                except ToolNotAllowed as e:
1171                    return e.message
1172                return "no error"
1173        """)
1174        tool = _make_python_code_tool(
1175            tmp_path,
1176            code,
1177            tool_allowlist=["kiln_tool::add_numbers"],
1178        )
1179        result = await tool.run(None, x="test")
1180        assert not result.is_error
1181        assert "not available" in result.output
1182        assert "'add'" in result.output
1183
1184    @pytest.mark.asyncio
1185    async def test_nonsense_name_not_allowed(self, tmp_path):
1186        """tools.bad_tool() raises ToolNotAllowed listing available names."""
1187        code = textwrap.dedent("""\
1188            from kiln.tools import ToolNotAllowed
1189            from kiln import tools
1190            def run(x):
1191                try:
1192                    tools.bad_tool(a=1)
1193                except ToolNotAllowed as e:
1194                    return e.message
1195                return "no error"
1196        """)
1197        tool = _make_python_code_tool(
1198            tmp_path,
1199            code,
1200            tool_allowlist=["kiln_tool::add_numbers"],
1201        )
1202        result = await tool.run(None, x="test")
1203        assert not result.is_error
1204        assert "not available" in result.output
1205        assert "'add'" in result.output
1206
1207    @pytest.mark.asyncio
1208    async def test_positional_args_error_message(self, tmp_path):
1209        """tools.add(1, 2) raises ToolCallError mentioning keyword args and params."""
1210        code = textwrap.dedent("""\
1211            from kiln.tools import ToolCallError
1212            from kiln import tools
1213            def run(x):
1214                try:
1215                    tools.add(1, 2)
1216                except ToolCallError as e:
1217                    return e.message
1218                return "no error"
1219        """)
1220        tool = _make_python_code_tool(
1221            tmp_path,
1222            code,
1223            tool_allowlist=["kiln_tool::add_numbers"],
1224        )
1225        result = await tool.run(None, x="test")
1226        assert not result.is_error
1227        assert "keyword arguments" in result.output
1228        assert "tools.add(" in result.output
1229        assert "a: number (required)" in result.output
1230        assert "b: number (required)" in result.output
1231
1232    @pytest.mark.asyncio
1233    async def test_wrong_kwargs_error_shows_schema(self, tmp_path):
1234        """tools.add(x=1) raises ToolCallError showing expected parameters."""
1235        code = textwrap.dedent("""\
1236            from kiln.tools import ToolCallError
1237            from kiln import tools
1238            def run(x):
1239                try:
1240                    tools.add(x=1)
1241                except ToolCallError as e:
1242                    return e.message
1243                return "no error"
1244        """)
1245        tool = _make_python_code_tool(
1246            tmp_path,
1247            code,
1248            tool_allowlist=["kiln_tool::add_numbers"],
1249        )
1250        result = await tool.run(None, x="test")
1251        assert not result.is_error
1252        assert "Expected parameters:" in result.output
1253        assert "a: number (required)" in result.output
1254
1255    @pytest.mark.asyncio
1256    async def test_name_consistency_across_all_builtins(self, tmp_path):
1257        """For every KilnBuiltInToolId the dispatch-map name matches tool.name()
1258        AND matches what list_tools reports."""
1259        from kiln_ai.datamodel.tool_id import KilnBuiltInToolId
1260        from kiln_ai.tools.tool_registry import tool_from_id_and_project
1261
1262        project = _make_project(tmp_path)
1263
1264        math_ids = [
1265            KilnBuiltInToolId.ADD_NUMBERS,
1266            KilnBuiltInToolId.SUBTRACT_NUMBERS,
1267            KilnBuiltInToolId.MULTIPLY_NUMBERS,
1268            KilnBuiltInToolId.DIVIDE_NUMBERS,
1269        ]
1270
1271        for builtin_id in math_ids:
1272            tool_id = builtin_id.value
1273            real_tool = tool_from_id_and_project(tool_id, project=project)
1274            real_name = await real_tool.name()
1275
1276            ct = _make_code_tool(
1277                'def run(x): return "ok"',
1278                tool_allowlist=[tool_id],
1279            )
1280            ct.parent = project
1281            server = NestedToolServer(
1282                allowlist=ct.tool_allowlist, project=project, task=None, context=None
1283            )
1284            dispatch_names = list((await server.name_map()).keys())
1285
1286            assert dispatch_names == [real_name], (
1287                f"For {builtin_id}: dispatch name {dispatch_names} != "
1288                f"tool.name() '{real_name}'"
1289            )
1290
1291    @pytest.mark.asyncio
1292    async def test_async_proxy_keyword_call(self, tmp_path):
1293        """async_tools.subtract(a=5, b=3) returns '2'."""
1294        code = textwrap.dedent("""\
1295            from kiln import async_tools
1296            async def run(x):
1297                return await async_tools.subtract(a=5, b=3)
1298        """)
1299        tool = _make_python_code_tool(
1300            tmp_path,
1301            code,
1302            tool_allowlist=["kiln_tool::subtract_numbers"],
1303        )
1304        result = await tool.run(None, x="test")
1305        assert not result.is_error, f"Expected success, got: {result.output}"
1306        assert result.output == "2"
1307
1308    @pytest.mark.asyncio
1309    async def test_async_proxy_positional_error(self, tmp_path):
1310        """async_tools.add(1, 2) raises ToolCallError with a helpful message."""
1311        code = textwrap.dedent("""\
1312            from kiln.tools import ToolCallError
1313            from kiln import async_tools
1314            async def run(x):
1315                try:
1316                    await async_tools.add(1, 2)
1317                except ToolCallError as e:
1318                    return e.message
1319                return "no error"
1320        """)
1321        tool = _make_python_code_tool(
1322            tmp_path,
1323            code,
1324            tool_allowlist=["kiln_tool::add_numbers"],
1325        )
1326        result = await tool.run(None, x="test")
1327        assert not result.is_error
1328        assert "keyword arguments" in result.output
1329
1330    @pytest.mark.asyncio
1331    async def test_positional_on_nonsense_name_still_not_allowed(self, tmp_path):
1332        """tools.bad_tool(1) raises ToolNotAllowed (not TypeError), regardless of args."""
1333        code = textwrap.dedent("""\
1334            from kiln.tools import ToolNotAllowed
1335            from kiln import tools
1336            def run(x):
1337                try:
1338                    tools.bad_tool(1, 2)
1339                except ToolNotAllowed as e:
1340                    return e.message
1341                return "no error"
1342        """)
1343        tool = _make_python_code_tool(
1344            tmp_path,
1345            code,
1346            tool_allowlist=["kiln_tool::add_numbers"],
1347        )
1348        result = await tool.run(None, x="test")
1349        assert not result.is_error
1350        assert "not available" in result.output
1351
1352
1353# ---------------------------------------------------------------------------
1354# UI example validation tests
1355#
1356# These tests execute the EXACT code strings shown in the "Code Tool Examples"
1357# modal (app/web_ui/src/lib/utils/code_tool_helpers.ts → generateExamples()).
1358# If you change those examples, you MUST update these tests to match.
1359# ---------------------------------------------------------------------------
1360
1361# The example code strings are duplicated here intentionally so any drift
1362# between the UI and these tests causes a test failure during review.
1363
1364EXAMPLE_PARALLEL_WITH_RETRIES = """\
1365import json
1366import time
1367from concurrent.futures import ThreadPoolExecutor, as_completed
1368from kiln import tools
1369
1370def run(urls: list[str], max_retries: int = 3) -> str:
1371    \"\"\"Fetch multiple URLs in parallel with retries.\"\"\"
1372    results = {}
1373
1374    def fetch_with_retry(url):
1375        for attempt in range(max_retries):
1376            try:
1377                result = tools.fetch_url(url=url)
1378                return url, json.loads(result)
1379            except Exception as e:
1380                if attempt == max_retries - 1:
1381                    return url, {"error": str(e)}
1382                time.sleep(0.5 * (attempt + 1))
1383
1384    with ThreadPoolExecutor(max_workers=5) as pool:
1385        futures = [pool.submit(fetch_with_retry, u) for u in urls]
1386        for future in as_completed(futures):
1387            url, data = future.result()
1388            results[url] = data
1389
1390    return json.dumps(results)
1391"""
1392
1393EXAMPLE_ASYNC_FAN_OUT = """\
1394import json
1395import asyncio
1396from kiln import async_tools
1397
1398async def run(user_ids: list[str]) -> str:
1399    \"\"\"Fetch user details concurrently using async_tools.\"\"\"
1400    async def fetch_user(uid):
1401        result = await async_tools.get_user(id=uid)
1402        return json.loads(result)
1403
1404    users = await asyncio.gather(*(fetch_user(uid) for uid in user_ids))
1405    return json.dumps(users)
1406"""
1407
1408EXAMPLE_FILTER_AND_TRANSFORM = """\
1409import json
1410from kiln import tools
1411
1412def run(query: str, max_results: int = 10) -> str:
1413    \"\"\"Search and filter results, returning only relevant fields.\"\"\"
1414    raw = tools.search(query=query)
1415    results = json.loads(raw)
1416
1417    filtered = [
1418        {"title": r["title"], "url": r["url"]}
1419        for r in results[:max_results]
1420        if "title" in r and "url" in r
1421    ]
1422
1423    return json.dumps(filtered)
1424"""
1425
1426
1427class TestUIExampleParallelWithRetries:
1428    """Validate the 'Parallel with Retries' example from the Code Tool Examples modal."""
1429
1430    @pytest.mark.asyncio
1431    async def test_parallel_with_retries_happy_path(self, tmp_path):
1432        fetch_url_responses = {
1433            "https://a.com": '{"status": "ok_a"}',
1434            "https://b.com": '{"status": "ok_b"}',
1435        }
1436        fake = FakeTool(
1437            "mcp::remote::test_server::fetch_url",
1438            "fetch_url",
1439            fn_desc="Fetch a URL",
1440            params={
1441                "type": "object",
1442                "properties": {"url": {"type": "string"}},
1443                "required": ["url"],
1444            },
1445        )
1446
1447        async def route_fetch(context=None, **kwargs):
1448            url = kwargs["url"]
1449            return ToolCallResult(output=fetch_url_responses[url])
1450
1451        fake.run = route_fetch  # type: ignore[assignment]
1452
1453        tool = _make_python_code_tool(
1454            tmp_path,
1455            EXAMPLE_PARALLEL_WITH_RETRIES,
1456            tool_allowlist=["mcp::remote::test_server::fetch_url"],
1457            parameters_schema={
1458                "type": "object",
1459                "properties": {
1460                    "urls": {"type": "array", "items": {"type": "string"}},
1461                    "max_retries": {"type": "integer"},
1462                },
1463                "required": ["urls"],
1464            },
1465        )
1466        with patch(
1467            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1468            return_value=fake,
1469        ):
1470            result = await tool.run(
1471                None, urls=["https://a.com", "https://b.com"], max_retries=1
1472            )
1473        assert not result.is_error, f"Expected success, got: {result.output}"
1474        parsed = json.loads(result.output)
1475        assert parsed["https://a.com"] == {"status": "ok_a"}
1476        assert parsed["https://b.com"] == {"status": "ok_b"}
1477
1478    @pytest.mark.asyncio
1479    async def test_parallel_with_retries_error_fallback(self, tmp_path):
1480        """When a tool call fails, the retry logic catches the exception and
1481        returns an error dict after exhausting retries."""
1482        fake = FakeTool(
1483            "mcp::remote::test_server::fetch_url",
1484            "fetch_url",
1485            fn_desc="Fetch a URL",
1486            params={
1487                "type": "object",
1488                "properties": {"url": {"type": "string"}},
1489                "required": ["url"],
1490            },
1491            result=ToolCallResult(
1492                output="connection refused",
1493                is_error=True,
1494                error_message="connection refused",
1495            ),
1496        )
1497        tool = _make_python_code_tool(
1498            tmp_path,
1499            EXAMPLE_PARALLEL_WITH_RETRIES,
1500            tool_allowlist=["mcp::remote::test_server::fetch_url"],
1501            parameters_schema={
1502                "type": "object",
1503                "properties": {
1504                    "urls": {"type": "array", "items": {"type": "string"}},
1505                    "max_retries": {"type": "integer"},
1506                },
1507                "required": ["urls"],
1508            },
1509        )
1510        with patch(
1511            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1512            return_value=fake,
1513        ):
1514            result = await tool.run(None, urls=["https://fail.com"], max_retries=1)
1515        assert not result.is_error, f"Expected success, got: {result.output}"
1516        parsed = json.loads(result.output)
1517        assert "error" in parsed["https://fail.com"]
1518
1519
1520class TestUIExampleAsyncFanOut:
1521    """Validate the 'Async Fan-Out' example from the Code Tool Examples modal."""
1522
1523    @pytest.mark.asyncio
1524    async def test_async_fan_out_happy_path(self, tmp_path):
1525        user_data = {
1526            "u1": '{"name": "Alice", "id": "u1"}',
1527            "u2": '{"name": "Bob", "id": "u2"}',
1528        }
1529        fake = FakeTool(
1530            "mcp::remote::test_server::get_user",
1531            "get_user",
1532            fn_desc="Get user details",
1533            params={
1534                "type": "object",
1535                "properties": {"id": {"type": "string"}},
1536                "required": ["id"],
1537            },
1538        )
1539
1540        async def route_user(context=None, **kwargs):
1541            uid = kwargs["id"]
1542            return ToolCallResult(output=user_data[uid])
1543
1544        fake.run = route_user  # type: ignore[assignment]
1545
1546        tool = _make_python_code_tool(
1547            tmp_path,
1548            EXAMPLE_ASYNC_FAN_OUT,
1549            tool_allowlist=["mcp::remote::test_server::get_user"],
1550            parameters_schema={
1551                "type": "object",
1552                "properties": {
1553                    "user_ids": {"type": "array", "items": {"type": "string"}},
1554                },
1555                "required": ["user_ids"],
1556            },
1557        )
1558        with patch(
1559            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1560            return_value=fake,
1561        ):
1562            result = await tool.run(None, user_ids=["u1", "u2"])
1563        assert not result.is_error, f"Expected success, got: {result.output}"
1564        parsed = json.loads(result.output)
1565        assert len(parsed) == 2
1566        assert parsed[0] == {"name": "Alice", "id": "u1"}
1567        assert parsed[1] == {"name": "Bob", "id": "u2"}
1568
1569
1570class TestUIExampleFilterAndTransform:
1571    """Validate the 'Filter & Transform' example from the Code Tool Examples modal."""
1572
1573    @pytest.mark.asyncio
1574    async def test_filter_and_transform_happy_path(self, tmp_path):
1575        search_results = json.dumps(
1576            [
1577                {"title": "Result 1", "url": "https://1.com", "score": 0.9},
1578                {"title": "Result 2", "url": "https://2.com", "score": 0.8},
1579                {"description": "no title or url"},
1580                {"title": "Result 3", "url": "https://3.com", "score": 0.7},
1581            ]
1582        )
1583        fake = FakeTool(
1584            "mcp::remote::test_server::search",
1585            "search",
1586            fn_desc="Search",
1587            params={
1588                "type": "object",
1589                "properties": {"query": {"type": "string"}},
1590                "required": ["query"],
1591            },
1592            result=ToolCallResult(output=search_results),
1593        )
1594        tool = _make_python_code_tool(
1595            tmp_path,
1596            EXAMPLE_FILTER_AND_TRANSFORM,
1597            tool_allowlist=["mcp::remote::test_server::search"],
1598            parameters_schema={
1599                "type": "object",
1600                "properties": {
1601                    "query": {"type": "string"},
1602                    "max_results": {"type": "integer"},
1603                },
1604                "required": ["query"],
1605            },
1606        )
1607        with patch(
1608            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1609            return_value=fake,
1610        ):
1611            result = await tool.run(None, query="test query")
1612        assert not result.is_error, f"Expected success, got: {result.output}"
1613        parsed = json.loads(result.output)
1614        assert len(parsed) == 3
1615        assert parsed[0] == {"title": "Result 1", "url": "https://1.com"}
1616        assert parsed[1] == {"title": "Result 2", "url": "https://2.com"}
1617        assert parsed[2] == {"title": "Result 3", "url": "https://3.com"}
1618
1619    @pytest.mark.asyncio
1620    async def test_filter_and_transform_respects_max_results(self, tmp_path):
1621        search_results = json.dumps(
1622            [{"title": f"R{i}", "url": f"https://{i}.com"} for i in range(20)]
1623        )
1624        fake = FakeTool(
1625            "mcp::remote::test_server::search",
1626            "search",
1627            fn_desc="Search",
1628            params={
1629                "type": "object",
1630                "properties": {"query": {"type": "string"}},
1631                "required": ["query"],
1632            },
1633            result=ToolCallResult(output=search_results),
1634        )
1635        tool = _make_python_code_tool(
1636            tmp_path,
1637            EXAMPLE_FILTER_AND_TRANSFORM,
1638            tool_allowlist=["mcp::remote::test_server::search"],
1639            parameters_schema={
1640                "type": "object",
1641                "properties": {
1642                    "query": {"type": "string"},
1643                    "max_results": {"type": "integer"},
1644                },
1645                "required": ["query"],
1646            },
1647        )
1648        with patch(
1649            "kiln_ai.tools.tool_registry.tool_from_id_and_project",
1650            return_value=fake,
1651        ):
1652            result = await tool.run(None, query="test", max_results=3)
1653        assert not result.is_error, f"Expected success, got: {result.output}"
1654        parsed = json.loads(result.output)
1655        assert len(parsed) == 3