diff --git a/packages/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 4ed2c00e..39f075d8 100644 --- a/packages/client/api-report/launchdarkly_ai_server.txt +++ b/packages/client/api-report/launchdarkly_ai_server.txt @@ -91,6 +91,7 @@ InputTokenDetails [class, @dataclass] cache_read: int = 0 uncached: int = 0 inspect_config(config_key: str, context: dict[str, Any]) -> dict[str, Any] [async] +invoke_tool_callable(fn: ToolCallable, arguments: Any, *, timeout: float = _DEFAULT_TIMEOUT_SECONDS) -> Any [async] is_content_blocks(content: MessageContent) -> bool JudgeResult [class, @dataclass] __init__(usage: UsageDict, response: str, score: float, event_key: str | None = None) @@ -224,6 +225,13 @@ RunUsage [class, @dataclass] add(turn: SpanUsage | None) -> None reported: bool [property] total: SpanUsage = field(default_factory=SpanUsage) +Sandbox [class] + __init__(fn: ToolCallable) + __call__(arguments: Any) -> Awaitable[Any] | Any + __repr__() -> str + __slots__ = ('_fn',) + wrapped: ToolCallable [property] +SandboxError(RuntimeError) [class] SDK_INFO_CONTEXT: LDContext = {'kind': 'ld_ai', 'key': 'ld-internal-tracking', 'anonymous': True} SDK_INFO_EVENT = '$ld:ai:sdk:info' SDKOptions [class, @dataclass(frozen=True)] diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index b170c1cc..52f5c6a9 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -77,6 +77,7 @@ from .skills_fs import _write_skills as _write_skills from .skills_source import config_skills from .skills_watch import SkillWatcher, watch_skills +from .tool_sandbox import Sandbox, SandboxError, invoke_tool_callable from .tracking import execute_and_stream, execute_and_track, wrap_tool_handlers from .types import ( NATIVE_TOOL_KEY, @@ -169,6 +170,9 @@ "HandlerStreamEvent", "InitClientOptions", "SDKOptions", + "Sandbox", + "SandboxError", + "invoke_tool_callable", "JudgeResult", "JudgeRunResult", "JudgeTask", diff --git a/packages/client/src/launchdarkly_ai_server/tool_sandbox.py b/packages/client/src/launchdarkly_ai_server/tool_sandbox.py new file mode 100644 index 00000000..174fc23e --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/tool_sandbox.py @@ -0,0 +1,95 @@ +"""Process-isolated execution for AI Config tool callables. + +The OpenCode handler serves plain Python tools through a loopback MCP bridge. +Those callables run in a spawned child process so a tool cannot read the +handler's in-process state (module globals, open connections, and similar). + +``NativeTool`` entries never pass through this module; OpenCode runs them +inside ``opencode serve``. +""" + +from __future__ import annotations + +import asyncio +import inspect +import multiprocessing +import pickle +from collections.abc import Awaitable, Callable +from concurrent.futures import ProcessPoolExecutor +from typing import Any + +ToolCallable = Callable[[Any], Awaitable[Any] | Any] + +_DEFAULT_TIMEOUT_SECONDS = 600.0 + + +class SandboxError(RuntimeError): + """The sandbox child failed or could not start.""" + + +class Sandbox: + """Runs *fn* in a spawned process when invoked. + + Authors may wrap a tool explicitly. The OpenCode loopback bridge also + wraps bridged callables before registration so every served Python tool + is isolated even when the author passes a plain function. + """ + + __slots__ = ("_fn",) + + def __init__(self, fn: ToolCallable) -> None: + self._fn = fn + + @property + def wrapped(self) -> ToolCallable: + """The inner callable (before tracking wrappers).""" + return self._fn + + def __call__(self, arguments: Any) -> Awaitable[Any] | Any: + return invoke_tool_callable(self._fn, arguments) + + def __repr__(self) -> str: + return f"Sandbox({self._fn!r})" + + +def _execute_pickled(payload: bytes) -> Any: + fn, arguments = pickle.loads(payload) + result = fn(arguments) + if inspect.iscoroutine(result): + return asyncio.run(result) + return result + + +def _spawn_sync(fn: ToolCallable, arguments: Any, *, timeout: float) -> Any: + try: + payload = pickle.dumps((fn, arguments)) + except (pickle.PicklingError, AttributeError) as exc: + raise SandboxError( + "Tool callable is not picklable; define it at module scope " + "instead of as a closure or lambda." + ) from exc + ctx = multiprocessing.get_context("spawn") + with ProcessPoolExecutor(max_workers=1, mp_context=ctx) as pool: + future = pool.submit(_execute_pickled, payload) + try: + return future.result(timeout=timeout) + except TimeoutError as exc: + raise SandboxError("Tool callable timed out in the sandbox") from exc + except Exception as exc: + if isinstance(exc, SandboxError): + raise + raise SandboxError(str(exc)) from exc + + +async def invoke_tool_callable( + fn: ToolCallable, + arguments: Any, + *, + timeout: float = _DEFAULT_TIMEOUT_SECONDS, +) -> Any: + """Run *fn* with *arguments* in a spawned child process.""" + loop = asyncio.get_running_loop() + return await asyncio.wait_for( + loop.run_in_executor(None, lambda: _spawn_sync(fn, arguments, timeout=timeout)), + timeout=timeout + 1.0, + ) diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index acc4f8b5..1518389b 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -9,6 +9,7 @@ from typing import Any from .agui import HANDLER_RICH_TYPES, TEXT_MESSAGE_CONTENT, TEXT_TYPES +from .tool_sandbox import Sandbox from .trajectory import HANDOFF_TOOL_PREFIX, TrajectoryRecorder, render_trajectory from .types import ( NATIVE_TOOL_KEY, @@ -94,6 +95,32 @@ def stub(*args: Any, **kwargs: Any) -> None: return stub wrapped[name] = _make_native_stub(name, native) + elif isinstance(fn, Sandbox): + sandbox = fn + + def _make_sandbox_wrapper( + tool_name: str, boxed: Sandbox + ) -> Callable[..., Any]: + def track_call() -> None: + get_client().track( + "$ld:ai:tool_call", + user_context, + {**track_data, "toolKey": tool_name}, + 1, + ) + + async def wrapper(*args: Any, **kwargs: Any) -> Any: + from .tool_sandbox import invoke_tool_callable + + track_call() + arguments = args[0] if args else kwargs + return await invoke_tool_callable(boxed.wrapped, arguments) + + setattr(wrapper, "_ld_tool_target", boxed.wrapped) + setattr(wrapper, "_ld_tool_track", track_call) + return wrapper + + wrapped[name] = _make_sandbox_wrapper(name, sandbox) else: def _make_regular_wrapper( @@ -109,16 +136,21 @@ def handoff_wrapper(*args: Any, **kwargs: Any) -> Any: return handoff_wrapper - async def wrapper(*args: Any, **kwargs: Any) -> Any: + def track_call() -> None: get_client().track( "$ld:ai:tool_call", user_context, {**track_data, "toolKey": tool_name}, 1, ) + + async def wrapper(*args: Any, **kwargs: Any) -> Any: + track_call() result = original(*args, **kwargs) return await result if inspect.isawaitable(result) else result + setattr(wrapper, "_ld_tool_target", original) + setattr(wrapper, "_ld_tool_track", track_call) return wrapper wrapped[name] = _make_regular_wrapper(name, fn) diff --git a/packages/client/tests/test_public_surface.py b/packages/client/tests/test_public_surface.py index bf106ef8..522356e0 100644 --- a/packages/client/tests/test_public_surface.py +++ b/packages/client/tests/test_public_surface.py @@ -45,6 +45,8 @@ "ProviderResponse", "Registry", "RunUsage", + "Sandbox", + "SandboxError", "SDK_INFO_CONTEXT", "SDK_INFO_EVENT", "SpanMessage", @@ -80,6 +82,7 @@ "image_block_to_url", "init_client", "inspect_config", + "invoke_tool_callable", "is_content_blocks", "lang_chain_content_text", "lang_chain_finish_reasons", diff --git a/packages/client/tests/test_tool_sandbox.py b/packages/client/tests/test_tool_sandbox.py new file mode 100644 index 00000000..fe6e3e75 --- /dev/null +++ b/packages/client/tests/test_tool_sandbox.py @@ -0,0 +1,43 @@ +"""Process sandbox for AI Config tool callables.""" + +from __future__ import annotations + +import pytest + +from launchdarkly_ai_server import Sandbox, SandboxError, invoke_tool_callable + +_RUNTIME_SECRET: str | None = None + + +def _runtime_secret_probe(_args: dict[str, object]) -> str: + return _RUNTIME_SECRET or "missing" + + +async def test_invoke_tool_callable_does_not_see_parent_globals() -> None: + global _RUNTIME_SECRET + _RUNTIME_SECRET = "parent-only" + result = await invoke_tool_callable(_runtime_secret_probe, {}) + assert result == "missing" + + +def _sync_echo(args: dict[str, object]) -> str: + return str(args.get("city")) + + +async def _async_echo(args: dict[str, object]) -> dict[str, object]: + return args + + +async def test_invoke_tool_callable_runs_sync_callables() -> None: + result = await invoke_tool_callable(_sync_echo, {"city": "Oslo"}) + assert result == "Oslo" + + +async def test_sandbox_delegates_to_inner() -> None: + boxed = Sandbox(_async_echo) + assert await boxed({"x": 1}) == {"x": 1} + + +async def test_non_picklable_callable_raises_sandbox_error() -> None: + with pytest.raises(SandboxError, match="picklable"): + await invoke_tool_callable(lambda _args: "nope", {}) diff --git a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py index bcd94a2d..a2364452 100644 --- a/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py +++ b/packages/opencode-agents/src/launchdarkly_ai_opencode_agents/tool_bridge.py @@ -322,13 +322,27 @@ def _call_tool( } +def _sandbox_target(fn: ToolCallable) -> ToolCallable: + from launchdarkly_ai_server.tool_sandbox import Sandbox + + if isinstance(fn, Sandbox): + return fn.wrapped + return fn + + async def _invoke(fn: ToolCallable, arguments: Any) -> tuple[str, bool]: from launchdarkly_ai_server import tool_result_string + from launchdarkly_ai_server.tool_sandbox import invoke_tool_callable try: - result = fn(arguments) - if asyncio.iscoroutine(result): - result = await result + target = getattr(fn, "_ld_tool_target", None) + if target is not None: + track = getattr(fn, "_ld_tool_track", None) + if track is not None: + track() + result = await invoke_tool_callable(_sandbox_target(target), arguments) + else: + result = await invoke_tool_callable(fn, arguments) return tool_result_string(result), False except Exception as exc: return str(exc), True diff --git a/packages/opencode-agents/tests/test_tool_bridge.py b/packages/opencode-agents/tests/test_tool_bridge.py index 236eedb4..79b53267 100644 --- a/packages/opencode-agents/tests/test_tool_bridge.py +++ b/packages/opencode-agents/tests/test_tool_bridge.py @@ -21,6 +21,24 @@ sanitize, ) +_BRIDGE_RUNTIME_SECRET: str | None = None + + +def _bridge_runtime_probe(_args: dict[str, object]) -> str: + return _BRIDGE_RUNTIME_SECRET or "missing" + + +def _weather_sync(args: dict[str, Any]) -> str: + return f"sunny in {args['city']}" + + +async def _echo_async(args: dict[str, Any]) -> dict[str, Any]: + return {"ok": True, "echo": args} + + +def _weather_boom(_args: dict[str, Any]) -> str: + raise ValueError("no such city") + def test_sanitize_matches_opencodes_own_regex() -> None: # mcp/catalog.ts: `value.replace(/[^a-zA-Z0-9_-]/g, "_")`. @@ -107,13 +125,41 @@ async def test_tools_list_reports_the_registered_schema(self) -> None: finally: await bridge.stop() + async def test_tools_call_runs_callables_in_a_sandbox(self) -> None: + global _BRIDGE_RUNTIME_SECRET + _BRIDGE_RUNTIME_SECRET = "parent-only" + bridge = ToolBridge() + bridge.register( + "probe", + description="", + parameters={}, + fn=_bridge_runtime_probe, + ) + await bridge.start() + try: + response = await self._post( + bridge.base_url, + bridge.token, + { + "jsonrpc": "2.0", + "id": 31, + "method": "tools/call", + "params": {"name": "probe", "arguments": {}}, + }, + ) + result = response.json()["result"] + assert result["isError"] is False + assert result["content"][0]["text"] == "missing" + finally: + await bridge.stop() + async def test_tools_call_runs_a_sync_handler_and_wraps_the_result(self) -> None: bridge = ToolBridge() bridge.register( "get_weather", description="", parameters={}, - fn=lambda args: f"sunny in {args['city']}", + fn=_weather_sync, ) await bridge.start() try: @@ -134,11 +180,8 @@ async def test_tools_call_runs_a_sync_handler_and_wraps_the_result(self) -> None await bridge.stop() async def test_tools_call_awaits_an_async_handler(self) -> None: - async def handler(args: dict[str, Any]) -> dict[str, Any]: - return {"ok": True, "echo": args} - bridge = ToolBridge() - bridge.register("echo", description="", parameters={}, fn=handler) + bridge.register("echo", description="", parameters={}, fn=_echo_async) await bridge.start() try: response = await self._post( @@ -162,11 +205,8 @@ async def handler(args: dict[str, Any]) -> dict[str, Any]: await bridge.stop() async def test_a_raising_handler_comes_back_as_an_mcp_tool_error(self) -> None: - def boom(args: dict[str, Any]) -> str: - raise ValueError("no such city") - bridge = ToolBridge() - bridge.register("get_weather", description="", parameters={}, fn=boom) + bridge.register("get_weather", description="", parameters={}, fn=_weather_boom) await bridge.start() try: response = await self._post( diff --git a/packages/opencode-agents/tests/test_tool_contract.py b/packages/opencode-agents/tests/test_tool_contract.py index 7aec305a..ebf752fb 100644 --- a/packages/opencode-agents/tests/test_tool_contract.py +++ b/packages/opencode-agents/tests/test_tool_contract.py @@ -48,6 +48,11 @@ from launchdarkly_ai_opencode_agents.server import OpenCodeHTTPError from launchdarkly_ai_server import NativeTool, wrap_tool_handlers + +async def _contract_get_weather(args: dict[str, Any]) -> str: + return f"sunny in {args['city']}" + + CONTEXT = {"kind": "user", "key": "user-1"} _TRACK_DATA: dict[str, Any] = { "runId": "run-1", @@ -576,9 +581,6 @@ class TestToolBridgeEndToEnd: async def test_the_registered_bridge_answers_a_real_tools_call( self, mock_ld_client: MagicMock ) -> None: - async def get_weather(args: dict[str, Any]) -> str: - return f"sunny in {args['city']}" - handler, script = playing( [ assistant_open("m1"), @@ -587,7 +589,10 @@ async def get_weather(args: dict[str, Any]) -> str: ) agen = ( await handler.stream( - _CONFIG_WITH_TOOLS, "hi", _wrapped({"get_weather": get_weather}), None + _CONFIG_WITH_TOOLS, + "hi", + _wrapped({"get_weather": _contract_get_weather}), + None, ) ).__aiter__() first = await agen.__anext__() diff --git a/uv.lock b/uv.lock index ec6a31c7..6407d96f 100644 --- a/uv.lock +++ b/uv.lock @@ -1322,7 +1322,7 @@ wheels = [ [[package]] name = "launchdarkly-ai-claude-agents" -version = "0.2.6" +version = "0.3.0" source = { editable = "packages/claude-agents" } dependencies = [ { name = "anthropic" }, @@ -1341,7 +1341,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-claude-messages" -version = "0.2.5" +version = "0.3.0" source = { editable = "packages/claude-messages" } dependencies = [ { name = "anthropic" }, @@ -1358,7 +1358,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-langchain-agents" -version = "0.2.6" +version = "0.3.0" source = { editable = "packages/langchain-agents" } dependencies = [ { name = "langchain-core" }, @@ -1377,7 +1377,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-langchain-messages" -version = "0.2.5" +version = "0.3.0" source = { editable = "packages/langchain-messages" } dependencies = [ { name = "langchain-core" }, @@ -1394,7 +1394,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-litellm-agents" -version = "0.1.1" +version = "0.2.0" source = { editable = "packages/litellm-agents" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1411,7 +1411,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-litellm-messages" -version = "0.1.1" +version = "0.2.0" source = { editable = "packages/litellm-messages" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1428,7 +1428,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-openai-agents" -version = "0.2.6" +version = "0.3.0" source = { editable = "packages/openai-agents" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1447,7 +1447,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-openai-messages" -version = "0.2.5" +version = "0.3.0" source = { editable = "packages/openai-messages" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1464,7 +1464,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-opencode-agents" -version = "0.2.0" +version = "0.3.0" source = { editable = "packages/opencode-agents" } dependencies = [ { name = "httpx" }, @@ -1481,7 +1481,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-python" -version = "0.3.0" +version = "0.4.0" source = { editable = "packages/ai" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1501,7 +1501,7 @@ provides-extras = ["otel"] [[package]] name = "launchdarkly-ai-server" -version = "0.4.0" +version = "0.5.0" source = { editable = "packages/client" } dependencies = [ { name = "opentelemetry-api" }, @@ -1523,7 +1523,7 @@ provides-extras = ["otel"] [[package]] name = "launchdarkly-ai-typesafe" -version = "0.1.1" +version = "0.2.0" source = { editable = "packages/typesafe" } dependencies = [ { name = "launchdarkly-ai-server" }, @@ -1540,7 +1540,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-vercel-agents" -version = "0.1.1" +version = "0.2.0" source = { editable = "packages/vercel-agents" } dependencies = [ { name = "ai" }, @@ -1557,7 +1557,7 @@ requires-dist = [ [[package]] name = "launchdarkly-ai-vercel-messages" -version = "0.1.1" +version = "0.2.0" source = { editable = "packages/vercel-messages" } dependencies = [ { name = "ai" },