aip-agents-binary 0.5.21__py3-none-macosx_13_0_arm64.whl → 0.6.8__py3-none-macosx_13_0_arm64.whl

aip_agents/agent/langgraph_react_agent.py

@@ -6,8 +6,11 @@ Authors:
     Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
     Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
     Raymond Christopher (raymond.christopher@gdplabs.id)
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
 """
 
+from __future__ import annotations
+
 import asyncio
 import time
 import uuid
@@ -15,11 +18,14 @@ from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import asdict, dataclass
 from functools import reduce
 from textwrap import dedent
-from typing import Annotated, Any
+from typing import TYPE_CHECKING, Annotated, Any, cast
+
+from deprecated import deprecated  # type: ignore[import-untyped]
 
-from deprecated import deprecated
-from gllm_core.event import EventEmitter
-from gllm_core.schema import Chunk
+if TYPE_CHECKING:
+    from aip_agents.guardrails.manager import GuardrailManager
+from gllm_core.event import EventEmitter  # type: ignore[import-untyped]
+from gllm_core.schema import Chunk  # type: ignore[import-untyped]
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import (
     AIMessage,
@@ -48,7 +54,7 @@ from aip_agents.schema.a2a import A2AStreamEventType
 from aip_agents.schema.hitl import ApprovalDecision, HitlMetadata
 from aip_agents.schema.langgraph import ToolCallResult, ToolStorageParams
 from aip_agents.schema.step_limit import MaxStepsExceededError, StepLimitConfig
-from aip_agents.tools.memory_search_tool import MEMORY_SEARCH_TOOL_NAME
+from aip_agents.tools.memory_search_tool import MEMORY_DELETE_TOOL_NAME, MEMORY_SEARCH_TOOL_NAME
 from aip_agents.tools.tool_config_injector import TOOL_CONFIGS_KEY
 from aip_agents.utils import add_references_chunks
 from aip_agents.utils.langgraph import (
@@ -81,6 +87,9 @@ from aip_agents.utils.token_usage_helper import (
     extract_token_usage_from_tool_output,
 )
 
+if TYPE_CHECKING:
+    from aip_agents.ptc import PTCSandboxConfig
+
 logger = get_logger(__name__)
 
 # Default instruction for ReAct agents
@@ -157,7 +166,9 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         tool_output_manager: ToolOutputManager | None = None,
         planning: bool = False,
         middlewares: Sequence[AgentMiddleware] | None = None,
+        guardrail: GuardrailManager | None = None,
         step_limit_config: StepLimitConfig | None = None,
+        ptc_config: PTCSandboxConfig | None = None,
         **kwargs: Any,
     ):
         """Initialize the LangGraph ReAct Agent.
@@ -178,10 +189,19 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             planning: Enable planning capabilities with TodoListMiddleware. Defaults to False.
             middlewares: Optional sequence of custom middleware to COMPOSE (not override) with built-in middleware.
                         Execution order: [TodoListMiddleware (if planning=True),
+                                         GuardrailMiddleware (if guardrail provided),
                                          ...custom middlewares in order provided]
                         All middleware hooks execute - this extends capabilities, never replaces them.
+            guardrail: Optional GuardrailManager for content filtering and safety checks.
+                     When provided, automatically wraps in GuardrailMiddleware for transparent
+                     input/output filtering during agent execution.
             enable_pii: Optional toggle to enable PII handling for tool inputs and outputs.
             step_limit_config: Optional configuration for step limits and delegation depth.
+            ptc_config: Optional configuration for PTC sandbox execution. See PTCSandboxConfig
+                for available options including enabled flag, sandbox timeout, and template settings.
+                PTC is enabled when ptc_config is not None and ptc_config.enabled is True.
+                When enabled, prompt guidance is automatically injected into the agent's instruction.
+                PTC runs in a sandbox only; there is no in-process trusted PTC path.
             **kwargs: Additional keyword arguments passed to BaseLangGraphAgent.
         """
         # Use LangGraph's standard AgentState for ReAct
@@ -201,6 +221,12 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             **kwargs,
         )
 
+        if self.model is None and self.lm_invoker is None:
+            logger.warning(
+                "Agent '%s': Model and LM invoker are both unset. Calls that require a model will fail.",
+                self.name,
+            )
+
         # Handle tool output management
         self.tool_output_manager = tool_output_manager
         self._pii_handlers_by_thread: dict[str, ToolPIIHandler] = {}
@@ -212,6 +238,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         # Setup middleware
         self._middleware_manager = self._setup_middleware(
             planning=planning,
+            guardrail=guardrail,
             custom_middlewares=middlewares,
         )
 
@@ -221,18 +248,32 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
         self.step_limit_config = step_limit_config
 
+        # Initialize PTC state (Programmatic Tool Calling)
+        self._ptc_config: PTCSandboxConfig | None = None
+        self._ptc_tool_synced = False
+        self._ptc_tool: BaseTool | None = None
+        self._ptc_prompt_hash: str = ""
+        # Capture instruction after middleware setup so middleware prompts are preserved
+        self._original_instruction: str = self.instruction
+
+        # Enable PTC if requested via constructor
+        if ptc_config is not None and ptc_config.enabled:
+            self.enable_ptc(ptc_config)
+
     def _setup_middleware(
         self,
         planning: bool,
+        guardrail: GuardrailManager | None,
         custom_middlewares: Sequence[AgentMiddleware] | None,
     ) -> MiddlewareManager | None:
         """Setup middleware based on configuration.
 
-        Creates auto-configured middleware (planning) and composes
+        Creates auto-configured middleware (planning, guardrails) and composes
         with custom middleware if provided.
 
         Args:
             planning: Whether to enable TodoListMiddleware.
+            guardrail: Optional GuardrailManager to wrap in GuardrailMiddleware.
             custom_middlewares: Optional custom middlewares to append.
 
         Returns:
@@ -242,7 +283,13 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
         # Auto-configure TodoListMiddleware if planning enabled
         if planning:
-            middleware_list.append(TodoListMiddleware())
+            middleware_list.append(cast(AgentMiddleware, TodoListMiddleware()))
+
+        # Auto-configure GuardrailMiddleware if guardrail provided
+        if guardrail:
+            from aip_agents.guardrails.middleware import GuardrailMiddleware
+
+            middleware_list.append(GuardrailMiddleware(guardrail))
 
         # Append custom middlewares
         if custom_middlewares:
@@ -399,9 +446,9 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             )
 
     def _rebuild_resolved_tools(self) -> None:
-        """Rebuild resolved tools including middleware tools.
+        """Rebuild resolved tools including middleware and PTC tools.
 
-        Overrides base class to ensure middleware tools are preserved
+        Overrides base class to ensure middleware tools and the PTC tool are preserved
         when tools are rebuilt (e.g., after update_regular_tools).
         """
         # Call base class to rebuild with regular, a2a, delegation, and mcp tools
@@ -411,6 +458,10 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         if hasattr(self, "_middleware_tools") and self._middleware_tools:
             self.resolved_tools.extend(self._middleware_tools)
 
+        # Add PTC tool if synced
+        if hasattr(self, "_ptc_tool") and self._ptc_tool is not None:
+            self.resolved_tools.append(self._ptc_tool)
+
     def _handle_tool_artifacts(
         self, tool_output: Any, pending_artifacts: list[dict[str, Any]]
     ) -> tuple[str, list[dict[str, Any]]]:
@@ -533,6 +584,29 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
         return memory_node
 
+    def _should_save_interaction(self, final_state: dict[str, Any] | None) -> bool:
+        """Return True when interaction should be saved to memory."""
+        if self._contains_memory_delete_action(final_state):
+            logger.info("Memory: Skipping save_interaction due to memory delete action in state.")
+            return False
+        return True
+
+    @staticmethod
+    def _contains_memory_delete_action(final_state: dict[str, Any] | None) -> bool:
+        """Return True when final state includes a delete memory action block."""
+        if not isinstance(final_state, dict):
+            return False
+        messages = final_state.get("messages")
+        if not isinstance(messages, list):
+            return False
+        for message in messages:
+            content = getattr(message, "content", None)
+            if not isinstance(content, str):
+                continue
+            if "<MEMORY_ACTION>" in content and "action=delete" in content:
+                return True
+        return False
+
     def _extract_user_query_from_messages(self, messages: list[Any]) -> str | None:
         """Get latest user query string from a list of messages.
 
@@ -579,14 +653,29 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             current_messages = state["messages"]
 
             # Execute LLM call
-            if self.lm_invoker:
-                result = await self._handle_lm_invoker_call(current_messages, state, config)
-            elif isinstance(self.model, BaseChatModel):
-                result = await self._handle_langchain_model_call(current_messages, state, config)
-            else:
-                raise ValueError(
-                    f"Agent '{self.name}': No valid LMInvoker or LangChain model configured for ReAct agent node."
-                )
+            try:
+                if self.lm_invoker:
+                    result = await self._handle_lm_invoker_call(current_messages, state, config)
+                elif isinstance(self.model, BaseChatModel):
+                    result = await self._handle_langchain_model_call(current_messages, state, config)
+                else:
+                    raise ValueError(
+                        f"Agent '{self.name}': No valid LMInvoker or LangChain model configured for ReAct agent node."
+                    )
+            except Exception as e:
+                # Lazy import to support optional guardrails dependency
+                from aip_agents.guardrails.exceptions import GuardrailViolationError
+
+                if isinstance(e, GuardrailViolationError):
+                    return {
+                        "messages": [
+                            AIMessage(
+                                content=f"⚠️ Guardrail violation: {e.result.reason}",
+                                response_metadata={"finish_reason": "stop"},
+                            )
+                        ]
+                    }
+                raise
 
             # Increment step counter after successful execution
             manager.increment_step()
@@ -697,7 +786,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         pending_artifacts: list[dict[str, Any]] = state.get("artifacts") or []
         reference_updates: list[Chunk] = []
         tool_map = {tool.name: tool for tool in self.resolved_tools}
-        pii_mapping = {}
+        pii_mapping: dict[str, str] = {}
 
         aggregated_metadata_delta: dict[str, Any] = {}
         total_tools_token_usage: list[UsageMetadata] = []
@@ -721,7 +810,8 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
                 ),
             )
 
-        tasks = [asyncio.create_task(run_tool(tc)) for tc in last_message.tool_calls]
+        normalized_tool_calls = [self._normalize_tool_call(tc) for tc in last_message.tool_calls]
+        tasks = [asyncio.create_task(run_tool(tc)) for tc in normalized_tool_calls]
 
         for coro in asyncio.as_completed(tasks):
             tool_result = await coro
@@ -744,6 +834,31 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             pii_mapping,
         )
 
+    def _normalize_tool_call(self, tool_call: Any) -> dict[str, Any]:
+        """Normalize tool call inputs into a dict with required keys."""
+        if isinstance(tool_call, dict):
+            normalized = dict(tool_call)
+        elif hasattr(tool_call, "model_dump"):
+            normalized = tool_call.model_dump()
+        elif hasattr(tool_call, "dict"):
+            normalized = tool_call.dict()
+        elif hasattr(tool_call, "name") and hasattr(tool_call, "args"):
+            normalized = {
+                "id": getattr(tool_call, "id", None),
+                "name": getattr(tool_call, "name", None),
+                "args": getattr(tool_call, "args", None),
+            }
+        else:
+            raise TypeError("Tool call must be a dict-like object or ToolCall instance.")
+
+        if not isinstance(normalized, dict):
+            raise TypeError("Tool call normalization did not produce a dict.")
+
+        if "name" not in normalized or "args" not in normalized:
+            raise TypeError("Tool call must include 'name' and 'args' fields.")
+
+        return normalized
+
     def _accumulate_tool_result(  # noqa: PLR0913
         self,
         tool_result: Any,
@@ -752,7 +867,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         aggregated_metadata_delta: dict[str, Any],
         reference_updates: list[Chunk],
         total_tools_token_usage: list[UsageMetadata],
-        pii_mapping: dict[str, str] | None,
+        pii_mapping: dict[str, str],
     ) -> None:  # noqa: PLR0913
         """Accumulate results from a single tool call.
 
@@ -1198,13 +1313,16 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
             # Create enhanced tool configuration with output management
             tool_config = self._create_enhanced_tool_config(config, state, tool_call["name"], tool_call_id)
+            if not isinstance(tool_config, dict):
+                raise TypeError("Tool configuration must be a dictionary.")
+            tool_config_runnable = tool_config
 
             arun_streaming_method = getattr(tool, TOOL_RUN_STREAMING_METHOD, None)
 
             if arun_streaming_method and callable(arun_streaming_method):
                 tool_output = await self._execute_tool_with_streaming(tool, tool_call, tool_config)
             else:
-                tool_output = await tool.ainvoke(resolved_args, tool_config)
+                tool_output = await tool.ainvoke(resolved_args, tool_config_runnable)
 
             references = extract_references_from_tool(tool, tool_output)
 
@@ -1478,7 +1596,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         tool_call: dict[str, Any],
         execution_time: float,
         pending_artifacts: list[dict[str, Any]],
-    ) -> tuple[list[BaseMessage], list[dict[str, Any]], dict[str, Any]]:
+    ) -> tuple[list[ToolMessage], list[dict[str, Any]], dict[str, Any]]:
         """Process tool output into messages, artifacts, and metadata.
 
         Args:
@@ -1506,7 +1624,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
     def _handle_command_output(
         self, tool_output: Command, tool_call: dict[str, Any], execution_time: float, metadata_delta: dict[str, Any]
-    ) -> tuple[list[BaseMessage], list[dict[str, Any]], dict[str, Any]]:
+    ) -> tuple[list[ToolMessage], list[dict[str, Any]], dict[str, Any]]:
         """Handle Command type tool outputs.
 
         Args:
@@ -1535,7 +1653,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
     def _handle_string_output(
         self, tool_output: str, tool_call: dict[str, Any], execution_time: float
-    ) -> tuple[list[BaseMessage], list[dict[str, Any]], dict[str, Any]]:
+    ) -> tuple[list[ToolMessage], list[dict[str, Any]], dict[str, Any]]:
         """Handle string type tool outputs.
 
         Args:
@@ -1561,7 +1679,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         execution_time: float,
         pending_artifacts: list[dict[str, Any]],
         metadata_delta: dict[str, Any],
-    ) -> tuple[list[BaseMessage], list[dict[str, Any]], dict[str, Any]]:
+    ) -> tuple[list[ToolMessage], list[dict[str, Any]], dict[str, Any]]:
         """Handle legacy dict and other tool outputs.
 
         Args:
@@ -1659,8 +1777,11 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             self._emit_default_tool_call_event(writer, tool_name, tool_call_id, tool_args)
 
             streaming_kwargs = self._build_streaming_kwargs(tool_args, tool_config)
+            arun_streaming_method = getattr(tool, TOOL_RUN_STREAMING_METHOD, None)
+            if not callable(arun_streaming_method):
+                raise RuntimeError(f"Tool '{tool_name}' does not implement streaming.")
 
-            async for chunk in tool.arun_streaming(**streaming_kwargs):
+            async for chunk in arun_streaming_method(**streaming_kwargs):
                 final_output, saw_tool_result = self._handle_streaming_chunk(
                     chunk=chunk,
                     writer=writer,
@@ -1954,6 +2075,47 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         )
         writer(a2a_event)
 
+    async def _execute_abefore_model_hook(self, state: dict[str, Any]) -> None:
+        """Asynchronously execute abefore_model middleware hook and update state.
+
+        Args:
+            state: Current agent state to potentially update.
+        """
+        if self._middleware_manager:
+            try:
+                before_updates = await self._middleware_manager.abefore_model(state)
+                if before_updates:
+                    state.update(before_updates)
+            except Exception as e:
+                # Lazy import to support optional guardrails dependency
+                from aip_agents.guardrails.exceptions import GuardrailViolationError
+
+                if isinstance(e, GuardrailViolationError):
+                    # Re-raise guardrail violations to be caught by the agent node
+                    raise
+                logger.error(f"Agent '{self.name}': Middleware abefore_model hook failed: {e}")
+
+    async def _execute_aafter_model_hook(self, state_updates: dict[str, Any], state: dict[str, Any]) -> None:
+        """Asynchronously execute aafter_model middleware hook.
+
+        Args:
+            state_updates: Updates to be merged into state.
+            state: Current agent state for context.
+        """
+        if self._middleware_manager:
+            try:
+                after_updates = await self._middleware_manager.aafter_model(state)
+                if after_updates:
+                    state_updates.update(after_updates)
+            except Exception as e:
+                # Lazy import to support optional guardrails dependency
+                from aip_agents.guardrails.exceptions import GuardrailViolationError
+
+                if isinstance(e, GuardrailViolationError):
+                    # Re-raise guardrail violations
+                    raise
+                logger.error(f"Agent '{self.name}': Middleware aafter_model hook failed: {e}")
+
     def _execute_before_model_hook(self, state: dict[str, Any]) -> None:
         """Execute before_model middleware hook and update state.
 
@@ -1966,6 +2128,12 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
                 if before_updates:
                     state.update(before_updates)
             except Exception as e:
+                # Lazy import to support optional guardrails dependency
+                from aip_agents.guardrails.exceptions import GuardrailViolationError
+
+                if isinstance(e, GuardrailViolationError):
+                    # Re-raise guardrail violations to be caught by the agent node
+                    raise
                 logger.error(f"Agent '{self.name}': Middleware before_model hook failed: {e}")
 
     def _execute_modify_model_request_hook(
@@ -2029,7 +2197,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             dict[str, Any]: A dictionary containing the new messages and updated token usage.
         """
         # Execute before_model middleware hook
-        self._execute_before_model_hook(state)
+        await self._execute_abefore_model_hook(state)
 
         # Build tool output aware instruction
         enhanced_instruction = self._build_tool_output_aware_instruction(self.instruction, state, config)
@@ -2043,6 +2211,9 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
 
         effective_event_emitter = state.get("event_emitter") or self.event_emitter
 
+        if self.lm_invoker is None:
+            raise RuntimeError("LM invoker is required for this execution path.")
+
         if self.resolved_tools:
             self.lm_invoker.set_tools(self.resolved_tools)
 
@@ -2063,7 +2234,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         state_updates.update(token_usage_updates)
 
         # Execute after_model middleware hook
-        self._execute_after_model_hook(state_updates, state)
+        await self._execute_aafter_model_hook(state_updates, state)
 
         return state_updates
 
@@ -2081,7 +2252,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
             dict[str, Any]: A dictionary containing the new messages and updated token usage.
         """
         # Execute before_model middleware hook
-        self._execute_before_model_hook(state)
+        await self._execute_abefore_model_hook(state)
 
         # Build tool output aware instruction
         enhanced_instruction = self._build_tool_output_aware_instruction(self.instruction, state, config)
@@ -2101,6 +2272,9 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         ):
             langchain_prompt = [SystemMessage(content=enhanced_instruction)] + list(current_messages)
 
+        if self.model is None:
+            raise RuntimeError("Model is required for this execution path.")
+
         model_with_tools = self.model.bind_tools(self.resolved_tools) if self.resolved_tools else self.model
 
         ai_message = await model_with_tools.ainvoke(langchain_prompt, config)
@@ -2113,7 +2287,7 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         state_updates.update(token_usage_updates)
 
         # Execute after_model middleware hook
-        self._execute_after_model_hook(state_updates, state)
+        await self._execute_aafter_model_hook(state_updates, state)
 
         return state_updates
 
@@ -2126,11 +2300,12 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         """
         try:
             tool_cfgs = metadata.get(TOOL_CONFIGS_KEY, {})
-            per_tool_config = tool_cfgs.get(MEMORY_SEARCH_TOOL_NAME)
-            if not isinstance(per_tool_config, dict):
-                per_tool_config = {}
-            per_tool_config["user_id"] = memory_user_id
-            tool_cfgs[MEMORY_SEARCH_TOOL_NAME] = per_tool_config
+            for tool_name in (MEMORY_SEARCH_TOOL_NAME, MEMORY_DELETE_TOOL_NAME):
+                per_tool_config = tool_cfgs.get(tool_name)
+                if not isinstance(per_tool_config, dict):
+                    per_tool_config = {}
+                per_tool_config["user_id"] = memory_user_id
+                tool_cfgs[tool_name] = per_tool_config
             metadata[TOOL_CONFIGS_KEY] = tool_cfgs
         except Exception as e:
             # Non-fatal; metadata injection is best-effort
@@ -2494,6 +2669,220 @@ class LangGraphReactAgent(LangGraphHitLMixin, BaseLangGraphAgent):
         if current_thread_id:
             self._pii_handlers_by_thread.pop(current_thread_id, None)
 
+    # ==========================================================================
+    # Programmatic Tool Calling (PTC) Methods
+    # ==========================================================================
+
+    def add_mcp_server(self, mcp_config: dict[str, dict[str, Any]]) -> None:
+        """Add MCP servers and refresh PTC tool state if needed."""
+        super().add_mcp_server(mcp_config)
+
+        if not self._ptc_config or not self._ptc_config.enabled:
+            return
+
+        if self._ptc_tool is not None:
+            self._ptc_tool = None
+
+        self._ptc_tool_synced = False
+        logger.debug(f"Agent '{self.name}': PTC tool will resync after MCP changes")
+
+    def enable_ptc(self, config: PTCSandboxConfig | None = None) -> None:
+        """Enable Programmatic Tool Calling (PTC) for this agent.
+
+        PTC allows the LLM to execute Python code that calls MCP tools
+        programmatically inside a sandboxed environment. This is useful for
+        chaining multiple tool calls with local data processing.
+
+        The execute_ptc_code tool is automatically added to the agent's tools
+        after MCP servers are configured. If no MCP servers are configured,
+        the tool sync is deferred until servers are added.
+
+        Args:
+            config: Optional configuration for PTC sandbox execution.
+                See PTCSandboxConfig for options like enabled flag and sandbox_timeout.
+                If None is passed, a default config with enabled=True will be created.
+
+        Example:
+            agent.enable_ptc(PTCSandboxConfig(enabled=True))
+            agent.add_mcp_server({"yfinance": {...}})
+            # execute_ptc_code tool is now available
+
+        Note:
+            PTC can also be enabled via the constructor by passing
+            ptc_config=PTCSandboxConfig(enabled=True, ...).
+        """
+        # Lazy import to avoid circular dependencies
+        from aip_agents.ptc.executor import PTCSandboxConfig
+
+        self._ptc_config = config or PTCSandboxConfig()
+        self._ptc_config.enabled = True
+        self._ptc_tool_synced = False
+
+        logger.info(f"Agent '{self.name}': PTC enabled")
+
+        # Attempt to sync PTC tool if MCP client is available
+        self._sync_ptc_tool()
+
+    def _enrich_custom_tool_metadata(self) -> None:
+        """Enrich custom tool definitions with metadata from agent's tools.
+
+        This method matches custom tool definitions in ptc_config with actual tool
+        objects from self.tools and enriches them with description and input_schema.
+        Called at PTC sync time to ensure tool definitions have accurate metadata.
+
+        The matching is done by comparing tool names (both original and sanitized).
+
+        Note: This method modifies self._ptc_config.custom_tools.tools in-place.
+        """
+        if not self._ptc_config:
+            return
+
+        # Lazy import to avoid circular dependencies
+        from aip_agents.ptc import enrich_custom_tools_from_agent
+
+        enrich_custom_tools_from_agent(
+            self._ptc_config.custom_tools,
+            self.tools,
+            agent_name=self.name,
+        )
+
+    def _sync_ptc_tool(self) -> None:
+        """Build and register the execute_ptc_code tool when MCP or custom tools are available.
+
+        This method is called after enable_ptc() and after MCP servers are added.
+        It creates the execute_ptc_code tool using the current MCP client
+        configuration and/or custom tools config and adds it to the agent's resolved tools.
+
+        The tool is only created once. Subsequent calls are no-ops if the tool
+        has already been synced.
+        """
+        if not self._ptc_config or not self._ptc_config.enabled:
+            return
+
+        if self._ptc_tool_synced:
+            return
+
+        # Check if we have custom tools enabled
+        has_custom_tools = self._ptc_config.custom_tools.enabled and self._ptc_config.custom_tools.tools
+        has_mcp_servers = bool(self.mcp_config)
+
+        if has_mcp_servers:
+            if not self.mcp_client:
+                logger.debug(f"Agent '{self.name}': PTC tool sync deferred - no MCP client yet")
+                return
+
+            if not self.mcp_client.is_initialized:
+                logger.debug(f"Agent '{self.name}': PTC tool sync deferred - MCP client not initialized")
+                return
+
+            mcp_client_to_use = self.mcp_client
+        elif has_custom_tools:
+            mcp_client_to_use = None
+        else:
+            return
+
+        # Enrich custom tool definitions with metadata from actual tool objects
+        if has_custom_tools:
+            self._enrich_custom_tool_metadata()
+
+        # Lazy import to avoid circular dependencies
+        from aip_agents.tools.execute_ptc_code import create_execute_ptc_code_tool
+
+        if has_custom_tools and not mcp_client_to_use:
+            logger.info(f"Agent '{self.name}': Syncing PTC tool with custom tools only (no MCP)")
+        else:
+            logger.info(f"Agent '{self.name}': Syncing PTC tool with MCP client")
+
+        # Create the execute_ptc_code tool with agent's tool configs
+        self._ptc_tool = create_execute_ptc_code_tool(
+            mcp_client_to_use, self._ptc_config, agent_tool_configs=self.tool_configs
+        )
+
+        # Rebuild graph to include PTC tool
+        self._rebuild_graph()
+
+        self._ptc_tool_synced = True
+        logger.info(f"Agent '{self.name}': PTC tool synced successfully")
+
+        # Sync PTC prompt guidance
+        self._sync_ptc_prompt()
+
+    def _sync_ptc_prompt(self) -> None:
+        """Sync PTC usage guidance into the agent instruction.
+
+        This method builds and injects a PTC usage block into the agent's
+        instruction when PTC is enabled. The prompt is refreshed when MCP
+        or custom tools configuration changes (detected via hash).
+        """
+        if not self._ptc_config or not self._ptc_config.enabled:
+            return
+
+        # Check if we have custom tools enabled
+        has_custom_tools = self._ptc_config.custom_tools.enabled and self._ptc_config.custom_tools.tools
+
+        # For custom-only configs, allow None mcp_client
+        # For MCP-only or MCP+custom, require mcp_client
+        if not has_custom_tools and not self.mcp_client:
+            return
+
+        # Lazy import to avoid circular dependencies
+        from aip_agents.ptc.prompt_builder import build_ptc_prompt, compute_ptc_prompt_hash
+
+        # Get prompt config and custom tools config from PTC sandbox config
+        prompt_config = self._ptc_config.prompt if self._ptc_config else None
+        custom_tools_config = self._ptc_config.custom_tools if self._ptc_config else None
+
+        # Use mcp_client if available, None for custom-only
+        mcp_client_to_use = self.mcp_client if self.mcp_client else None
+
+        # Check if MCP or custom tools config has changed
+        current_hash = compute_ptc_prompt_hash(
+            mcp_client_to_use, config=prompt_config, custom_tools_config=custom_tools_config
+        )
+        if current_hash == self._ptc_prompt_hash:
+            logger.debug(f"Agent '{self.name}': PTC prompt unchanged, skipping refresh")
+            return
+
+        # Build and inject the prompt
+        ptc_prompt = build_ptc_prompt(mcp_client_to_use, config=prompt_config, custom_tools_config=custom_tools_config)
+
+        # Rebuild instruction from original + PTC guidance
+        self.instruction = f"{self._original_instruction}\n\n{ptc_prompt}"
+        self._ptc_prompt_hash = current_hash
+
+        logger.info(f"Agent '{self.name}': PTC prompt guidance injected")
+
+    async def _register_mcp_tools(self) -> None:
+        """Override to sync PTC tool after MCP tools are registered.
+
+        This extends the base implementation to ensure the execute_ptc_code
+        tool is added after MCP servers are initialized.
+        """
+        await super()._register_mcp_tools()
+
+        # Sync PTC tool after MCP tools are registered
+        if self._ptc_config and self._ptc_config.enabled and not self._ptc_tool_synced:
+            self._sync_ptc_tool()
+
+    async def cleanup(self) -> None:
+        """Cleanup agent resources including PTC sandbox.
+
+        Extends base cleanup to also cleanup the PTC sandbox runtime if
+        execute_ptc_code tool was created.
+        """
+        # Cleanup PTC tool's sandbox runtime if present
+        if self._ptc_tool is not None:
+            try:
+                cleanup_method = getattr(self._ptc_tool, "cleanup", None)
+                if cleanup_method and callable(cleanup_method):
+                    await cleanup_method()
+                    logger.debug(f"Agent '{self.name}': PTC sandbox cleanup completed")
+            except Exception as e:
+                logger.warning(f"Agent '{self.name}': Error during PTC sandbox cleanup: {e}")
+
+        # Call parent cleanup for MCP client
+        await super().cleanup()
+
     def _format_graph_output(self, final_state_result: dict[str, Any]) -> Any:
         """Convert final graph state to user-friendly output.