kiln_ai.sandbox.worker

  1"""Child-process entry point for code-tool execution.
  2
  3Stdlib only — no Pydantic / Kiln-model / DB / UI imports.
  4
  5The parent spawns this via ``multiprocessing`` (spawn context).
  6Communication uses two queues (requests child->parent, responses parent->child).
  7"""
  8
  9from __future__ import annotations
 10
 11import io
 12import json
 13import sys
 14import traceback
 15from multiprocessing import Queue
 16from typing import Any
 17
 18from kiln_ai.sandbox.entrypoint import call_entrypoint
 19from kiln_ai.sandbox.tools_api import install_tools_modules
 20
 21_TRUNCATION_LIMIT = 64 * 1024  # 64 KB
 22_TRUNCATION_MARKER = "\n...[truncated]"
 23
 24
 25def child_main(
 26    code: str,
 27    kwargs: dict[str, Any],
 28    requests: Queue,  # type: ignore[type-arg]
 29    responses: Queue,  # type: ignore[type-arg]
 30) -> None:
 31    """Entry point for the code-tool child process.
 32
 33    Puts exactly one ``result`` message on *requests* and then returns.
 34    """
 35    captured_stdout = io.StringIO()
 36    captured_stderr = io.StringIO()
 37    old_stdout = sys.stdout
 38    old_stderr = sys.stderr
 39    try:
 40        sys.stdout = captured_stdout  # type: ignore[assignment]
 41        sys.stderr = captured_stderr  # type: ignore[assignment]
 42
 43        install_tools_modules(requests, responses)
 44
 45        namespace: dict[str, Any] = {}
 46        exec(compile(code, "<code_tool>", "exec"), namespace)
 47
 48        run_fn = namespace.get("run")
 49        if run_fn is None:
 50            _put_error(
 51                requests,
 52                "Code does not define a 'run()' function",
 53                None,
 54                captured_stdout,
 55                captured_stderr,
 56            )
 57            return
 58        if not callable(run_fn):
 59            _put_error(
 60                requests,
 61                "'run' is defined but is not callable",
 62                None,
 63                captured_stdout,
 64                captured_stderr,
 65            )
 66            return
 67
 68        result = call_entrypoint(run_fn, kwargs)
 69        serialized = _serialize_result(result)
 70
 71        requests.put(
 72            {
 73                "type": "result",
 74                "ok": serialized,
 75                "stdout": _truncate(captured_stdout.getvalue()),
 76                "stderr": _truncate(captured_stderr.getvalue()),
 77            }
 78        )
 79    except Exception:
 80        tb = _trim_traceback()
 81        exc_info = sys.exc_info()
 82        _put_error(
 83            requests,
 84            str(exc_info[1]),
 85            tb,
 86            captured_stdout,
 87            captured_stderr,
 88        )
 89    finally:
 90        sys.stdout = old_stdout
 91        sys.stderr = old_stderr
 92
 93
 94# ---------------------------------------------------------------------------
 95# Helpers
 96# ---------------------------------------------------------------------------
 97
 98
 99def _serialize_result(value: Any) -> str:
100    """Serialize the return value of ``run()`` per the spec.
101
102    - ``str`` passes through as-is.
103    - ``dict/list/int/float/bool/None`` -> ``json.dumps``.
104    - Anything else -> raise with a clear message.
105    """
106    if isinstance(value, str):
107        return value
108    if isinstance(value, (dict, list, int, float, bool)) or value is None:
109        try:
110            return json.dumps(value, ensure_ascii=False)
111        except (TypeError, ValueError) as exc:
112            raise TypeError(
113                f"run() returned a value containing non-JSON-serializable data: {exc}"
114            ) from exc
115    raise TypeError(
116        f"run() must return str or JSON-serializable data, got {type(value).__name__}"
117    )
118
119
120def _trim_traceback() -> str:
121    """Format the current exception traceback, keeping only ``<code_tool>`` frames."""
122    _, exc_value, exc_tb = sys.exc_info()
123    if exc_value is None or exc_tb is None:
124        return ""
125
126    entries = traceback.extract_tb(exc_tb)
127    first_user_idx = None
128    for i, frame in enumerate(entries):
129        if frame.filename == "<code_tool>":
130            first_user_idx = i
131            break
132
133    if first_user_idx is not None:
134        entries = entries[first_user_idx:]
135
136    lines = ["Traceback (most recent call last):\n"]
137    lines.extend(traceback.format_list(entries))
138    lines.extend(traceback.format_exception_only(type(exc_value), exc_value))
139    return "".join(lines)
140
141
142def _truncate(text: str) -> str:
143    if len(text) <= _TRUNCATION_LIMIT:
144        return text
145    return text[:_TRUNCATION_LIMIT] + _TRUNCATION_MARKER
146
147
148def _put_error(
149    queue: Queue,  # type: ignore[type-arg]
150    error: str,
151    tb: str | None,
152    stdout: io.StringIO,
153    stderr: io.StringIO,
154) -> None:
155    queue.put(
156        {
157            "type": "result",
158            "error": error,
159            "traceback": tb,
160            "stdout": _truncate(stdout.getvalue()),
161            "stderr": _truncate(stderr.getvalue()),
162        }
163    )