aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/ptc/executor.py

@@ -0,0 +1,627 @@
+"""PTC Executor implementations.
+
+This module provides the sandboxed executor for Programmatic Tool Calling.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import inspect
+import json
+import shlex
+from dataclasses import dataclass, field
+from pathlib import PurePosixPath
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from aip_agents.mcp.client.base_mcp_client import BaseMCPClient
+from aip_agents.ptc.custom_tools import PTCCustomToolConfig
+from aip_agents.ptc.exceptions import PTCToolError
+from aip_agents.ptc.prompt_builder import PromptConfig
+from aip_agents.sandbox.defaults import DEFAULT_PTC_PACKAGES, DEFAULT_PTC_TEMPLATE
+from aip_agents.utils.logger import get_logger
+
+# Lazy import to avoid circular dependencies
+# These are only needed for PTCSandboxExecutor
+try:
+    from aip_agents.ptc.sandbox_bridge import build_sandbox_payload, wrap_ptc_code
+    from aip_agents.sandbox.types import SandboxExecutionResult
+
+    _SANDBOX_DEPS_AVAILABLE = True
+except ImportError:
+    _SANDBOX_DEPS_AVAILABLE = False
+
+logger = get_logger(__name__)
+
+STRUCTURED_OUTPUT_ENV = "AIP_AGENTS_PTC_STRUCTURED_OUTPUT_PATH"
+STRUCTURED_OUTPUT_DIR = "/home/user/.aip_agents_ptc/structured_outputs"
+MAX_STRUCTURED_OUTPUT_BYTES = 10 * 1024 * 1024
+MAX_STRUCTURED_OUTPUT_TOTAL_BYTES = 10 * MAX_STRUCTURED_OUTPUT_BYTES
+
+if TYPE_CHECKING:
+    from aip_agents.middleware.backends.sandbox import SandboxBackend
+    from aip_agents.ptc.payload import SandboxPayload
+    from aip_agents.sandbox.types import PTCStructuredOutput
+
+
+@dataclass
+class PTCSandboxConfig:
+    """Configuration for PTC sandbox executor.
+
+    Attributes:
+        enabled: Whether PTC is enabled. When False, PTC is disabled.
+        default_tool_timeout: Optional default timeout per nested MCP tool call in seconds.
+            When None, PTC does not inject a default nested tool-call timeout.
+        sandbox_template: Optional E2B sandbox template ID. Ignored for runtime
+            selection when filesystem_backend is set; the backend's template is
+            the source of truth for the shared sandbox runtime.
+        sandbox_timeout: Sandbox lease timeout in seconds (VM TTL). This is not
+            the same clock as a single code execution budget.
+        execution_timeout: Per-run PTC execution timeout in seconds. When unset,
+            it inherits sandbox_timeout for backward compatibility.
+        package_install_timeout: Timeout in seconds for sandbox package installation.
+            When None, the runtime default applies.
+        auto_build_template: Whether to automatically build the sandbox template if it's missing.
+        ptc_packages: List of packages to install in sandbox. Defaults to DEFAULT_PTC_PACKAGES.
+        prompt: Prompt configuration for PTC usage guidance.
+        custom_tools: Configuration for custom LangChain tools in sandbox.
+        filesystem_backend: Optional caller-provided filesystem backend for shared sandbox execution.
+            When set, PTC configures packages on that shared backend runtime on
+            first execution and leaves runtime cleanup to the backend owner.
+    """
+
+    enabled: bool = False
+    default_tool_timeout: float | None = None
+    sandbox_template: str | None = DEFAULT_PTC_TEMPLATE
+    sandbox_timeout: float = 3600.0
+    execution_timeout: float | None = None
+    package_install_timeout: float | None = None
+    auto_build_template: bool = False
+    ptc_packages: list[str] | None = field(default_factory=lambda: list(DEFAULT_PTC_PACKAGES))
+    prompt: PromptConfig = field(default_factory=PromptConfig)
+    custom_tools: PTCCustomToolConfig = field(default_factory=PTCCustomToolConfig)
+    filesystem_backend: SandboxBackend | None = None
+
+    def __post_init__(self) -> None:
+        """Validate timeout relationships after dataclass initialization."""
+        if self.execution_timeout is None:
+            return
+        if self.execution_timeout <= 0:
+            msg = "PTCSandboxConfig.execution_timeout must be > 0."
+            raise ValueError(msg)
+        if self.execution_timeout > self.sandbox_timeout:
+            msg = "PTCSandboxConfig.execution_timeout must be <= sandbox_timeout."
+            raise ValueError(msg)
+
+    def effective_execution_timeout(self) -> float:
+        """Return the per-run execution timeout, inheriting sandbox_timeout when unset."""
+        timeout = self.sandbox_timeout if self.execution_timeout is None else self.execution_timeout
+        return float(timeout)
+
+
+class PTCSandboxExecutor:
+    r"""Executes PTC code inside an E2B sandbox.
+
+    This executor is used for LLM-generated code that requires sandboxing.
+    It builds a sandbox payload (MCP server config + generated tool modules)
+    and executes the code using the E2B runtime.
+
+    Static bundle caching:
+        The executor tracks whether the static bundle (wrappers, registry, sources)
+        has been uploaded. On the first run, it uploads the full payload. On subsequent
+        runs, it only uploads per-run files (e.g., tools/custom_defaults.json) to
+        reduce upload overhead.
+
+        If the sandbox is destroyed/recreated, call reset_static_bundle() to force
+        a full re-upload on the next execution.
+
+    Thread-safety:
+        The static bundle upload is guarded by an asyncio.Lock. Concurrent calls will
+        serialize while the initial upload completes. Per-run diff/execute/cache-update
+        is guarded by an executor-local asyncio.Lock to keep dedupe state consistent.
+
+    Example:
+        runtime = E2BSandboxRuntime()
+        executor = PTCSandboxExecutor(mcp_client, runtime)
+        result = await executor.execute_code("from tools.yfinance import get_stock\\nprint(get_stock('AAPL'))")
+    """
+
+    def __init__(
+        self,
+        mcp_client: BaseMCPClient | None,
+        runtime: Any,
+        config: PTCSandboxConfig | None = None,
+        owns_runtime: bool = True,
+    ) -> None:
+        """Initialize PTCSandboxExecutor.
+
+        Args:
+            mcp_client: The MCP client with configured servers. Can be None for custom-only configs.
+            runtime: The E2B sandbox runtime instance.
+            config: Optional sandbox executor configuration.
+            owns_runtime: Whether this executor owns and may clean up the runtime.
+
+        Raises:
+            ImportError: If sandbox dependencies are not available.
+        """
+        if not _SANDBOX_DEPS_AVAILABLE:
+            raise ImportError(
+                "Sandbox dependencies not available. "
+                "PTCSandboxExecutor requires sandbox_bridge and e2b_runtime modules."
+            )
+
+        self._mcp_client = mcp_client
+        self._runtime = runtime
+        self._config = config or PTCSandboxConfig()
+        self._owns_runtime = owns_runtime
+        self._static_bundle_uploaded = False
+        self._per_run_hash_cache: dict[str, str] = {}
+        self._bundle_lock = asyncio.Lock()
+        self._per_run_lock = asyncio.Lock()
+        self._cleanup_timeout_s = 5.0
+        self._entrypoint_file_path = f".aip_ptc_entrypoint_{id(self):x}.py"
+
+    async def _refresh_runtime_state(self) -> None:
+        """Refresh shared runtimes and reset upload caches when they change."""
+        refresh_runtime = getattr(self._runtime, "refresh_runtime", None)
+        if not callable(refresh_runtime):
+            return
+        runtime_changed = refresh_runtime()
+        if asyncio.iscoroutine(runtime_changed):
+            runtime_changed = await runtime_changed
+        if runtime_changed:
+            logger.info("PTC runtime changed, resetting static bundle state")
+            self._reset_cached_upload_state()
+
+    def _clear_upload_caches(self) -> None:
+        """Clear static and per-run upload caches."""
+        self._static_bundle_uploaded = False
+        self._per_run_hash_cache.clear()
+
+    def _reset_cached_upload_state(self) -> None:
+        """Reset static/per-run upload cache state."""
+        self._clear_upload_caches()
+
+    def _reset_bundle_state_if_inactive(self) -> None:
+        """Reset cached bundle state when the runtime is inactive."""
+        if self._runtime.is_active:
+            return
+        if not self._static_bundle_uploaded and not self._per_run_hash_cache:
+            return
+
+        logger.info("Runtime is inactive, resetting static bundle state")
+        self._reset_cached_upload_state()
+
+    def _build_per_run_hashes(self, payload: SandboxPayload) -> dict[str, str]:
+        """Build deterministic hash map for per-run files."""
+        return {
+            file_path: hashlib.sha256(content.encode("utf-8")).hexdigest()
+            for file_path, content in payload.per_run_files.items()
+        }
+
+    async def _cleanup_runtime_for_removed_paths(self, removed_paths: set[str]) -> None:
+        """Cleanup runtime when per-run path shrink is detected.
+
+        Raises:
+            RuntimeError: If cleanup cannot be confirmed as successful.
+        """
+        logger.info(f"Per-run path shrink detected ({sorted(removed_paths)}), running conservative cleanup")
+        if not self._owns_runtime:
+            cleanup_removed_paths = getattr(self._runtime, "cleanup_removed_paths", None)
+            if not callable(cleanup_removed_paths):
+                raise RuntimeError("Caller-owned runtime cannot remove stale PTC files")
+            cleanup_result = cleanup_removed_paths(removed_paths)
+            if asyncio.iscoroutine(cleanup_result):
+                await cleanup_result
+            self._reset_cached_upload_state()
+            return
+        try:
+            cleanup_result = await asyncio.wait_for(
+                self._runtime.cleanup(),
+                timeout=self._cleanup_timeout_s,
+            )
+        except TimeoutError as exc:
+            raise RuntimeError("Runtime cleanup timed out") from exc
+        if isinstance(cleanup_result, bool) and not cleanup_result:
+            raise RuntimeError("Runtime cleanup reported failure")
+        self._reset_cached_upload_state()
+
+    def _should_include_packages_path(self) -> bool:
+        """Check if the packages path should be added to sys.path."""
+        custom_tools = self._config.custom_tools
+        if not custom_tools.enabled or not custom_tools.tools:
+            return False
+
+        return any(tool.get("package_path") or tool.get("kind") == "file" for tool in custom_tools.tools)
+
+    async def _build_payload(self, tool_configs: dict[str, dict] | None) -> SandboxPayload:
+        """Build the sandbox payload for execution."""
+        logger.info("Building sandbox payload")
+        return await build_sandbox_payload(
+            self._mcp_client,
+            self._config.default_tool_timeout,
+            custom_tools_config=self._config.custom_tools,
+            tool_configs=tool_configs,
+        )
+
+    def _wrap_code(self, code: str) -> str:
+        """Wrap code with required imports and setup."""
+        logger.info("Wrapping PTC code")
+        return wrap_ptc_code(code, include_packages_path=self._should_include_packages_path())
+
+    def _inject_structured_output_env(self, wrapped_code: str, path: str) -> str:
+        """Set the structured-output path inside the sandbox process.
+
+        Some E2B/gllm transport combinations have dropped per-call env values
+        under pytest-asyncio session-loop reuse. The host still passes env to
+        runtimes that honor it, but also sets this internal variable in the
+        executed Python process so the custom-tool helper sees the same path.
+        """
+        return (
+            "# PTC structured output bridge - host generated\n"
+            "import os as _ptc_structured_output_os\n"
+            f"_ptc_structured_output_os.environ[{json.dumps(STRUCTURED_OUTPUT_ENV)}] = {json.dumps(path)}\n"
+            + wrapped_code
+        )
+
+    def _execution_timeout(self) -> float:
+        """Return the per-run execution timeout for PTC code."""
+        return self._config.effective_execution_timeout()
+
+    def _entrypoint_path(self, wrapped_code: str) -> str:
+        """Build the executor-specific entrypoint path for command execution."""
+        _ = wrapped_code
+        return self._entrypoint_file_path
+
+    def _runtime_prefers_python_file_execution(self) -> bool:
+        return getattr(self._runtime, "prefer_python_file_execution", False) is True and callable(
+            getattr(self._runtime, "execute_python_file", None)
+        )
+
+    async def _execute_runtime(
+        self,
+        *,
+        wrapped_code: str,
+        entrypoint_path: str,
+        files: dict[str, str] | None,
+        env: dict[str, str],
+    ) -> SandboxExecutionResult:
+        execution_timeout = self._execution_timeout()
+        if self._runtime_prefers_python_file_execution():
+            return await self._runtime.execute_python_file(
+                code=wrapped_code,
+                entrypoint_path=entrypoint_path,
+                timeout=execution_timeout,
+                files=files,
+                env=env,
+                template=self._config.sandbox_template,
+            )
+        return await self._runtime.execute(
+            code=wrapped_code,
+            timeout=execution_timeout,
+            files=files,
+            env=env,
+            template=self._config.sandbox_template,
+        )
+
+    async def _execute_payload(
+        self,
+        payload: SandboxPayload,
+        wrapped_code: str,
+        *,
+        entrypoint_path: str | None = None,
+        upload_static_bundle: bool,
+    ) -> SandboxExecutionResult:
+        """Execute the wrapped code with the provided payload."""
+        entrypoint_path = entrypoint_path or self._entrypoint_path(wrapped_code)
+        if upload_static_bundle:
+            logger.info("Uploading static bundle and per-run files")
+            files_to_upload = {**payload.files, **payload.per_run_files}
+        else:
+            logger.debug("Static bundle already uploaded, uploading only per-run files")
+            files_to_upload = dict(payload.per_run_files)
+        if self._runtime_prefers_python_file_execution():
+            files_to_upload[entrypoint_path] = wrapped_code
+
+        logger.info(
+            "Executing code in sandbox "
+            f"(sandbox_timeout: {self._config.sandbox_timeout}s, execution_timeout: {self._execution_timeout()}s)"
+        )
+        return await self._execute_runtime(
+            wrapped_code=wrapped_code,
+            entrypoint_path=entrypoint_path,
+            files=files_to_upload if files_to_upload else None,
+            env=payload.env,
+        )
+
+    def _new_structured_output_path(self) -> str:
+        """Create an executor-controlled side-channel path for this run."""
+        return f"{STRUCTURED_OUTPUT_DIR}/{uuid4().hex}.jsonl"
+
+    def _is_valid_structured_output_path(self, path: str) -> bool:
+        """Validate that a side-channel path is under the reserved directory."""
+        parsed = PurePosixPath(path)
+        if parsed.parent != PurePosixPath(STRUCTURED_OUTPUT_DIR) or not parsed.name.endswith(".jsonl"):
+            return False
+        stem = parsed.name[: -len(".jsonl")]
+        return len(stem) == 32 and all(char in "0123456789abcdef" for char in stem)
+
+    async def _read_structured_output_text(self, path: str) -> str | None:
+        """Read bounded structured output text from a reserved side-channel file."""
+        quoted_path = shlex.quote(path)
+        async with asyncio.timeout(self._cleanup_timeout_s):
+            command_result = self._runtime.execute_command(
+                f"test -f {quoted_path} && head -c {MAX_STRUCTURED_OUTPUT_TOTAL_BYTES} {quoted_path} || true",
+            )
+            if not inspect.isawaitable(command_result):
+                logger.warning("Skipping PTC structured output read because runtime command was not awaitable")
+                return None
+            result = await command_result
+        if result.exit_code != 0 or not result.stdout:
+            return None
+        return result.stdout
+
+    @staticmethod
+    def _structured_output_lines(stdout: str) -> list[str]:
+        """Split JSONL output while dropping an invalid truncated final record."""
+        lines = stdout.splitlines()
+        if not lines or len(stdout.encode("utf-8")) < MAX_STRUCTURED_OUTPUT_TOTAL_BYTES:
+            return lines
+        if stdout.endswith(("\n", "\r")):
+            return lines  # pragma: no cover
+
+        trailing_line = lines[-1]
+        if len(trailing_line.encode("utf-8")) > MAX_STRUCTURED_OUTPUT_BYTES:
+            return lines  # pragma: no cover
+        try:
+            json.loads(trailing_line)
+        except json.JSONDecodeError:
+            logger.warning("Skipping truncated PTC structured output line")
+            return lines[:-1]
+        return lines  # pragma: no cover
+
+    @staticmethod
+    def _parse_structured_output_line(line: str) -> tuple[dict[str, Any] | None, str | None]:
+        """Parse one bounded structured output JSONL record and surface skipped records."""
+        line_size = len(line.encode("utf-8"))
+        if line_size > MAX_STRUCTURED_OUTPUT_BYTES:
+            warning = (
+                "Skipped oversized PTC structured output line "
+                f"({line_size} bytes exceeds {MAX_STRUCTURED_OUTPUT_BYTES} byte limit)"
+            )
+            logger.warning(warning)
+            return None, warning
+        try:
+            item = json.loads(line)
+        except json.JSONDecodeError:
+            logger.warning("Skipping malformed PTC structured output line")
+            return None, None
+        return (item, None) if isinstance(item, dict) else (None, None)
+
+    async def _read_structured_output(self, path: str) -> PTCStructuredOutput:
+        """Read and aggregate supported structured updates from the reserved JSONL file."""
+        from aip_agents.sandbox.types import PTCStructuredOutput
+
+        output = PTCStructuredOutput()
+        if not self._is_valid_structured_output_path(path):
+            logger.warning("Skipping invalid PTC structured output path")
+            return output
+
+        stdout = await self._read_structured_output_text(path)
+        if not stdout:
+            return output
+
+        for line in self._structured_output_lines(stdout):
+            item, warning = self._parse_structured_output_line(line)
+            if warning is not None:
+                output.add_warning(warning)
+            if item is not None:
+                output.merge(item)
+        return output
+
+    async def _cleanup_structured_output(self, path: str) -> None:
+        """Best-effort cleanup of the reserved side-channel file."""
+        if not self._is_valid_structured_output_path(path):
+            return
+        try:
+            async with asyncio.timeout(self._cleanup_timeout_s):
+                command_result = self._runtime.execute_command(f"rm -f {shlex.quote(path)}")
+                if not inspect.isawaitable(command_result):
+                    return
+                result = await command_result
+            exit_code = getattr(result, "exit_code", 0)
+            if not isinstance(exit_code, int) or exit_code == 0:
+                return
+
+            # In E2B, execute_command can run as a different sandbox user than
+            # execute_code. If it can read but not unlink a user-owned file, run
+            # the cleanup in the same Python execution plane that created it.
+            cleanup_code = (
+                "import os\n"
+                f"path = {json.dumps(path)}\n"
+                "try:\n"
+                "    os.remove(path)\n"
+                "except FileNotFoundError:\n"
+                "    pass\n"
+            )
+            execute_result = self._runtime.execute(code=cleanup_code, timeout=self._cleanup_timeout_s)
+            if inspect.isawaitable(execute_result):
+                await execute_result
+        except Exception as exc:
+            logger.warning(f"Failed to cleanup PTC structured output side-channel: {exc}")
+
+    async def _attach_structured_output(
+        self,
+        result: SandboxExecutionResult,
+        path: str,
+    ) -> SandboxExecutionResult:
+        """Attach parsed structured output and cleanup its side-channel file."""
+        try:
+            result.structured_output = await self._read_structured_output(path)
+        except Exception as exc:
+            logger.warning(f"Failed to read PTC structured output side-channel: {exc}")
+        finally:
+            await self._cleanup_structured_output(path)
+        return result
+
+    def _changed_per_run_files(
+        self,
+        payload: SandboxPayload,
+        current_hashes: dict[str, str],
+        entrypoint_path: str,
+        wrapped_code: str,
+    ) -> dict[str, str]:
+        """Return changed per-run files plus entrypoint when file execution is active."""
+        changed_per_run_files = {
+            file_path: content
+            for file_path, content in payload.per_run_files.items()
+            if self._per_run_hash_cache.get(file_path) != current_hashes[file_path]
+        }
+        if self._runtime_prefers_python_file_execution():
+            changed_per_run_files[entrypoint_path] = wrapped_code
+        return changed_per_run_files
+
+    async def _execute_incremental(
+        self,
+        payload: SandboxPayload,
+        wrapped_code: str,
+        current_hashes: dict[str, str],
+        entrypoint_path: str,
+    ) -> SandboxExecutionResult:
+        """Execute with only the changed per-run files (static bundle already uploaded)."""
+        changed_per_run_files = self._changed_per_run_files(payload, current_hashes, entrypoint_path, wrapped_code)
+        return await self._execute_runtime(
+            wrapped_code=wrapped_code,
+            entrypoint_path=entrypoint_path,
+            files=changed_per_run_files or None,
+            env=payload.env,
+        )
+
+    async def _finalize_cached_run(
+        self,
+        result: SandboxExecutionResult,
+        current_hashes: dict[str, str],
+        should_read_structured_output: bool,
+        structured_output_path: str,
+    ) -> SandboxExecutionResult:
+        """Update per-run hash cache and attach structured output if enabled."""
+        self._per_run_hash_cache = current_hashes
+        if should_read_structured_output:
+            return await self._attach_structured_output(result, structured_output_path)
+        return result
+
+    async def _execute_with_bundle(
+        self,
+        payload: SandboxPayload,
+        wrapped_code: str,
+        entrypoint_path: str | None = None,
+    ) -> SandboxExecutionResult:
+        """Execute code using the cached static bundle when possible."""
+        structured_output_path = self._new_structured_output_path()
+        should_read_structured_output = self._config.custom_tools.enabled and bool(self._config.custom_tools.tools)
+        if should_read_structured_output:
+            payload.env = {**payload.env, STRUCTURED_OUTPUT_ENV: structured_output_path}
+            wrapped_code = self._inject_structured_output_env(wrapped_code, structured_output_path)
+        entrypoint_path = entrypoint_path or self._entrypoint_path(wrapped_code)
+
+        async with self._per_run_lock:
+            try:
+                current_hashes = self._build_per_run_hashes(payload)
+                removed_paths = set(self._per_run_hash_cache) - set(current_hashes)
+                if removed_paths:
+                    try:
+                        await self._cleanup_runtime_for_removed_paths(removed_paths)
+                    except Exception:
+                        self._reset_cached_upload_state()
+                        raise
+
+                if self._static_bundle_uploaded:
+                    result = await self._execute_incremental(payload, wrapped_code, current_hashes, entrypoint_path)
+                    return await self._finalize_cached_run(
+                        result, current_hashes, should_read_structured_output, structured_output_path
+                    )
+
+                async with self._bundle_lock:
+                    if self._static_bundle_uploaded:
+                        result = await self._execute_incremental(payload, wrapped_code, current_hashes, entrypoint_path)
+                        return await self._finalize_cached_run(
+                            result, current_hashes, should_read_structured_output, structured_output_path
+                        )
+
+                    result = await self._execute_payload(
+                        payload,
+                        wrapped_code,
+                        entrypoint_path=entrypoint_path,
+                        upload_static_bundle=True,
+                    )
+                    self._update_static_bundle_cache()
+                    return await self._finalize_cached_run(
+                        result, current_hashes, should_read_structured_output, structured_output_path
+                    )
+            except Exception:
+                if should_read_structured_output:
+                    await self._cleanup_structured_output(structured_output_path)
+                raise
+
+    def _update_static_bundle_cache(self) -> None:
+        """Update cached bundle state after a successful upload."""
+        self._static_bundle_uploaded = True
+        logger.debug("Static bundle successfully uploaded and cached")
+
+    def _log_execution_result(self, result: SandboxExecutionResult) -> None:
+        """Log execution results based on exit status."""
+        if result.exit_code == 0:
+            logger.info("Sandbox execution completed successfully")
+        else:
+            logger.warning(f"Sandbox execution failed with exit code {result.exit_code}")
+
+    async def execute_code(
+        self,
+        code: str,
+        tool_configs: dict[str, dict] | None = None,
+    ) -> SandboxExecutionResult:
+        """Execute code inside the sandbox with MCP access.
+
+        This method:
+        1. Builds the sandbox payload (MCP config + generated tool modules + custom tools)
+        2. Wraps the user code with necessary imports and setup
+        3. Executes the code in the E2B sandbox
+        4. Returns the execution result (stdout/stderr/exit_code)
+
+        Args:
+            code: Python code to execute in the sandbox.
+            tool_configs: Optional per-tool config values for custom LangChain tools.
+                These are merged with agent defaults and written to custom_defaults.json.
+
+        Returns:
+            SandboxExecutionResult with stdout, stderr, and exit_code.
+
+        Raises:
+            PTCToolError: If sandbox execution fails.
+        """
+        try:
+            await self._refresh_runtime_state()
+            self._reset_bundle_state_if_inactive()
+            payload = await self._build_payload(tool_configs)
+            wrapped_code = self._wrap_code(code)
+            entrypoint_path = self._entrypoint_path(wrapped_code)
+            result = await self._execute_with_bundle(payload, wrapped_code, entrypoint_path)
+            self._log_execution_result(result)
+            return result
+
+        except Exception as exc:
+            logger.error(f"Sandbox execution failed: {exc}")
+            raise PTCToolError(f"Sandbox execution failed: {exc}") from exc
+
+    def reset_static_bundle(self) -> None:
+        """Reset the static bundle upload state.
+
+        Call this method when the sandbox is destroyed/recreated to force
+        a full re-upload of the static bundle on the next execution.
+        """
+        self._clear_upload_caches()
+        logger.debug("Static bundle state reset, next execution will re-upload full payload")

aip_agents/ptc/mcp/__init__.py

@@ -0,0 +1,45 @@
+"""Programmatic Tool Calling (PTC) module for MCP tools.
+
+This module provides programmatic tool calling capabilities for MCP tools,
+allowing code-based tool invocation instead of JSON tool calls.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from aip_agents.ptc.exceptions import PTCError, PTCToolError
+from aip_agents.ptc.mcp.sandbox_bridge import (
+    ServerConfig,
+    build_mcp_payload,
+)
+from aip_agents.ptc.naming import (
+    json_type_to_python,
+    sanitize_function_name,
+    sanitize_module_name,
+    sanitize_param_name,
+    schema_to_params,
+)
+from aip_agents.ptc.payload import SandboxPayload
+from aip_agents.ptc.prompt_builder import (
+    build_ptc_prompt,
+    compute_ptc_prompt_hash,
+)
+
+__all__ = [
+    # Exceptions
+    "PTCError",
+    "PTCToolError",
+    # Naming utilities
+    "json_type_to_python",
+    "sanitize_function_name",
+    "sanitize_module_name",
+    "sanitize_param_name",
+    "schema_to_params",
+    # Prompt builder
+    "build_ptc_prompt",
+    "compute_ptc_prompt_hash",
+    # Sandbox Bridge
+    "SandboxPayload",
+    "ServerConfig",
+    "build_mcp_payload",
+]