From d3da0dce6802cd125c1e915258ea9bc96e8dd062 Mon Sep 17 00:00:00 2001 From: christop <825583681@qq.com> Date: Wed, 17 Jun 2026 12:28:20 +0800 Subject: [PATCH] fix(strategy/tot): use ast.literal_eval for model-generated thoughts ThoughtSolverBase.generate_thoughts parsed the model's generated thought block with eval() after only a cosmetic markdown-fence strip. Because the Tree-of-Thoughts solver feeds raw model output (rsp -> CodeParser.parse_code) into eval(), a model that emits non-literal Python (e.g. a list whose element calls __import__('os').system(...)) executes arbitrary code on the host. No sandbox, try/except, or literal check guarded the call. The model is instructed (OUTPUT_FORMAT) to return a plain JSON/Python list of thought nodes, so the value is always a literal data structure. Replace eval() with ast.literal_eval(), which only evaluates literals and raises on code. Parse failures are logged and degrade to an empty thought list (matching update_node's default), preserving behaviour for well-formed output. Adds a regression test that fails on the previous eval() path (model-supplied code executes) and passes with literal-only parsing, plus a benign-output test confirming valid thought lists still parse. The regression test injects a mocked LLM, but the session-wide autouse llm_mock fixture (tests/conftest.py) still builds a real OpenAILLM for every test, which validates config.llm.proxy via httpx. An earlier-collected module assigns a non-URL object onto the shared config.llm instance and never restores it, so that leaked proxy made llm_mock raise "Proxy protocol must be ..." at setup of these tests. A module-scoped autouse fixture restores config.llm.proxy to a sane value before llm_mock runs and puts it back afterwards, isolating these tests without affecting any other module. Signed-off-by: christop <825583681@qq.com> --- metagpt/strategy/tot.py | 10 +- .../test_tot_generate_thoughts_eval.py | 134 ++++++++++++++++++ 2 files changed, 143 insertions(+), 1 deletion(-) create mode 100644 tests/metagpt/strategy/test_tot_generate_thoughts_eval.py diff --git a/metagpt/strategy/tot.py b/metagpt/strategy/tot.py index 17ce632118..ac1c42e7e1 100644 --- a/metagpt/strategy/tot.py +++ b/metagpt/strategy/tot.py @@ -4,6 +4,7 @@ # @Desc : from __future__ import annotations +import ast import asyncio from typing import Any, List, Optional @@ -63,7 +64,14 @@ async def generate_thoughts(self, current_state="", current_node=None) -> List[T ) rsp = await self.llm.aask(msg=state_prompt + "\n" + OUTPUT_FORMAT) thoughts = CodeParser.parse_code(text=rsp) - thoughts = eval(thoughts) + # The model is asked (see OUTPUT_FORMAT) to return a plain list literal of thought nodes. + # Parse it with ast.literal_eval so that only literal data structures are accepted; this + # prevents arbitrary code execution if the model output contains non-literal Python. + try: + thoughts = ast.literal_eval(thoughts) + except (ValueError, SyntaxError) as e: + logger.warning(f"Failed to parse thoughts as a literal structure, ignoring generated thoughts: {e}") + thoughts = [] # fixme 避免不跟随，生成过多nodes # valid_thoughts = [_node for idx, _node in enumerate(thoughts) if idx < self.n_generate_sample] return self.thought_tree.update_node(thoughts, current_node=current_node) diff --git a/tests/metagpt/strategy/test_tot_generate_thoughts_eval.py b/tests/metagpt/strategy/test_tot_generate_thoughts_eval.py new file mode 100644 index 0000000000..70c90ca6f7 --- /dev/null +++ b/tests/metagpt/strategy/test_tot_generate_thoughts_eval.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- +""" +Regression test for metagpt-eval-exec-tot-generate-thoughts-eval. + +ThoughtSolverBase.generate_thoughts used to `eval()` the model's generated +thought block, which allowed arbitrary code execution from model output. +It now uses `ast.literal_eval`, which only accepts literal data structures. + +These tests verify: + * a benign list-of-dicts thought block is still parsed correctly, and + * a malicious (non-literal) payload does NOT execute and is ignored. + +The Tree-of-Thoughts solver is exercised with a fully mocked LLM. The real +``metagpt.llm.LLM`` factory builds a provider from ``Config`` (api key, proxy, +base url, ...), which touches a real configuration and the network. To stay +config-independent and offline, the mock LLM is passed straight into the +constructor: ``ThoughtSolverBase.llm`` declares ``default_factory=LLM``, and +pydantic only calls that factory when no value is supplied, so providing +``llm=`` means the real factory is never invoked (assigning ``solver.llm`` +afterwards would be too late -- the factory runs during ``__init__``). +Following the existing suite (e.g. ``tests/metagpt/test_role.py``) the mock is a +``MagicMock(spec=BaseLLM)`` with an ``AsyncMock`` ``aask`` returning canned text. +As a safety net ``create_llm_instance`` is patched to fail loudly if anything +still tries to build a real provider. +""" + +import ast +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from metagpt.config2 import config +from metagpt.provider.base_llm import BaseLLM +from metagpt.strategy.base import BaseParser, ThoughtNode, ThoughtTree +from metagpt.strategy.tot import ThoughtSolverBase +from metagpt.strategy.tot_schema import ThoughtSolverConfig + + +@pytest.fixture(autouse=True, scope="module") +def _restore_global_llm_proxy(): + """Keep the global ``config.llm.proxy`` sane while this module runs. + + The session-wide autouse ``llm_mock`` fixture (tests/conftest.py) builds a + real ``MockLLM`` -> ``OpenAILLM`` for *every* test, which constructs an httpx + client and validates ``config.llm.proxy``. Some earlier-collected test + modules assign a non-URL object straight onto the shared ``config.llm`` + instance (e.g. ``config.get_openai_llm().proxy = mocker.PropertyMock(...)``); + that assignment is a plain attribute write, so it is never rolled back and + leaks into later modules. When ``llm_mock`` then runs for our tests the bad + proxy makes httpx raise ``ValueError: Proxy protocol must be ...``. + + Being module-scoped, this autouse fixture is set up before the function-scoped + ``llm_mock``, so it neutralises any such leaked value before ``MockLLM`` is + built, and restores the original afterwards. It does not change behaviour for + any other test module. + """ + saved_proxy = config.llm.proxy + if not isinstance(config.llm.proxy, str): + config.llm.proxy = None + yield + config.llm.proxy = saved_proxy + + +class _StubParser(BaseParser): + def propose(self, current_state: str, **kwargs) -> str: + return "propose" + + +@pytest.fixture +def no_real_llm(mocker): + """Guarantee no real LLM provider is constructed (no Config/proxy/network).""" + + def _boom(*args, **kwargs): + raise AssertionError("create_llm_instance must not be called: real LLM was built") + + mocker.patch("metagpt.provider.llm_provider_registry.create_llm_instance", side_effect=_boom) + mocker.patch("metagpt.context.create_llm_instance", side_effect=_boom) + + +def _make_solver(response: str) -> ThoughtSolverBase: + """Build a solver whose LLM is fully mocked. + + The mock is passed via the ``llm=`` constructor argument so pydantic never + falls back to ``default_factory=LLM`` -- i.e. no real provider, Config, + proxy resolution, or network access happens at construction time. + """ + mock_llm = MagicMock(spec=BaseLLM) + mock_llm.aask = AsyncMock(return_value=response) + mock_llm.use_system_prompt = True + + solver = ThoughtSolverBase(llm=mock_llm, config=ThoughtSolverConfig(parser=_StubParser())) + assert solver.llm is mock_llm # guard: the real LLM factory was never used + solver.thought_tree = ThoughtTree(ThoughtNode(name="root")) + return solver + + +@pytest.mark.asyncio +async def test_generate_thoughts_parses_benign_list(no_real_llm): + response = ( + "```json\n" + '[{"node_id": "1", "node_state_instruction": "first thought"},\n' + ' {"node_id": "2", "node_state_instruction": "second thought"}]\n' + "```" + ) + solver = _make_solver(response) + nodes = await solver.generate_thoughts(current_state="", current_node=None) + assert [n.name for n in nodes] == ["first thought", "second thought"] + assert [n.id for n in nodes] == [1, 2] + + +@pytest.mark.asyncio +async def test_generate_thoughts_does_not_execute_model_code(no_real_llm, tmp_path: Path): + sentinel = tmp_path / "tot_pwned" + # A literal-only parser must refuse to run this; with eval() it would execute. + payload = f"```python\n[__import__('os').system('touch {sentinel}')]\n```" + solver = _make_solver(payload) + + # Swallow any downstream error so the discriminating signal is purely the + # side effect: on the vulnerable (eval) code path the sentinel file is + # created before any exception; on the fixed (ast.literal_eval) path it is not. + try: + nodes = await solver.generate_thoughts(current_state="", current_node=None) + except Exception: + nodes = None + + assert not sentinel.exists(), "model-supplied code was executed (eval sink still live)" + assert nodes == [] + + +def test_literal_eval_rejects_code(): + # Guard the underlying primitive directly. + with pytest.raises((ValueError, SyntaxError)): + ast.literal_eval("__import__('os').system('id')")