"""(c) Kevin Dunn, 2010-2026. MIT License.
Safe execution wrapper for process-improve tool calls.
Adds the four guard rails needed to expose the tool registry over an
untrusted transport (public MCP server, hosted REST API, etc.):
1. Input-size validation (reject oversize arrays/strings before work).
2. Wall-clock timeout via subprocess isolation.
3. Memory cap per subprocess (POSIX; best-effort on Windows).
4. Structured error types so callers can distinguish failure modes.
The in-process :func:`process_improve.tool_spec.execute_tool_call` is
left untouched for callers that trust their input (notebooks, tests,
the stdio MCP server running on the user's own machine). Hosted
callers should use :func:`safe_execute_tool_call` instead.
Configuration
-------------
Every knob is read lazily from :mod:`process_improve.config.settings`,
which in turn picks up the corresponding ``PROCESS_IMPROVE_*``
environment variable on first access (ENG-09 / ENG-27). Tests can
override a knob in-process via ``settings.tool_timeout = 5.0``; CI
deployments can still set env vars at startup. See
``process_improve/config.py`` for the canonical defaults table.
"""
from __future__ import annotations
import contextlib
import multiprocessing
import os
import sys
import threading
import warnings
from collections.abc import Callable
from concurrent.futures import Executor, Future, ProcessPoolExecutor
from concurrent.futures import TimeoutError as FuturesTimeoutError
from concurrent.futures.process import BrokenProcessPool
from pathlib import Path
from typing import Any
from pydantic import BaseModel, ValidationError
from process_improve._warnings import user_stacklevel
from process_improve.config import settings
def _default_mp_context(
platform: str = sys.platform, free_threaded: bool | None = None
) -> multiprocessing.context.BaseContext | None:
"""Return the multiprocessing context the worker pool uses on ``platform``.
Prefer ``fork`` on Linux: the worker inherits the parent's imported
numpy/registry/etc., which makes startup and tool-dispatch much cheaper
than re-importing on every spawn. macOS is excluded because Apple's
Accelerate framework (used by numpy) is not fork-safe and Python 3.13
emits a DeprecationWarning when a multi-threaded parent forks; Windows
does not support fork at all. On those platforms, and wherever ``fork``
is unavailable, return None: the pool then uses the platform default
(spawn).
A free-threaded build (``free_threaded``, by default whether the GIL is
off) uses ``forkserver`` instead. Its threads run truly in parallel, so a
fork from a thread while another holds a lock leaves that lock held in the
child for ever: concurrent tool calls then deadlock and time out.
``forkserver`` forks each worker from a single-threaded server process.
"""
if not platform.startswith("linux"):
return None
if free_threaded is None:
free_threaded = not getattr(sys, "_is_gil_enabled", lambda: True)()
try:
return multiprocessing.get_context("forkserver" if free_threaded else "fork")
except ValueError:
return None
_DEFAULT_MP_CONTEXT: multiprocessing.context.BaseContext | None = _default_mp_context()
#: True on a thread while it submits a tool call to a worker pool: a ``fork`` pool forks its
#: worker inside ``submit``. The fork hooks below act on those forks only.
_forking_a_worker = threading.local()
def _hold_warnings_lock() -> None:
"""Hold the warnings module's lock while this module forks a worker, so no other thread holds it.
From Python 3.14 every warning, and ``warnings.catch_warnings``, takes one lock for the
whole process. A worker forked while another thread holds that lock inherits it held by a
thread the worker does not have, so the worker's first warning waits forever. A parent
with threads warns at each ``fork()``, so two threads forking at once is enough, and
``execute_tool_call`` records warnings on every call: on the free-threaded build about one
worker in five hung. Python 3.13 and earlier have no such lock.
"""
lock = getattr(warnings, "_lock", None)
if lock is not None and getattr(_forking_a_worker, "active", False):
lock.__enter__()
def _release_warnings_lock() -> None:
"""Release the lock :func:`_hold_warnings_lock` took, in the parent and in the worker.
The worker's one thread is the thread that forked it, so it owns the lock it inherited.
"""
lock = getattr(warnings, "_lock", None)
if lock is not None and getattr(_forking_a_worker, "active", False):
lock.__exit__(None, None, None)
if hasattr(os, "register_at_fork"): # POSIX
os.register_at_fork(
before=_hold_warnings_lock, after_in_parent=_release_warnings_lock, after_in_child=_release_warnings_lock
)
# ---------------------------------------------------------------------------
# Legacy default-name compatibility shims (ENG-22: deprecated since 2.0.0,
# remove in 3.0).
# ---------------------------------------------------------------------------
# Some callers imported these by name (``from process_improve.tool_safety
# import DEFAULT_TIMEOUT_S``). The new home for the actual values is
# ``process_improve.config.settings``. The shims below keep existing imports
# resolving, to whatever the settings object currently reports rather than
# the value at import time. Until 2.0.0 they forwarded silently, so the
# deprecation policy's announce and warn cycle starts with that release.
#: Each legacy module-level name and the :data:`settings` attribute it forwards to.
_LEGACY_TO_SETTINGS: dict[str, str] = {
"DEFAULT_TIMEOUT_S": "tool_timeout",
"DEFAULT_MAX_CELLS": "max_cells",
"DEFAULT_MAX_STRING": "max_string",
"DEFAULT_MAX_DEPTH": "max_depth",
"DEFAULT_MEMORY_MB": "max_memory_mb",
}
def __getattr__(name: str) -> Any: # noqa: ANN401
"""Forward legacy ``DEFAULT_*`` reads to :data:`settings`, with a deprecation warning.
Lets ``from process_improve.tool_safety import DEFAULT_TIMEOUT_S``
keep working, but the value is not frozen at module import: a test
that overrides ``settings.tool_timeout`` sees the new value through
both the new and the legacy name.
.. deprecated:: 2.0.0
The ``DEFAULT_*`` names will be removed in 3.0. Read the
attribute of :data:`process_improve.config.settings` instead.
"""
if name in _LEGACY_TO_SETTINGS:
attr = _LEGACY_TO_SETTINGS[name]
warnings.warn(
f"process_improve.tool_safety.{name} is deprecated since 2.0.0 and will be removed in 3.0; "
f"use process_improve.config.settings.{attr} instead.",
category=DeprecationWarning,
stacklevel=user_stacklevel(),
)
return getattr(settings, attr)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
# Keys whose numeric value scales the cost of the underlying algorithm.
# A malicious caller can otherwise request a huge SVD or a long ESD loop
# with a tiny input array. SEC-19 (#268) extends the historic six keys
# with caps for the design / experiment cost knobs.
_SCALAR_CAPS: dict[str, float] = {
"n_components": 50,
"max_outliers_to_detect": 20,
"n_iter": 10_000,
"max_iter": 10_000,
"n_boot": 1_000,
"n_permutations": 1_000,
"budget": 10_000,
# SEC-19 additions.
"n_steps": 100,
"n_additional_runs": 500,
"n_center_points": 50,
"n_replicates": 50,
"n_factors": 15, # mirrors settings.max_factors_combinatorial
}
# ---------------------------------------------------------------------------
# Structured errors
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Input validation
# ---------------------------------------------------------------------------
def _count_numeric_leaves(value: Any, depth: int, max_depth: int) -> int: # noqa: ANN401
if depth > max_depth:
raise ToolInputTooLargeError(
f"Input nesting depth exceeds limit of {max_depth}",
details={"limit": max_depth},
)
if isinstance(value, (list, tuple)):
return sum(_count_numeric_leaves(v, depth + 1, max_depth) for v in value)
if isinstance(value, dict):
return sum(_count_numeric_leaves(v, depth + 1, max_depth) for v in value.values())
if isinstance(value, (int, float, bool)):
return 1
# Strings, None, and anything else do not count as numeric cells.
return 0
def _check_strings(value: Any, limit: int, depth: int, max_depth: int) -> None: # noqa: ANN401
if depth > max_depth:
# Depth check is already handled by _count_numeric_leaves; keep this
# function simple so it can be reused independently.
raise ToolInputTooLargeError(
f"Input nesting depth exceeds limit of {max_depth}",
details={"limit": max_depth},
)
if isinstance(value, str):
if len(value) > limit:
raise ToolInputTooLargeError(
f"String length {len(value)} exceeds limit of {limit}",
details={"limit": limit, "observed": len(value)},
)
return
if isinstance(value, (list, tuple)):
for v in value:
_check_strings(v, limit, depth + 1, max_depth)
elif isinstance(value, dict):
for v in value.values():
_check_strings(v, limit, depth + 1, max_depth)
# ---------------------------------------------------------------------------
# Pydantic-derived input validation
# ---------------------------------------------------------------------------
#
# Every tool's input is described by a ``BaseModel`` with
# ``ConfigDict(extra="forbid")`` (ENG-04 / ENG-10). Validation lives on the
# model: types, bounds, enums, required fields and "no unknown keys" all come
# from ``model_validate``. ``safe_execute_tool_call`` runs the model
# pre-submission so an oversize or malformed payload is rejected before the
# subprocess pool is even touched.
def _lookup_input_model(tool_name: str) -> type[BaseModel] | None:
"""Return the pydantic input model for *tool_name*, or None if not registered."""
from process_improve.tool_spec import _TOOL_REGISTRY, discover_tools # noqa: PLC0415
discover_tools()
func = _TOOL_REGISTRY.get(tool_name)
if func is None:
return None
return getattr(func, "_input_model", None)
def _validate_against_model(tool_name: str, tool_input: dict[str, Any], model_cls: type[BaseModel]) -> None:
"""Run ``model_cls.model_validate`` and translate ``ValidationError`` -> ``ToolInputInvalidError``."""
try:
model_cls.model_validate(tool_input)
except ValidationError as exc:
raise ToolInputInvalidError(
f"Input to tool {tool_name!r} failed validation: {exc}",
details={"tool_name": tool_name, "errors": exc.errors()},
) from exc
# ---------------------------------------------------------------------------
# Subprocess worker
# ---------------------------------------------------------------------------
def _address_space_in_use() -> int:
"""Return the bytes of address space this process has mapped; 0 where ``/proc`` does not say (macOS)."""
try:
return int(Path("/proc/self/statm").read_text(encoding="ascii").split()[0]) * os.sysconf("SC_PAGE_SIZE")
except (OSError, ValueError):
return 0
def _apply_memory_limit(memory_mb: int) -> None:
"""Allow the current process ``memory_mb`` more of address space than it has mapped (POSIX).
``RLIMIT_AS`` counts address space reserved, not memory in use, and an allocator may
reserve far ahead: the free-threaded build's maps about 1.5 GiB on importing the package
while using about 250 MiB. A cap of ``memory_mb`` alone was then below what the worker
already had, and every new mapping failed, however small.
"""
try:
import resource # noqa: PLC0415 POSIX-only
except ImportError:
return # Windows: no-op
limit_bytes = _address_space_in_use() + memory_mb * 1024 * 1024
# Some sandboxes disallow raising the hard limit; ignore silently.
with contextlib.suppress(ValueError, OSError):
resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
def _pool_initializer(memory_mb: int) -> None:
"""Run inside each worker process before it accepts tasks.
Imports run *before* the memory cap is applied so that what the
package and its dependencies load at import time does not count
against the per-call budget.
"""
# Warm the registry so the first real call doesn't pay discovery cost.
from process_improve.tool_spec import discover_tools # noqa: PLC0415
discover_tools()
_apply_memory_limit(memory_mb)
def _worker_ready() -> None:
"""Do nothing in a worker: the parent waits for this task to know the worker has started."""
def _worker_run(tool_name: str, tool_input: dict[str, Any]) -> Any: # noqa: ANN401
"""Target function executed inside a worker process."""
from process_improve.tool_spec import execute_tool_call # noqa: PLC0415
return execute_tool_call(tool_name, tool_input)
#: Seconds a call's fresh worker may take to start, before the call is refused; not part of the
#: tool's own timeout. A worker started with ``spawn`` (macOS, Windows) imports the package and
#: its tools again, which took 2.1 to 3.5 s on Linux and up to 4.6 s cold on a fast machine
#: (#1130): about 14 s on a machine three times slower. 60 s leaves room for a cold, loaded host
#: and still ends a worker that hangs while it starts.
_WORKER_STARTUP_TIMEOUT_S = 60.0
def _submit(pool: Executor, fn: Callable[..., Any], *args: object) -> Future:
"""Submit ``fn(*args)`` to ``pool``, flagging this thread while a ``fork`` pool forks its worker."""
_forking_a_worker.active = True
try:
return pool.submit(fn, *args)
finally:
_forking_a_worker.active = False
def _start_worker(pool: Executor, tool_name: str, timeout: float) -> None:
"""Start ``pool``'s worker and wait, up to :data:`_WORKER_STARTUP_TIMEOUT_S`, until it is ready.
A pool starts its worker on the first task, and the worker runs ``_pool_initializer``, which
imports the tool registry, before it takes the task. Waiting for a task that does nothing
keeps that start-up out of the tool's ``timeout``, which then measures the tool's own work.
"""
try:
_submit(pool, _worker_ready).result(timeout=_WORKER_STARTUP_TIMEOUT_S)
except FuturesTimeoutError as exc:
raise ToolTimeoutError(
f"Tool {tool_name!r} did not run: its worker process took more than {_WORKER_STARTUP_TIMEOUT_S:g}s "
f"to start. Start-up is timed apart from the tool's {timeout:g}s timeout, which had not begun.",
details={"tool_name": tool_name, "timeout": _WORKER_STARTUP_TIMEOUT_S, "stage": "startup"},
) from exc
# ---------------------------------------------------------------------------
# Pool management
# ---------------------------------------------------------------------------
#: The module-level pool and the memory cap it was built with, or None.
_pool_state: tuple[ProcessPoolExecutor, int] | None = None
#: Guards the module-level ``_pool_state``. The MCP
#: server runs tool calls on executor threads, so pool creation and teardown
#: race without it: two threads could each construct an executor (orphaning
#: one worker forever), and one thread's teardown could SIGKILL a worker
#: still running another thread's task, which then surfaced as a bogus
#: "likely exceeded memory limit" BrokenProcessPool diagnosis.
_pool_lock = threading.Lock()
def _create_pool(memory_mb: int, max_workers: int = 1) -> ProcessPoolExecutor:
"""Construct a new worker pool with the standard initializer/context."""
kwargs: dict[str, Any] = {
"max_workers": max_workers,
"initializer": _pool_initializer,
"initargs": (memory_mb,),
}
if _DEFAULT_MP_CONTEXT is not None:
kwargs["mp_context"] = _DEFAULT_MP_CONTEXT
return ProcessPoolExecutor(**kwargs)
[docs]
def get_pool(memory_mb: int | None = None, max_workers: int = 1) -> ProcessPoolExecutor:
"""Return a lazily-initialised module-level :class:`ProcessPoolExecutor`.
The pool is recreated if ``memory_mb`` changes (e.g. tests override it).
Passing ``None`` resolves the cap from :data:`settings.max_memory_mb` at
call time (ENG-09 / ENG-27). Thread-safe.
"""
if memory_mb is None:
memory_mb = settings.max_memory_mb
global _pool_state # noqa: PLW0603
with _pool_lock:
if _pool_state is None or _pool_state[1] != memory_mb:
_shutdown_pool_locked()
_pool_state = (_create_pool(memory_mb, max_workers=max_workers), memory_mb)
return _pool_state[0]
def _terminate_workers(pool: ProcessPoolExecutor) -> None:
"""Force-kill a pool's worker processes (best-effort).
``ProcessPoolExecutor.shutdown(wait=False, cancel_futures=True)`` does not
interrupt a worker that is *already running* a task: ``cancel_futures`` only
drops not-yet-started futures. So a runaway tool (infinite loop, pathological
input) would keep a CPU core busy after a timeout. Here we reach into the
CPython-internal process table and terminate each worker, escalating to
``kill()`` if ``terminate()`` does not take. Guarded with ``suppress`` so that
an internal-API change degrades to the old (non-killing) behaviour instead of
raising.
"""
processes = getattr(pool, "_processes", None)
if not processes:
return
workers = list(processes.values())
for proc in workers:
with contextlib.suppress(Exception):
if proc.is_alive():
proc.terminate()
for proc in workers:
with contextlib.suppress(Exception):
proc.join(timeout=1.0)
if proc.is_alive():
proc.kill()
def _shutdown_pool_locked() -> None:
"""Tear down the module-level pool. Caller must hold ``_pool_lock``."""
global _pool_state # noqa: PLW0603
if _pool_state is not None:
pool = _pool_state[0]
_terminate_workers(pool)
pool.shutdown(wait=False, cancel_futures=True)
_pool_state = None
[docs]
def shutdown_pool() -> None:
"""Shut down the module-level pool, if any. Safe to call repeatedly.
Worker processes are force-terminated first so a runaway task cannot keep
holding a CPU after the executor is torn down. Thread-safe.
"""
with _pool_lock:
_shutdown_pool_locked()
# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------
