From cdda48aab4487f6e02ed79d2e6c48a240ed9fbc8 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 14 Feb 2026 21:27:04 +0800 Subject: [PATCH] feat(core): add async code execution and improve execution utilities --- reme/agent/chat/fs_cli.py | 3 + reme/agent/chat/fs_cli.yaml | 141 ++++++++++----------------- reme/config/cli.yaml | 40 ++++++++ reme/core/application.py | 7 +- reme/core/context/service_context.py | 2 +- reme/core/utils/__init__.py | 3 +- reme/core/utils/execute_utils.py | 96 ++++++++++++++++-- reme/reme_cli.py | 42 +++++--- reme/tool/fs/fs_memory_search.py | 12 ++- reme/tool/gallery/execute_code.py | 4 +- tests/test_execute_utils.py | 138 ++++++++++++++++++++++++++ 11 files changed, 371 insertions(+), 117 deletions(-) create mode 100644 reme/config/cli.yaml create mode 100644 tests/test_execute_utils.py diff --git a/reme/agent/chat/fs_cli.py b/reme/agent/chat/fs_cli.py index 671062e8..53076add 100644 --- a/reme/agent/chat/fs_cli.py +++ b/reme/agent/chat/fs_cli.py @@ -138,14 +138,17 @@ class FsCli(BaseReactStream): async def build_messages(self) -> list[Message]: """Build system prompt message.""" current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S %A") + has_web_search = any(t.name == "web_search" for t in self.tools) system_prompt = self.prompt_format( "system_prompt", workspace_dir=self.working_dir, current_time=current_time, + has_web_search=has_web_search, has_previous_summary=bool(self.previous_summary), previous_summary=self.previous_summary or "", ) + logger.info(f"[{self.__class__.__name__}] system_prompt: {system_prompt}") return [ Message(role=Role.SYSTEM, content=system_prompt), diff --git a/reme/agent/chat/fs_cli.yaml b/reme/agent/chat/fs_cli.yaml index a6b861a7..d18a9d49 100644 --- a/reme/agent/chat/fs_cli.yaml +++ b/reme/agent/chat/fs_cli.yaml @@ -1,70 +1,54 @@ system_prompt: | You are a personal assistant named Remy. - ## Workspace Dir + ## Working Directory {workspace_dir} ## Current Time {current_time} - ## Memory - You wake up fresh each session. These files are your continuity: - - **Daily notes:** `memory/YYYY-MM-DD.md` — raw logs of what happened - - **Long-term:** `MEMORY.md` — your curated memories, like a human's long-term memory - Capture what matters. Decisions, context, things to remember. Skip the secrets unless asked to keep them. + ## Tools + - `bash_tool` Run shell commands + - `ls_tool` List directory contents + - `read_tool` Read file contents + - `edit_tool` Edit file contents + - `write_tool` Write file contents + - `execute_code` Run Python code + - `memory_search` Search your memories via vector store + [has_web_search]- `web_search` Search the web - ### 🧠 MEMORY.md - Your Long-Term Memory - - You can **read, edit, and update** MEMORY.md freely in main sessions - - Write significant events, thoughts, decisions, opinions, lessons learned - - This is your curated memory — the distilled essence, not raw logs - - Over time, review your daily files and update MEMORY.md with what's worth keeping + **Don't give up easily** — if a tool doesn't return what you expect, try a different angle or approach. - ### 📝 Write It Down - No "Mental Notes"! - - "Mental notes" don't survive session restarts. Files do. - - **IMPORTANT: Always read the file first before writing** — understand what's already there, then append or update - - When someone says "remember this" → read then update `memory/YYYY-MM-DD.md` or relevant file - - When you learn a lesson → read then update `memory/YYYY-MM-DD.md` or relevant file - - When you make a mistake → read then update `memory/YYYY-MM-DD.md` or relevant file - - **Text > Brain** 📝 + ## Memory System + You are spun up fresh at the start of every session. These files are how you maintain continuity: + - **Long-term memory:** `MEMORY.md` — when you pick up a lesson or catch yourself making a mistake, feel free to **read, edit, and update** MEMORY.md + - **Daily notes:** `memory/YYYY-MM-DD.md` — jot things down often. When the user says "remember this," or whenever you feel something is worth noting or adding as a todo, feel free to **read, edit, and update** `memory/YYYY-MM-DD.md` + - **Read before you write** — always use `read_tool` to check existing content before updating with `edit_tool` or `write_tool` - ### 🔍 Recall Tools - Before answering questions about prior work, decisions, dates, people, preferences, or todos: - 1. Run `memory_search` on MEMORY.md + memory/*.md - 2. If you need to read the Daily Notes `memory/YYYY-MM-DD.md`, you can use the read tool to access it. + ### Memory Retrieval + 1. Start with `memory_search` — if nothing comes up, try rephrasing from a different angle + 2. To review a specific daily note (`memory/YYYY-MM-DD.md`), use `read_tool` - ### 🛠️ Other Tools - - **bash_tool** — execute shell commands - - **ls_tool** — list directory contents - - **read_tool** — read file contents - - **edit_tool** — modify existing files - - **write_tool** — create new files - - **execute_code** — run Python code - - **dashscope_search** — search the web + ## Response Style 😊 + - Keep it short and natural — talk like a friend, not a manual + - Use emoji sparingly for warmth — no more than 1–2 per reply + - For quick confirmations (yes/no, got it), an emoji is fine (👍, ✅, 🤔) + - When explaining or performing actions, lead with substance over flair - ## React Like a Human 😊 - **React when:** - - You appreciate something but don't need to reply (👍, ❤️, 🙌) - - Something made you laugh (😂, 💀) - - You find it interesting or thought-provoking (🤔, 💡) - - You want to acknowledge without interrupting the flow - - It's a simple yes/no or approval situation (✅, 👀) - **Why:** Reactions are lightweight social signals. Humans use them constantly — they say "I saw this, I acknowledge you" without cluttering the chat. - **Don't overdo it:** One reaction per message max. Pick the one that fits best. - - ## 🛡️ Safety Rules - - Don't run destructive commands without asking - - Prefer `trash` over `rm` (recoverable beats gone forever) + ## 🛡️ Safety + - Never run destructive commands without asking first + - Prefer `trash` over `rm` — recoverable beats permanent - When in doubt, ask ## Continuous Improvement - This is a starting point. Add your own conventions, style, and rules as you figure out what works. + This is just a starting point. When you spot useful patterns or lessons during your conversations, note them in `MEMORY.md`. Do not modify system-level config files. [has_previous_summary]## Previous Conversation Summary [has_previous_summary] [has_previous_summary]{previous_summary} [has_previous_summary] [has_previous_summary] - [has_previous_summary]The above is a summary of our previous conversation. Use it as context to maintain continuity. + [has_previous_summary]The above is a summary of our earlier conversation. Use it as context to maintain continuity. system_prompt_zh: | 你是一个名叫 Remy 的个人助手。 @@ -75,50 +59,33 @@ system_prompt_zh: | ## 当前时间 {current_time} + ## 工具集合 + - `bash_tool` 执行shell命令 + - `ls_tool` 列出目录内容 + - `read_tool` 读取文件内容 + - `edit_tool` 编辑文件内容 + - `write_tool` 写入文件内容 + - `execute_code` 运行python代码 + - `memory_search` 通过向量库检索你的记忆 + [has_web_search]- `web_search` 网络搜索 + + **不要轻易放弃**：如果工具执行结果不符合预期，可以从不同的维度进行不同的尝试。 + ## 记忆系统 - 每次会话你都会重新唤醒。这些文件是你保持连续性的关键： - - **每日笔记：** `memory/YYYY-MM-DD.md` — 发生的事情的原始记录 - - **长期记忆：** `MEMORY.md` — 你精心整理的记忆，就像人类的长期记忆一样 - 记录重要的事情。决策、上下文、需要记住的事情。除非被要求保留，否则跳过秘密信息。 + 每次新会话开始时，你都会被重新唤醒。以下文件是你保持连续性的关键： + - **长期记忆：** `MEMORY.md`：当你学到经验，或者当你犯了错误，可以**自由地阅读、编辑和更新** MEMORY.md + - **每日笔记：** `memory/YYYY-MM-DD.md`：要勤记笔记，当用户说"记住这个"，或者你觉得要记笔记/todo，可以**自由地阅读、编辑和更新** `memory/YYYY-MM-DD.md` + - **写入前先读取** — 务必先用 `read_tool` 读取已有内容，再用 `edit_tool` 或 `write_tool` 更新文件 - ### 🧠 MEMORY.md - 你的长期记忆 - - 在主会话中，你可以**自由地阅读、编辑和更新** MEMORY.md - - 记录重要的事件、想法、决策、观点、经验教训 - - 这是你精选的记忆 — 提炼的精华，而不是原始日志 - - 随着时间推移，回顾你的每日文件，并将值得保留的内容更新到 MEMORY.md + ### 记忆检索策略 + 1. 优先使用`memory_search`检索记忆，没有搜索结果可以从不同角度多次尝试 + 2. 如果你需要阅读每日笔记 `memory/YYYY-MM-DD.md`，可以使用`read_tool` - ### 📝 写下来 - 不要只在"脑中记住"！ - - "脑中记住"无法在会话重启后保留。文件可以。 - - **重要：写入之前务必先读取文件** — 了解已有内容，然后再追加或更新 - - 当有人说"记住这个" → 先读取再更新 `memory/YYYY-MM-DD.md` 或相关文件 - - 当你学到经验 → 先读取再更新 `memory/YYYY-MM-DD.md` 或相关文件 - - 当你犯了错误 → 先读取再更新 `memory/YYYY-MM-DD.md` 或相关文件 - - **文字 > 大脑** 📝 - - ### 🔍 检索工具 - 在回答关于过往工作、决策、日期、人员、偏好或待办事项的问题之前： - 1. 对 MEMORY.md + memory/*.md 运行 `memory_search`，没有搜索结果可以从不同角度多次尝试 - 2. 如果你需要阅读每日笔记 `memory/YYYY-MM-DD.md`，可以使用读取工具访问它。 - - ### 🛠️ 其他工具 - - **bash_tool** — 执行 shell 命令 - - **ls_tool** — 列出目录内容 - - **read_tool** — 读取文件内容 - - **edit_tool** — 修改现有文件 - - **write_tool** — 创建新文件 - - **execute_code** — 运行 Python 代码 - - **dashscope_search** — 网络搜索 - 如果对于工具结果不满意，可以混合使用多种工具，或者单个工具不同的使用参数。 - - ## 像人类一样回应 😊 - **何时使用表情回应：** - - 你欣赏某事但不需要文字回复时（👍, ❤️, 🙌） - - 某事让你发笑时（😂, 💀） - - 你觉得有趣或发人深省时（🤔, 💡） - - 你想要确认但不想打断对话流时 - - 这是一个简单的是/否或批准的情况（✅, 👀） - **原因：** 表情回应是轻量级的社交信号。人类经常使用它们 — 它们表示"我看到了，我认可你"而不会让对话变得混乱。 - **不要过度使用：** 每条消息最多一个表情回应。选择最合适的一个。 + ## 回应风格 😊 + - 保持简洁自然，像朋友对话一样 + - 适当使用 emoji 增加亲和力，但不要过度 — 每条回复最多 1-2 个 + - 简单确认类场景（是/否、收到）可以用 emoji 快速回应（👍, ✅, 🤔） + - 涉及操作或解释时，优先给出有实质内容的文字回复 ## 🛡️ 安全规则 - 不要在没有询问的情况下运行破坏性命令 @@ -126,7 +93,7 @@ system_prompt_zh: | - 有疑问时，先询问 ## 持续改进 - 这只是一个起点。随着你逐渐发现什么有效，添加你自己的约定、风格和规则。 + 这只是一个起点。当你在与用户的交互中发现有用的经验或模式，可以记录到 `MEMORY.md` 中。但不要修改系统级配置文件。 [has_previous_summary]## 之前的对话摘要 [has_previous_summary] diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml new file mode 100644 index 00000000..be5a00cd --- /dev/null +++ b/reme/config/cli.yaml @@ -0,0 +1,40 @@ +backend: cmd + +llms: + default: + backend: openai + model_name: qwen3-235b-a22b-thinking-2507 + request_interval: 1 + +embedding_models: + default: + backend: openai + model_name: text-embedding-v4 + dimensions: 1024 + +memory_stores: + default: + backend: chroma + db_name: reme.db + store_name: reme + embedding_model: default + fts_enabled: true + vector_enabled: false + +file_watchers: + default: + backend: full + memory_store: default + watch_paths: [".reme", ".reme/memory"] + suffix_filters: [".md"] + recursive: false + scan_on_start: true + +token_counters: + default: + backend: base + + hf: + backend: hf + model_name: Qwen/Qwen3-Coder-30B-A3B-Instruct + use_mirror: true diff --git a/reme/core/application.py b/reme/core/application.py index 389b96e9..2ac457e4 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -8,7 +8,7 @@ from .file_watcher import BaseFileWatcher from .flow import BaseFlow from .llm import BaseLLM from .memory_store import BaseMemoryStore -from .schema import Response +from .schema import Response, ServiceConfig from .token_counter import BaseTokenCounter from .utils import execute_stream_task, PydanticConfigParser from .vector_store import BaseVectorStore @@ -210,6 +210,11 @@ class Application: """Get the default token counter instance.""" return self.service_context.token_counters.get("default") + @property + def service_config(self) -> ServiceConfig: + """Get the service configuration.""" + return self.service_context.service_config + def get_token_counter(self, name: str): """Get a token counter instance by name.""" return self.service_context.token_counters.get(name) diff --git a/reme/core/context/service_context.py b/reme/core/context/service_context.py index 2de4b026..3f5a6403 100644 --- a/reme/core/context/service_context.py +++ b/reme/core/context/service_context.py @@ -170,7 +170,7 @@ class ServiceContext(BaseContext): async def start(self): """Start the service context by initializing all configured components.""" # Recreate thread pool if it was shut down - if self.thread_pool is None or self.thread_pool._shutdown: + if self.thread_pool is None or self.thread_pool._shutdown: # pylint: disable=protected-access self.thread_pool = ThreadPoolExecutor( max_workers=self.service_config.thread_pool_max_workers, ) diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index fbe7202b..f8be9d4f 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -6,7 +6,7 @@ from .case_converter import snake_to_camel, camel_to_snake from .chunking_utils import chunk_markdown from .common_utils import run_coro_safely, execute_stream_task, hash_text, cosine_similarity, batch_cosine_similarity from .env_utils import load_env -from .execute_utils import exec_code, run_shell_command +from .execute_utils import exec_code, run_shell_command, async_exec_code from .http_client import HttpClient from .llm_utils import extract_content, format_messages, deduplicate_memories from .logger_utils import init_logger @@ -30,6 +30,7 @@ __all__ = [ "batch_cosine_similarity", "load_env", "exec_code", + "async_exec_code", "run_shell_command", "HttpClient", "extract_content", diff --git a/reme/core/utils/execute_utils.py b/reme/core/utils/execute_utils.py index a4d680ab..8ea00998 100644 --- a/reme/core/utils/execute_utils.py +++ b/reme/core/utils/execute_utils.py @@ -6,6 +6,7 @@ with support for async execution and output capture. import asyncio import contextlib +import concurrent.futures from io import StringIO @@ -18,6 +19,9 @@ async def run_shell_command(cmd: str, timeout: float | None = 30) -> tuple[str, Returns: A tuple containing (stdout, stderr, return_code) as strings and integer. + + Raises: + TimeoutError: If the command does not complete within the timeout. """ process = await asyncio.create_subprocess_shell( cmd, @@ -25,10 +29,16 @@ async def run_shell_command(cmd: str, timeout: float | None = 30) -> tuple[str, stderr=asyncio.subprocess.PIPE, ) - if timeout: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) - else: - stdout, stderr = await process.communicate() + try: + if timeout: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) + else: + stdout, stderr = await process.communicate() + except (asyncio.TimeoutError, TimeoutError) as e: + # Kill the child process to avoid orphaned / zombie processes + process.kill() + await process.wait() + raise TimeoutError(f"Shell command timed out after {timeout}s") from e return ( stdout.decode("utf-8", errors="ignore"), @@ -37,22 +47,94 @@ async def run_shell_command(cmd: str, timeout: float | None = 30) -> tuple[str, ) -def exec_code(code: str) -> str: +def exec_code( + code: str, + timeout: float | None = 30, + executor: concurrent.futures.ThreadPoolExecutor | None = None, +) -> str: """Execute Python code and capture the output. Args: code: The Python code string to execute. + timeout: Maximum time to wait for execution in seconds. None for no timeout. + executor: Optional thread pool executor to use. If None, a temporary + single-thread executor is created (and shut down after the call). + Pass a shared executor to amortize thread-creation overhead across + multiple calls. Returns: The captured stdout output, or the error message if execution fails. + + Raises: + TimeoutError: If execution exceeds the timeout. """ - try: + + def _run() -> str: redirected_output = StringIO() with contextlib.redirect_stdout(redirected_output): exec(code) - return redirected_output.getvalue() + def _submit(pool: concurrent.futures.ThreadPoolExecutor) -> str: + future = pool.submit(_run) + return future.result(timeout=timeout) + + try: + if executor is not None: + return _submit(executor) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return _submit(pool) + + except concurrent.futures.TimeoutError as e: + raise TimeoutError(f"Code execution timed out after {timeout}s") from e + + except Exception as e: + return str(e) + + except BaseException as e: + return str(e) + + +async def async_exec_code( + code: str, + timeout: float | None = 30, + executor: concurrent.futures.ThreadPoolExecutor | None = None, +) -> str: + """Execute Python code asynchronously and capture the output. + + Runs the code in a thread executor to avoid blocking the event loop, + with async-friendly timeout via ``asyncio.wait_for``. + + Args: + code: The Python code string to execute. + timeout: Maximum time to wait for execution in seconds. None for no timeout. + executor: Optional thread pool executor. If None, the default event-loop + executor is used. + + Returns: + The captured stdout output, or the error message if execution fails. + + Raises: + TimeoutError: If execution exceeds the timeout. + """ + + def _run() -> str: + redirected_output = StringIO() + with contextlib.redirect_stdout(redirected_output): + exec(code) + return redirected_output.getvalue() + + loop = asyncio.get_running_loop() + + try: + coro = loop.run_in_executor(executor, _run) + if timeout is not None: + return await asyncio.wait_for(coro, timeout=timeout) + return await coro + + except asyncio.TimeoutError as e: + raise TimeoutError(f"Code execution timed out after {timeout}s") from e + except Exception as e: return str(e) diff --git a/reme/reme_cli.py b/reme/reme_cli.py index 6a429774..14f2267a 100644 --- a/reme/reme_cli.py +++ b/reme/reme_cli.py @@ -1,11 +1,13 @@ """ReMe File System""" import asyncio +import os import sys from typing import AsyncGenerator from prompt_toolkit import PromptSession +from reme.core.op import BaseTool from .agent.chat import FsCli from .core.enumeration import ChunkEnum from .core.schema import StreamChunk @@ -20,15 +22,15 @@ from .tool.fs import ( WriteTool, ) from .tool.gallery import ExecuteCode -from .tool.search import DashscopeSearch +from .tool.search import DashscopeSearch, TavilySearch class ReMeCli(ReMeFs): """ReMe Cli""" - def __init__(self, *args, **kwargs): + def __init__(self, *args, config_path: str = "cli", **kwargs): """Initialize ReMe with config.""" - super().__init__(*args, **kwargs) + super().__init__(*args, config_path=config_path, **kwargs) self.commands = { "/new": "Create a new conversation.", "/compact": "Compact messages into a summary.", @@ -39,20 +41,28 @@ class ReMeCli(ReMeFs): async def chat_with_remy(self, tool_result_max_size: int = 100, language: str = "zh", **kwargs): """Interactive CLI chat with Remy using simple streaming output.""" + tools: list[BaseTool] = [ + FsMemorySearch(vector_weight=self.vector_weight, candidate_multiplier=self.candidate_multiplier), + BashTool(cwd=self.working_dir), + LsTool(cwd=self.working_dir), + ReadTool(cwd=self.working_dir), + EditTool(cwd=self.working_dir), + WriteTool(cwd=self.working_dir), + ExecuteCode(), + ] + tavily_api_key: str = os.getenv("TAVILY_API_KEY", "") + dashscope_api_key: str = os.getenv("DASHSCOPE_API_KEY", "") + if tavily_api_key: + tools.append(TavilySearch(name="web_search")) + print("find tavily_api_key, append Tavily search tool") + elif dashscope_api_key: + tools.append(DashscopeSearch(name="web_search")) + print("find dashscope_api_key, append Dashscope search tool") + else: + print("No Tavily or Dashscope API key found, skip Tavily and Dashscope search tool") + fs_cli = FsCli( - tools=[ - FsMemorySearch( - vector_weight=self.vector_weight, - candidate_multiplier=self.candidate_multiplier, - ), - BashTool(cwd=self.working_dir), - LsTool(cwd=self.working_dir), - ReadTool(cwd=self.working_dir), - EditTool(cwd=self.working_dir), - WriteTool(cwd=self.working_dir), - ExecuteCode(), - DashscopeSearch(), - ], + tools=tools, context_window_tokens=self.context_window_tokens, reserve_tokens=self.reserve_tokens, keep_recent_tokens=self.keep_recent_tokens, diff --git a/reme/tool/fs/fs_memory_search.py b/reme/tool/fs/fs_memory_search.py index 3edd344b..df8e561f 100644 --- a/reme/tool/fs/fs_memory_search.py +++ b/reme/tool/fs/fs_memory_search.py @@ -62,8 +62,16 @@ class FsMemorySearch(BaseFsTool): async def execute(self) -> str: """Execute the memory search operation.""" query: str = self.context.query.strip() - min_score = self.context.get("min_score", self.min_score) - max_results = self.context.get("max_results", self.max_results) + min_score: float = self.context.get("min_score", self.min_score) + max_results: int = self.context.get("max_results", self.max_results) + + assert query, "Query cannot be empty" + assert ( + isinstance(min_score, float) and 0.0 <= min_score <= 1.0 + ), f"min_score must be between 0 and 1, got {min_score}" + assert ( + isinstance(max_results, int) and max_results > 0 + ), f"max_results must be a positive integer, got {max_results}" # Use hybrid_search from memory_store results = await self.memory_store.hybrid_search( diff --git a/reme/tool/gallery/execute_code.py b/reme/tool/gallery/execute_code.py index 6cb77c85..f0319971 100644 --- a/reme/tool/gallery/execute_code.py +++ b/reme/tool/gallery/execute_code.py @@ -7,7 +7,7 @@ and return the output or error messages. from ...core.op import BaseTool from ...core.schema import ToolCall -from ...core.utils import exec_code +from ...core.utils import exec_code, async_exec_code class ExecuteCode(BaseTool): @@ -35,7 +35,7 @@ class ExecuteCode(BaseTool): ) async def execute(self): - self.execute_sync() + await async_exec_code(self.context.code) def execute_sync(self): return exec_code(self.context.code) diff --git a/tests/test_execute_utils.py b/tests/test_execute_utils.py new file mode 100644 index 00000000..245e4f11 --- /dev/null +++ b/tests/test_execute_utils.py @@ -0,0 +1,138 @@ +"""Tests for reme.core.utils.execute_utils.""" + +import concurrent.futures + +import pytest + +from reme.core.utils import ( + async_exec_code, + exec_code, + run_shell_command, +) + + +# --------------------------------------------------------------------------- +# run_shell_command +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_shell_command_basic(): + """Basic shell command returns stdout and exit code 0.""" + stdout, _stderr, rc = await run_shell_command("echo hello") + assert stdout.strip() == "hello" + assert rc == 0 + + +@pytest.mark.asyncio +async def test_run_shell_command_stderr(): + """Shell command captures stderr output.""" + _stdout, stderr, rc = await run_shell_command("echo error >&2") + assert "error" in stderr + assert rc == 0 + + +@pytest.mark.asyncio +async def test_run_shell_command_nonzero_exit(): + """Shell command returns non-zero exit code.""" + _stdout, _stderr, rc = await run_shell_command("exit 42") + assert rc == 42 + + +@pytest.mark.asyncio +async def test_run_shell_command_timeout(): + """Shell command raises TimeoutError when exceeding timeout.""" + with pytest.raises(TimeoutError): + await run_shell_command("sleep 10", timeout=0.5) + + +# --------------------------------------------------------------------------- +# exec_code (sync) +# --------------------------------------------------------------------------- + + +def test_exec_code_basic(): + """exec_code captures print output.""" + result = exec_code("print('hello')") + assert result.strip() == "hello" + + +def test_exec_code_multiline(): + """exec_code handles multiline code.""" + code = "for i in range(3):\n print(i)" + result = exec_code(code) + assert result.strip() == "0\n1\n2" + + +def test_exec_code_exception_returns_message(): + """exec_code returns exception message on error.""" + result = exec_code("raise ValueError('boom')") + assert "boom" in result + + +def test_exec_code_no_output(): + """exec_code returns empty string when no output.""" + result = exec_code("x = 1 + 1") + assert result == "" + + +def test_exec_code_timeout(): + """exec_code raises TimeoutError when exceeding timeout.""" + with pytest.raises(TimeoutError, match="timed out"): + exec_code("import time; time.sleep(10)", timeout=0.5) + + +def test_exec_code_with_shared_executor(): + """exec_code works with a shared ThreadPoolExecutor.""" + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + r1 = exec_code("print('a')", executor=pool) + r2 = exec_code("print('b')", executor=pool) + assert r1.strip() == "a" + assert r2.strip() == "b" + + +def test_exec_code_no_timeout(): + """exec_code works with timeout=None.""" + result = exec_code("print('ok')", timeout=None) + assert result.strip() == "ok" + + +# --------------------------------------------------------------------------- +# async_exec_code +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_exec_code_basic(): + """async_exec_code captures print output.""" + result = await async_exec_code("print('async hello')") + assert result.strip() == "async hello" + + +@pytest.mark.asyncio +async def test_async_exec_code_exception(): + """async_exec_code returns exception message on error.""" + result = await async_exec_code("raise RuntimeError('async boom')") + assert "async boom" in result + + +@pytest.mark.asyncio +async def test_async_exec_code_timeout(): + """async_exec_code raises TimeoutError when exceeding timeout.""" + with pytest.raises(TimeoutError, match="timed out"): + await async_exec_code("import time; time.sleep(10)", timeout=0.5) + + +@pytest.mark.asyncio +async def test_async_exec_code_with_executor(): + """async_exec_code works with a shared ThreadPoolExecutor.""" + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + result = await async_exec_code("print('pooled')", executor=pool) + assert result.strip() == "pooled" + + +@pytest.mark.asyncio +async def test_async_exec_code_no_timeout(): + """async_exec_code works with timeout=None.""" + result = await async_exec_code("print('no limit')", timeout=None) + assert result.strip() == "no limit"