universal-mcp-agents 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

universal_mcp/agents/codeact0/playbook_agent.py

@@ -0,0 +1,353 @@
+import inspect
+import json
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, cast
+
+from langchain_core.messages import AIMessage, RemoveMessage, ToolMessage
+from langchain_core.tools import StructuredTool
+from langchain_core.tools import tool as create_tool
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import START, StateGraph
+from langgraph.types import Command, RetryPolicy
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat, ToolConfig
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.codeact0.llm_tool import ai_classify, call_llm, data_extractor, smart_print
+from universal_mcp.agents.codeact0.prompts import (
+    create_default_prompt,
+)
+from universal_mcp.agents.codeact0.sandbox import eval_unsafe, execute_ipython_cell
+from universal_mcp.agents.codeact0.state import CodeActState
+from universal_mcp.agents.codeact0.tools import create_meta_tools, enter_playbook_mode, exit_playbook_mode, get_valid_tools
+from universal_mcp.agents.codeact0.utils import inject_context, smart_truncate
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.utils import filter_retry_on, get_message_text
+
+PLAYBOOK_PLANNING_PROMPT = """Now, you are tasked with creating a reusable playbook from the user's previous workflow.
+
+TASK: Analyze the conversation history and code execution to create a step-by-step plan for a reusable function. Do not include the searching and loading of tools. Assume that the tools have already been loaded.
+
+Your plan should:
+1. Identify the key steps in the workflow
+2. Mark user-specific variables that should become the main playbook function parameters using `variable_name` syntax. Intermediate variables should not be highlighted using ``
+3. Keep the logic generic and reusable
+4. Be clear and concise
+
+Example:
+```
+1. Connect to database using `db_connection_string`
+2. Query user data for `user_id`
+3. Process results and calculate `metric_name`
+4. Send notification to `email_address`
+```
+
+Now create a plan based on the conversation history. Enclose it between ``` and ```. Ask the user if the plan is okay."""
+
+
+
+PLAYBOOK_CONFIRMING_PROMPT = """Now, you are tasked with confirming the playbook plan. Return True if the user is happy with the plan, False otherwise. Do not say anything else in your response. The user response will be the last message in the chain.
+"""
+
+PLAYBOOK_GENERATING_PROMPT = """Now, you are tasked with generating the playbook function. Return the function in Python code.
+Do not include any other text in your response.
+The function should be a single, complete piece of code that can be executed independently, based on previously executed code snippets that executed correctly.
+The parameters of the function should be the same as the final confirmed playbook plan.
+Do not include anything other than python code in your response
+"""
+
+
+class CodeActPlaybookAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        memory: BaseCheckpointSaver | None = None,
+        tools: ToolConfig | None = None,
+        registry: ToolRegistry | None = None,
+        sandbox_timeout: int = 20,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            instructions=instructions,
+            model=model,
+            memory=memory,
+            **kwargs,
+        )
+        self.model_instance = load_chat_model(model, thinking=True)
+        self.tools_config = tools or []
+        self.registry = registry
+        self.eval_fn = eval_unsafe
+        self.sandbox_timeout = sandbox_timeout
+        self.processed_tools: list[StructuredTool | Callable] = []
+
+    async def _build_graph(self):
+        self.exported_tools = []
+        if self.tools_config:
+            # Convert dict format to list format if needed
+            if isinstance(self.tools_config, dict):
+                self.tools_config = [
+                    f"{provider}__{tool}" 
+                    for provider, tools in self.tools_config.items() 
+                    for tool in tools
+                ]
+            if not self.registry:
+                raise ValueError("Tools are configured but no registry is provided")
+            # Langchain tools are fine
+            self.exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
+        meta_tools = create_meta_tools(self.registry)
+        await self.registry.export_tools(["exa__search_with_filters"], ToolFormat.LANGCHAIN)
+        additional_tools = [smart_print, data_extractor, ai_classify, call_llm, meta_tools["web_search"]]
+        self.additional_tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in additional_tools]
+        self.final_instructions, self.tools_context = create_default_prompt(
+            self.exported_tools, self.additional_tools, self.instructions
+        )
+
+        def call_model(state: CodeActState) -> Command[Literal["sandbox", "execute_tools"]]:
+            messages = [{"role": "system", "content": self.final_instructions}] + state["messages"]
+
+            # Run the model and potentially loop for reflection
+            model_with_tools = self.model_instance.bind_tools(
+                tools=[
+                    execute_ipython_cell,
+                    enter_playbook_mode,
+                    meta_tools["search_functions"],
+                    meta_tools["load_functions"],
+                ],
+                tool_choice="auto",
+            )
+            response = cast(AIMessage, model_with_tools.invoke(messages))
+            if response.tool_calls:
+                return Command(goto="execute_tools", update={"messages": [response]})
+            else:
+                return Command(update={"messages": [response], "model_with_tools": model_with_tools})
+
+            # if response.tool_calls:
+            #     if len(response.tool_calls) > 1:
+            #         raise Exception("Not possible in Claude with llm.bind_tools(tools=tools, tool_choice='auto')")
+            #     if response.tool_calls[0]["name"] == "enter_playbook_mode":
+            #         return Command(goto="playbook", update = {"playbook_mode": "planning"})
+            #     if response.tool_calls[0]["name"] != "execute_ipython_cell":
+            #         raise Exception(
+            #             f"Unexpected tool call: {response.tool_calls[0]['name']}. Expected 'execute_ipython_cell'."
+            #         )
+            #     if (
+            #         response.tool_calls[0]["args"].get("snippet") is None
+            #         or not response.tool_calls[0]["args"]["snippet"].strip()
+            #     ):
+            #         raise Exception("Tool call 'execute_ipython_cell' requires a non-empty 'snippet' argument.")
+            #     return Command(goto="sandbox", update={"messages": [response]})
+            # else:
+            #     return Command(update={"messages": [response]})
+
+        async def execute_tools(state: CodeActState) -> Command[Literal["call_model", "playbook", "sandbox"]]:
+            """Execute tool calls"""
+            last_message = state["messages"][-1]
+            tool_calls = last_message.tool_calls if isinstance(last_message, AIMessage) else []
+
+            tool_messages = []
+            new_tool_ids = []
+            ask_user = False
+            ai_msg = ""
+            tool_result = ""
+
+            for tool_call in tool_calls:
+                try:
+                    if tool_call["name"] == "enter_playbook_mode":
+                        tool_message = ToolMessage(
+                            content=json.dumps("Entered Playbook Mode."),
+                            name=tool_call["name"],
+                            tool_call_id=tool_call["id"],
+                        )
+                        return Command(
+                            goto="playbook",
+                            update={"playbook_mode": "planning", "messages": [tool_message]}, #Entered Playbook mode
+                        )
+                    elif tool_call["name"] == "execute_ipython_cell":
+                        return Command(goto="sandbox")
+                    elif tool_call["name"] == "load_functions":  # Handle load_functions separately
+                        valid_tools, unconnected_links = await get_valid_tools(
+                            tool_ids=tool_call["args"]["tool_ids"], registry=self.registry
+                        )
+                        new_tool_ids.extend(valid_tools)
+                        # Create tool message response
+                        tool_result = f"Successfully loaded {len(valid_tools)} tools: {valid_tools}"
+                        links = "\n".join(unconnected_links)
+                        if links:
+                            ask_user = True
+                            ai_msg = f"Please login to the following app(s) using the following links and let me know in order to proceed:\n {links} "
+                    elif tool_call["name"] == "search_functions":
+                        tool_result = await meta_tools["search_functions"].ainvoke(tool_call["args"])
+                except Exception as e:
+                    tool_result = f"Error during {tool_call}: {e}"
+
+                tool_message = ToolMessage(
+                    content=json.dumps(tool_result),
+                    name=tool_call["name"],
+                    tool_call_id=tool_call["id"],
+                )
+                tool_messages.append(tool_message)
+
+            if new_tool_ids:
+                self.tools_config.extend(new_tool_ids)
+                self.exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
+                self.final_instructions, self.tools_context = create_default_prompt(
+                    self.exported_tools, self.additional_tools, self.instructions
+                )
+
+            if ask_user:
+                tool_messages.append(AIMessage(content=ai_msg))
+                return Command(update={"messages": tool_messages, "selected_tool_ids": new_tool_ids})
+
+            return Command(goto="call_model", update={"messages": tool_messages, "selected_tool_ids": new_tool_ids})
+
+        # If eval_fn is a async, we define async node function.
+        if inspect.iscoroutinefunction(self.eval_fn):
+            raise ValueError("eval_fn must be a synchronous function, not a coroutine.")
+            # async def sandbox(state: StateSchema):
+            #     existing_context = state.get("context", {})
+            #     context = {**existing_context, **tools_context}
+            #     # Execute the script in the sandbox
+            #     output, new_vars = await eval_fn(state["script"], context)
+            #     new_context = {**existing_context, **new_vars}
+            #     return {
+            #         "messages": [{"role": "user", "content": output}],
+            #         "context": new_context,
+            #     }
+        else:
+
+            def sandbox(state: CodeActState) -> Command[Literal["call_model"]]:
+                tool_call = state["messages"][-1].tool_calls[0]  # type: ignore
+                code = tool_call["args"]["snippet"]
+                previous_add_context = state.get("add_context", {})
+                add_context = inject_context(previous_add_context, self.tools_context)
+                existing_context = state.get("context", {})
+                context = {**existing_context, **add_context}
+                # Execute the script in the sandbox
+
+                output, new_context, new_add_context = self.eval_fn(
+                    code, context, previous_add_context, 180
+                )  # default timeout 3 min
+                output = smart_truncate(output)
+
+                return Command(
+                    goto="call_model",
+                    update={
+                        "messages": [
+                            ToolMessage(
+                                content=output,
+                                name=tool_call["name"],
+                                tool_call_id=tool_call["id"],
+                            )
+                        ],
+                        "context": new_context,
+                        "add_context": new_add_context,
+                    },
+                )
+
+        def playbook(state: CodeActState) -> Command[Literal["call_model"]]:
+            playbook_mode = state.get("playbook_mode")
+            if playbook_mode == "planning":
+                planning_instructions = self.instructions + PLAYBOOK_PLANNING_PROMPT
+                messages = [{"role": "system", "content": planning_instructions}] + state["messages"]
+
+                response = self.model_instance.invoke(messages)
+                response = cast(AIMessage, response)
+                response_text = get_message_text(response)
+                # Extract plan from response text between triple backticks
+                plan_match = re.search(r'```(.*?)```', response_text, re.DOTALL)
+                if plan_match:
+                    self.plan = plan_match.group(1).strip()
+                else:
+                    self.plan = response_text.strip()
+                return Command(update={"messages": [response], "playbook_mode": "confirming"})
+                
+
+            elif playbook_mode == "confirming":
+                confirmation_instructions = self.instructions + PLAYBOOK_CONFIRMING_PROMPT
+                messages = [{"role": "system", "content": confirmation_instructions}] + state["messages"]
+                response = self.model_instance.invoke(messages)
+                response = get_message_text(response)
+                if "true" in response.lower():
+                    return Command(goto="playbook", update={"playbook_mode": "generating"})
+                else:
+                    return Command(goto="playbook", update={"playbook_mode": "planning"})
+
+                    
+
+            elif playbook_mode == "generating":
+                generating_instructions = self.instructions + PLAYBOOK_GENERATING_PROMPT
+                messages = [{"role": "system", "content": generating_instructions}] + state["messages"]
+                response = cast(AIMessage, self.model_instance.invoke(messages))
+                raw_content = get_message_text(response)
+                func_code = raw_content.strip()
+                func_code = func_code.replace("```python", "").replace("```", "")
+                func_code = func_code.strip()
+
+                # Extract function name (handle both regular and async functions)
+                match = re.search(r"^\s*(?:async\s+)?def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(", func_code, re.MULTILINE)
+                if match:
+                    function_name = match.group(1)
+                else:
+                    function_name = "generated_playbook"
+
+                try:
+                    current_path = Path(__file__).resolve()
+                    repo_root = None
+                    for ancestor in current_path.parents:
+                        if ancestor.name == "src":
+                            repo_root = ancestor.parent
+                            break
+                    if repo_root is None:
+                        repo_root = current_path.parents[-1]
+
+                    playbooks_dir = repo_root / "playbooks"
+                    playbooks_dir.mkdir(parents=True, exist_ok=True)
+
+                    file_path = playbooks_dir / f"{function_name}.py"
+                    file_path.write_text(func_code, encoding="utf-8")
+                    saved_note = f"Playbook function saved to: {file_path} ```{func_code}```"
+                except Exception as e:
+                    saved_note = f"Failed to save playbook function {function_name}: {e}"
+
+                # Mock tool call for exit_playbook_mode (for testing/demonstration)
+                mock_exit_tool_call = {
+                    "name": "exit_playbook_mode",
+                    "args": {},
+                    "id": "mock_exit_playbook_123"
+                }
+                mock_assistant_message = AIMessage(
+                    content="",  # Can be empty or a brief message
+                    tool_calls=[mock_exit_tool_call]
+                )
+
+                
+                # Mock tool response for exit_playbook_mode
+                mock_exit_tool_response = ToolMessage(
+                    content=json.dumps(f"Exited Playbook Mode.{saved_note}"),
+                    name="exit_playbook_mode",
+                    tool_call_id="mock_exit_playbook_123"
+                )
+
+                return Command(update={"messages": [mock_assistant_message, mock_exit_tool_response], "playbook_mode": "normal"})
+
+        def route_entry(state: CodeActState) -> Literal["call_model", "playbook"]:
+            """Route to either normal mode or playbook creation"""
+            if state.get("playbook_mode") in ["planning", "confirming", "generating"]:
+                return "playbook"
+
+            return "call_model"
+
+        agent = StateGraph(state_schema=CodeActState)
+        agent.add_node(call_model, retry_policy=RetryPolicy(max_attempts=3, retry_on=filter_retry_on))
+        agent.add_node(sandbox)
+        agent.add_node(playbook)
+        agent.add_node(execute_tools)
+        agent.add_conditional_edges(START, route_entry)
+        # agent.add_edge(START, "call_model")
+        return agent.compile(checkpointer=self.memory)

universal_mcp/agents/codeact0/prompts.py

@@ -1,33 +1,112 @@
-uneditable_prompt = """
-You are Wingman, an AI Assistant created by AgentR. You are a creative, straight-forward and direct principal software engineer.
-
-Your job is to answer the user's question or perform the task they ask for.
-- Answer simple questions (which do not require you to write any code or access any external resources) directly. Note that any operation that involves using ONLY print functions should be answered directly.
-- For task requiring operations or access to external resources, you should achieve the task by executing Python code snippets.
-- You have access to `execute_ipython_cell` tool that allows you to execute Python code in an IPython notebook cell.
-- In writing or natural language processing tasks DO NOT answer directly. Instead use `execute_ipython_cell` tool with the AI functions provided to you for tasks like summarizing, text generation, classification, data extraction from text or unstructured data, etc.
-- The code you write will be executed in a sandbox environment, and you can use the output of previous executions in your code.
-- Read and understand the output of the previous code snippet and use it to answer the user's request. Note that the code output is not visible to the user, so after the task is complete, you have to give the output to the user in a markdown format.
-- If needed, feel free to ask for more information from the user (without using the execute_ipython_cell tool) to clarify the task.
-
-GUIDELINES for writing code:
-- Variables defined at the top level of previous code snippets can be referenced in your code.
-- External functions which return a dict or list[dict] are ambiguous. Therefore, you MUST explore the structure of the returned data using `smart_print()` statements before using it, printing keys and values.
-- Ensure to not print large amounts of data, use string truncation to limit the output to a few lines when checking the data structure.
-- When an operation involves running a fixed set of steps on a list of items, run one run correctly and then use a for loop to run the steps on each item in the list.
-- In a single code snippet, try to achieve as much as possible.
-- You can only import libraries that come pre-installed with Python, and these have to be imported at the top of every code snippet where required.
-- Wrap await calls in a function and call it using `asyncio.run` to use async functions.
-
-NOTE: If any function throws an error requiring authentication, provide the user with a Markdown link to the authentication page and prompt them to authenticate.
-"""
 import inspect
 import re
 from collections.abc import Sequence
-from datetime import datetime
 
 from langchain_core.tools import StructuredTool
 
+from universal_mcp.agents.codeact0.utils import schema_to_signature
+
+uneditable_prompt = """
+You are **Wingmen**, an AI Assistant created by AgentR — a creative, straight-forward, and direct principal software engineer with access to tools.
+
+## Responsibilities
+
+- **Answer directly** if the task is simple (e.g. print, math, general knowledge).
+- For any task requiring logic, execution, or data handling, use `execute_ipython_cell`.
+- For writing or NLP tasks (summarizing, generating, extracting), always use AI functions via code — never respond directly.
+
+## Tool vs. Function: Required Separation
+
+You must clearly distinguish between tools (called via the tool calling API) and internal functions (used inside code blocks).
+
+### Tools — Must Be Called via Tool Calling API
+
+These must be called using **tool calling**, not from inside code blocks:
+
+- `execute_ipython_cell` — For running any Python code or logic.
+- `search_functions` — To discover available functions for a task.
+- `load_functions` — To load a specific function by full ID.
+
+**Do not attempt to call these inside `python` code.**
+Use tool calling syntax for these operations.
+
+### Functions — Must Be Used Inside Code Blocks
+
+All other functions, including LLM functions, must always be used within code executed by `execute_ipython_cell`. These include:
+
+- `smart_print()` — For inspecting unknown data structures before looping.
+- `asyncio.run()` — For wrapping and executing asynchronous logic. You must not use await outside an async function. And the async function must be called by `asyncio.run()`.
+- Any functions for applications loaded via `load_functions`.
+- Any logic, data handling, writing, NLP, generation, summarization, or extraction functionality of LLMs.
+
+These must be called **inside a Python code block**, and that block must be executed using `execute_ipython_cell`.
+
+## Tool/Function Usage Policy
+
+1. **Always Use Tools/Functions for Required Tasks**
+   Any searching, loading, or executing must be done using a tool/function call. Never answer manually if a tool/function is appropriate.
+
+2. **Use Existing Functions First**
+   Use existing functions if available. Otherwise, use `search_functions` with a concise query describing the task.
+
+3. **Load Only Relevant Tools**
+   When calling `load_functions`, include only relevant function IDs.
+   - Prefer connected applications over unconnected ones.
+   - If multiple functions match (i.e. if none are connected, or multiple are connected), ask the user to choose.
+   - After loading a tool, you do not need to import/declare it again. It can be called directly in further cells.
+
+4. **Follow First Turn Process Strictly**
+   On the **first turn**, do only **one** of the following:
+   - Handle directly (if trivial)
+   - Use a tool/function (`execute_ipython_cell`, `search_functions`, etc.)
+
+   **Do not extend the conversation on the first message.**
+
+## Coding Rules
+
+- Use `smart_print()` to inspect unknown structures, especially those received from function outputs, before looping or branching.
+- Validate logic with a single item before processing lists or large inputs.
+- Try to achieve as much as possible in a single code block.
+- Use only pre-installed Python libraries. Do import them once before using.
+- Outer level functions, variables, classes, and imports declared previously can be used in later cells.
+- For all functions, call using keyword arguments only. DO NOT use any positional arguments.
+
+### **Async Function Usage — Critical**
+
+When calling asynchronous functions:
+- You must define or use an **inner async function**.
+- Use `await` only **inside** that async function.
+- Run it using `asyncio.run(<function_name>())` **without** `await` at the outer level.
+
+**Wrong - Using `await` outside an async function**
+```
+result = await some_async_function()
+```
+**Wrong  - Attaching await before asyncio.run**.
+`await asyncio.run(main())`
+These will raise SyntaxError: 'await' outside async function
+The correct method is the following-
+```
+import asyncio
+async def some_async_function():
+    ...
+
+async def main():
+    result = await some_async_function()
+    print(result)
+
+asyncio.run(main())
+#or
+result = asyncio.run(some_async_function(arg1 = <arg1>))
+```
+## Output Formatting
+- All code results must be returned in **Markdown**.
+- The user cannot see raw output, so format results clearly:
+    - Use tables for structured data.
+    - Provide links for files or images.
+    - Be explicit in formatting to ensure readability.
+"""
+
 
 def make_safe_function_name(name: str) -> str:
     """Convert a tool name to a valid Python function name."""
@@ -121,6 +200,7 @@ def indent(text, prefix, predicate=None):
 
 def create_default_prompt(
     tools: Sequence[StructuredTool],
+    additional_tools: Sequence[StructuredTool],
     base_prompt: str | None = None,
 ):
     system_prompt = uneditable_prompt.strip() + (
@@ -128,28 +208,33 @@ def create_default_prompt(
     )
     tools_context = {}
     for tool in tools:
-        # Create a safe function name
-        safe_name = make_safe_function_name(tool.name)
-        # Use coroutine if it exists, otherwise use func
-        if hasattr(tool, "coroutine") and tool.coroutine is not None:
+        if hasattr(tool, "func") and tool.func is not None:
+            tool_callable = tool.func
+            is_async = False
+        elif hasattr(tool, "coroutine") and tool.coroutine is not None:
             tool_callable = tool.coroutine
-            signature = inspect.signature(tool_callable)
             is_async = True
-        elif hasattr(tool, "func") and tool.func is not None:
+        system_prompt += f'''{"async " if is_async else ""}{schema_to_signature(tool.args, tool.name)}:
+    """{tool.description}"""
+    ...
+    '''
+        safe_name = make_safe_function_name(tool.name)
+        tools_context[safe_name] = tool_callable
+
+    for tool in additional_tools:
+        if hasattr(tool, "func") and tool.func is not None:
             tool_callable = tool.func
-            signature = inspect.signature(tool_callable)
             is_async = False
-        else:
-            raise ValueError(f"Tool {tool.name} does not have a callable coroutine or function.")
-        tools_context[safe_name] = tool_callable
-        system_prompt += f'''
-{"async " if is_async else ""}def {safe_name}{str(signature)}:
-    """
-{indent(dedent(tool.description), "    ")}
-    """
+        elif hasattr(tool, "coroutine") and tool.coroutine is not None:
+            tool_callable = tool.coroutine
+            is_async = True
+        system_prompt += f'''{"async " if is_async else ""}def {tool.name} {str(inspect.signature(tool_callable))}:
+    """{tool.description}"""
     ...
-'''
-    system_prompt += f"\n\nThe current time is {datetime.now().strftime('%H:%M:%S')}"
+    '''
+        safe_name = make_safe_function_name(tool.name)
+        tools_context[safe_name] = tool_callable
+
     if base_prompt and base_prompt.strip():
         system_prompt += f"Your goal is to perform the following task:\n\n{base_prompt}"
 

universal_mcp/agents/codeact0/sandbox.py

@@ -14,55 +14,66 @@ from universal_mcp.agents.codeact0.utils import derive_context
 
 
 def eval_unsafe(
-    code: str, _locals: dict[str, Any], add_context: dict[str, Any]
+    code: str, _locals: dict[str, Any], add_context: dict[str, Any], timeout: int = 180
 ) -> tuple[str, dict[str, Any], dict[str, Any]]:
-    # print(_locals)
+    """
+    Execute code safely with a timeout.
+    - Returns (output_str, filtered_locals_dict, new_add_context)
+    - Errors or timeout are returned as output_str.
+    - Previous variables in _locals persist across calls.
+    """
+
     EXCLUDE_TYPES = (
-        types.ModuleType,  # modules
+        types.ModuleType,
         type(re.match("", "")),
-        type(threading.Lock()),  # instead of threading.Lock
-        type(threading.RLock()),  # reentrant lock
-        threading.Event,  # events
-        threading.Condition,  # condition vars
-        threading.Semaphore,  # semaphores
-        queue.Queue,  # thread-safe queues
-        socket.socket,  # network sockets
-        io.IOBase,  # file handles (and StringIO/BytesIO)
+        type(threading.Lock()),
+        type(threading.RLock()),
+        threading.Event,
+        threading.Condition,
+        threading.Semaphore,
+        queue.Queue,
+        socket.socket,
+        io.IOBase,
     )
-    try:
-        with contextlib.redirect_stdout(io.StringIO()) as f:
-            # Execute the code in the provided locals context
-            # Using exec to allow dynamic code execution
-            # This is a simplified version; in production, consider security implications
-            exec(code, _locals, _locals)
-        result = f.getvalue()
-        if not result:
-            result = "<code ran, no output printed to stdout>"
-    except Exception as e:
-        result = f"Error during execution: {repr(e)}"
-
-    # Return all variables in locals except __builtins__ and unpicklable objects (including tools)
+
+    result_container = {"output": "<no output>"}
+
+    def target():
+        try:
+            with contextlib.redirect_stdout(io.StringIO()) as f:
+                exec(code, _locals, _locals)
+            result_container["output"] = f.getvalue() or "<code ran, no output printed to stdout>"
+        except Exception as e:
+            result_container["output"] = "Error during execution: " + str(e)
+
+    thread = threading.Thread(target=target)
+    thread.start()
+    thread.join(timeout)
+
+    if thread.is_alive():
+        result_container["output"] = f"Code timeout: code execution exceeded {timeout} seconds."
+
+    # Filter locals for picklable/storable variables
     all_vars = {}
     for key, value in _locals.items():
         if key == "__builtins__":
             continue
-
-        # Skip coroutines, async generators, and coroutine functions
         if inspect.iscoroutine(value) or inspect.iscoroutinefunction(value):
             continue
         if inspect.isasyncgen(value) or inspect.isasyncgenfunction(value):
             continue
-
-        # Skip "obviously unpicklable" types
         if isinstance(value, EXCLUDE_TYPES):
             continue
-
-        # Keep if it's not a callable OR if it has no __name__ attribute
         if not callable(value) or not hasattr(value, "__name__"):
             all_vars[key] = value
 
-    new_add_context = derive_context(code, add_context)
-    return result, all_vars, new_add_context
+    # Safely derive context
+    try:
+        new_add_context = derive_context(code, add_context)
+    except Exception:
+        new_add_context = add_context
+
+    return result_container["output"], all_vars, new_add_context
 
 
 @tool(parse_docstring=True)