kiln_ai.sandbox.tools_api

  1"""Sandbox-side ``kiln.tools`` bridge for code-tool children.
  2
  3Stdlib only — no Pydantic / Kiln-model / DB / UI imports.
  4
  5The synthetic-module *surface* (proxy behavior, ``list_tools`` wiring, and the
  6typed exception classes ``ToolNotAllowed`` / ``ToolTimeout`` / ``ToolCallError``)
  7lives in :mod:`kiln_ai.sandbox.tools_surface`, shared with the test shim so
  8runtime and tests present one definition of ``kiln.tools``. This module provides
  9the sandbox's real IPC bridge and wires it into that surface.
 10
 11Provides:
 12
 13* ``ToolCallBridge``: thread-safe IPC bridge (child -> parent tool calls)
 14* ``install_tools_modules()``: injects ``kiln``, ``kiln.tools``,
 15  ``kiln.async_tools`` into ``sys.modules`` so user code can
 16  ``from kiln import tools`` etc.
 17
 18``ToolNotAllowed`` / ``ToolTimeout`` / ``ToolCallError`` are re-exported here
 19(they are defined in ``tools_surface``) so existing importers keep working.
 20"""
 21
 22from __future__ import annotations
 23
 24import json
 25import threading
 26from multiprocessing import Queue
 27from typing import Any
 28
 29from kiln_ai.sandbox.tools_surface import (
 30    ToolCallError,
 31    ToolNotAllowed,
 32    ToolTimeout,
 33    install_tools_modules_for_bridge,
 34)
 35
 36__all__ = [
 37    "ToolCallBridge",
 38    "ToolCallError",
 39    "ToolNotAllowed",
 40    "ToolTimeout",
 41    "install_tools_modules",
 42]
 43
 44
 45# ---------------------------------------------------------------------------
 46# Child-side IPC bridge
 47# ---------------------------------------------------------------------------
 48
 49
 50class ToolCallBridge:
 51    """Thread-safe bridge between user code and the parent message pump.
 52
 53    One bridge per child process.  ``call()`` and ``list_tools()`` may be
 54    invoked from any user thread (or from ``asyncio.to_thread`` for the
 55    async mirror).
 56    """
 57
 58    def __init__(self, requests: Queue, responses: Queue) -> None:  # type: ignore[type-arg]
 59        self._requests = requests
 60        self._responses = responses
 61        self._lock = threading.Lock()
 62        self._next_id = 0
 63        self._pending: dict[int, _PendingCall] = {}
 64
 65    # -- public API --
 66
 67    def call(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
 68        """Issue a ``tool_call`` to the parent; block until reply.
 69
 70        Returns the raw output string on success.  Raises one of the
 71        typed exceptions on failure.
 72        """
 73        try:
 74            safe_kwargs = json.loads(json.dumps(kwargs, ensure_ascii=False))
 75        except (TypeError, ValueError) as exc:
 76            raise ToolCallError(
 77                tool=name,
 78                message=f"tool arguments must be JSON-serializable: {exc}",
 79            ) from exc
 80
 81        # Serialize positional args so the parent can produce a helpful error
 82        try:
 83            safe_args = json.loads(json.dumps(list(args), ensure_ascii=False))
 84        except (TypeError, ValueError):
 85            safe_args = [repr(a) for a in args]
 86
 87        call_id, pending = self._allocate()
 88        msg: dict[str, Any] = {
 89            "type": "tool_call",
 90            "call_id": call_id,
 91            "tool_name": name,
 92            "arguments": safe_kwargs,
 93        }
 94        if safe_args:
 95            msg["positional_args"] = safe_args
 96        self._requests.put(msg)
 97        pending.event.wait()
 98        return self._resolve(call_id, name, pending)
 99
100    def list_tools(self) -> list[dict[str, Any]]:
101        """Ask the parent for the allowlisted tool definitions."""
102        call_id, pending = self._allocate()
103        self._requests.put({"type": "list_tools", "call_id": call_id})
104        pending.event.wait()
105        msg = pending.result
106        assert msg is not None
107        if "ok_list" in msg:
108            return msg["ok_list"]  # type: ignore[return-value]
109        err = msg.get("error", {})
110        raise ToolCallError(
111            tool="list_tools",
112            message=err.get("message", "unknown error"),
113        )
114
115    def start_dispatcher(self) -> None:
116        """Start the daemon thread that reads parent replies."""
117        t = threading.Thread(target=self._dispatch_loop, daemon=True)
118        t.start()
119
120    # -- internals --
121
122    def _allocate(self) -> tuple[int, "_PendingCall"]:
123        with self._lock:
124            cid = self._next_id
125            self._next_id += 1
126            p = _PendingCall()
127            self._pending[cid] = p
128            return cid, p
129
130    def _dispatch_loop(self) -> None:
131        try:
132            while True:
133                try:
134                    msg = self._responses.get()
135                except (EOFError, OSError, ValueError):
136                    break
137                cid = msg.get("call_id")
138                with self._lock:
139                    pending = self._pending.pop(cid, None)
140                if pending is not None:
141                    pending.result = msg
142                    pending.event.set()
143        finally:
144            with self._lock:
145                orphans = list(self._pending.values())
146                self._pending.clear()
147            for p in orphans:
148                p.result = {
149                    "error": {
150                        "kind": "call_error",
151                        "message": "Parent process disconnected or queue closed.",
152                    }
153                }
154                p.event.set()
155
156    def _resolve(self, call_id: int, name: str, pending: "_PendingCall") -> str:
157        msg = pending.result
158        assert msg is not None
159        if "ok" in msg:
160            return msg["ok"]
161        err = msg.get("error", {})
162        kind = err.get("kind", "call_error")
163        message = err.get("message", "unknown error")
164        raw = err.get("raw")
165        if kind == "not_allowed":
166            raise ToolNotAllowed(tool=name, message=message)
167        if kind == "timeout":
168            raise ToolTimeout(tool=name, message=message)
169        raise ToolCallError(tool=name, message=message, raw=raw)
170
171
172class _PendingCall:
173    __slots__ = ("event", "result")
174
175    def __init__(self) -> None:
176        self.event = threading.Event()
177        self.result: dict[str, Any] | None = None
178
179
180# ---------------------------------------------------------------------------
181# Module installation
182# ---------------------------------------------------------------------------
183
184
185def install_tools_modules(
186    requests: Queue,
187    responses: Queue,  # type: ignore[type-arg]
188) -> ToolCallBridge:
189    """Create and install the synthetic ``kiln`` package into ``sys.modules``.
190
191    Returns the bridge so the caller can inspect it if needed.
192    """
193    bridge = ToolCallBridge(requests, responses)
194    bridge.start_dispatcher()
195    install_tools_modules_for_bridge(bridge)
196    return bridge