237 lines
9.6 KiB
Python
237 lines
9.6 KiB
Python
from __future__ import annotations
|
||
|
||
import sys
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from core.llm_transport import ( # noqa: E402
|
||
PreferNonstreamToolCall,
|
||
)
|
||
from core.llm_transport import (
|
||
malformed_tool_calls as _malformed_tool_calls,
|
||
)
|
||
from core.loop import AgentLoop # noqa: E402
|
||
|
||
|
||
def _resp(arguments: str):
|
||
"""最小 response 桩:单条 tool_call,arguments 由测试指定。"""
|
||
tc = SimpleNamespace(function=SimpleNamespace(name="write", arguments=arguments))
|
||
msg = SimpleNamespace(tool_calls=[tc], content=None)
|
||
return SimpleNamespace(choices=[SimpleNamespace(message=msg)])
|
||
|
||
|
||
GOOD = '{"path": "a.md", "content": "ok"}'
|
||
BAD = '].cells[1].merge({"path": "a.md"' # 流式 delta 错位的典型形态
|
||
|
||
|
||
def _make_loop(stream_results, nonstream_results):
|
||
"""绕开构造器搭 AgentLoop:只装 _stream_llm 用到的属性,打桩两条取流路径。"""
|
||
loop = object.__new__(AgentLoop)
|
||
loop.caps = SimpleNamespace(reliable_context=64_000, family="test", variant="t")
|
||
loop.session = SimpleNamespace(messages=[], task_id="test-task")
|
||
loop.user_id = "test-user" # 无 DB 环境:_log_malformed_args 的落库路径应静默跳过
|
||
loop.user_root = None
|
||
# salvage(0.58.24)接在畸形重试前:空 schemas → 未知工具无白名单 → _try_salvage 返 False,
|
||
# 这些用例本就测「salvage 救不了 → 落非流式重试」,给个空 executor 即可走到该分支。
|
||
loop.executor = SimpleNamespace(schemas=lambda: [])
|
||
loop.events = []
|
||
loop._emit = loop.events.append
|
||
calls = {"stream": 0, "nonstream": 0}
|
||
|
||
def collect_stream_once(llm_messages):
|
||
calls["stream"] += 1
|
||
return stream_results.pop(0), False
|
||
|
||
def nonstream_once(llm_messages):
|
||
calls["nonstream"] += 1
|
||
return nonstream_results.pop(0)
|
||
|
||
loop._collect_stream_once = collect_stream_once
|
||
loop._nonstream_once = nonstream_once
|
||
return loop, calls
|
||
|
||
|
||
class TestMalformedRetry(unittest.TestCase):
|
||
def test_deepseek_write_name_resets_streamed_reasoning_before_reroute(self):
|
||
"""真实 collect 路径实时发 reasoning；见到 write 后重置再改走非流式。"""
|
||
reasoning = SimpleNamespace(
|
||
choices=[SimpleNamespace(delta=SimpleNamespace(
|
||
reasoning_content="thinking", content=None, tool_calls=None,
|
||
))]
|
||
)
|
||
write = SimpleNamespace(
|
||
choices=[SimpleNamespace(delta=SimpleNamespace(
|
||
reasoning_content=None,
|
||
content=None,
|
||
tool_calls=[SimpleNamespace(
|
||
function=SimpleNamespace(name="write", arguments=""),
|
||
)],
|
||
))]
|
||
)
|
||
loop = object.__new__(AgentLoop)
|
||
loop.caps = SimpleNamespace(
|
||
family="deepseek_v4", default_reasoning_effort=None,
|
||
)
|
||
loop.executor = SimpleNamespace(schemas=lambda: [])
|
||
loop.cancel_check = None
|
||
loop.events = []
|
||
loop._emit = loop.events.append
|
||
loop.llm = SimpleNamespace(
|
||
chat_stream=lambda **_kwargs: iter([reasoning, write]),
|
||
)
|
||
|
||
with self.assertRaises(PreferNonstreamToolCall):
|
||
loop._collect_stream_once([])
|
||
self.assertEqual(loop.events, [
|
||
{"type": "reasoning", "delta": "thinking"},
|
||
{"type": "reasoning_reset"},
|
||
])
|
||
|
||
def test_deepseek_reasoning_is_emitted_before_output_route_is_known(self):
|
||
"""普通 DeepSeek 回答不能等首个正文 token 才释放已经到达的推理。"""
|
||
reasoning = SimpleNamespace(
|
||
choices=[SimpleNamespace(delta=SimpleNamespace(
|
||
reasoning_content="thinking", content=None, tool_calls=None,
|
||
))]
|
||
)
|
||
text = SimpleNamespace(
|
||
choices=[SimpleNamespace(delta=SimpleNamespace(
|
||
reasoning_content=None, content="answer", tool_calls=None,
|
||
))]
|
||
)
|
||
loop = object.__new__(AgentLoop)
|
||
loop.caps = SimpleNamespace(
|
||
family="deepseek_v4", default_reasoning_effort=None,
|
||
)
|
||
loop.executor = SimpleNamespace(schemas=lambda: [])
|
||
loop.cancel_check = None
|
||
loop.events = []
|
||
loop._emit = loop.events.append
|
||
loop.llm = SimpleNamespace(
|
||
chat_stream=lambda **_kwargs: iter([reasoning, text]),
|
||
)
|
||
|
||
with unittest.mock.patch(
|
||
"core.loop.litellm.stream_chunk_builder", return_value=object()
|
||
):
|
||
loop._collect_stream_once([])
|
||
|
||
self.assertEqual(loop.events, [
|
||
{"type": "reasoning", "delta": "thinking"},
|
||
{"type": "text", "delta": "answer"},
|
||
])
|
||
|
||
def test_preferred_tool_reroutes_before_arguments_stream(self):
|
||
"""首包识别 write/edit 后直接非流式，不把它算作一次畸形失败。"""
|
||
loop, calls = _make_loop([], [_resp(GOOD)])
|
||
|
||
def prefer_nonstream(_messages):
|
||
calls["stream"] += 1
|
||
raise PreferNonstreamToolCall("write")
|
||
|
||
loop._collect_stream_once = prefer_nonstream
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertFalse(cancelled)
|
||
self.assertEqual(_malformed_tool_calls(resp), [])
|
||
self.assertEqual(calls, {"stream": 1, "nonstream": 1})
|
||
infos = [e for e in loop.events if e.get("level") == "info"]
|
||
self.assertEqual(len(infos), 1)
|
||
self.assertIn("write", infos[0]["msg"])
|
||
starts = [e for e in loop.events if e.get("type") == "llm_start"]
|
||
self.assertEqual(len(starts), 2)
|
||
|
||
def test_clean_stream_no_retry(self):
|
||
loop, calls = _make_loop([_resp(GOOD)], [])
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertFalse(cancelled)
|
||
self.assertEqual(_malformed_tool_calls(resp), [])
|
||
self.assertEqual(calls, {"stream": 1, "nonstream": 0})
|
||
|
||
def test_first_failure_falls_back_to_nonstream(self):
|
||
"""核心断言:流式只试 1 次,首败即降级非流式(不再流式重 roll)。"""
|
||
loop, calls = _make_loop([_resp(BAD)], [_resp(GOOD)])
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertFalse(cancelled)
|
||
self.assertEqual(_malformed_tool_calls(resp), [])
|
||
self.assertEqual(calls, {"stream": 1, "nonstream": 1})
|
||
warns = [e for e in loop.events if e.get("type") == "warn"]
|
||
self.assertEqual(len(warns), 1)
|
||
self.assertIn("非流式", warns[0]["msg"])
|
||
# 每个 attempt 重发 llm_start:warn 关掉前端占位段后,靠它重建「思考中」跳秒
|
||
starts = [e for e in loop.events if e.get("type") == "llm_start"]
|
||
self.assertEqual(len(starts), 2)
|
||
|
||
def test_all_attempts_exhausted_returns_last(self):
|
||
"""全部尝试耗尽仍畸形 → 返回最后一次 response(交 executor 返错),不死循环。"""
|
||
n_nonstream = AgentLoop._MAX_MALFORMED_ATTEMPTS - 1
|
||
loop, calls = _make_loop([_resp(BAD)], [_resp(BAD)] * n_nonstream)
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertFalse(cancelled)
|
||
self.assertTrue(_malformed_tool_calls(resp))
|
||
self.assertEqual(calls, {"stream": 1, "nonstream": n_nonstream})
|
||
warns = [e for e in loop.events if e.get("type") == "warn"]
|
||
self.assertEqual(len(warns), AgentLoop._MAX_MALFORMED_ATTEMPTS)
|
||
|
||
def test_cancel_mid_stream_short_circuits(self):
|
||
loop, calls = _make_loop([], [])
|
||
loop._collect_stream_once = lambda m: (None, True)
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertTrue(cancelled)
|
||
self.assertIsNone(resp)
|
||
self.assertEqual(calls["nonstream"], 0)
|
||
|
||
def test_cancel_during_nonstream_retry(self):
|
||
"""非流式重试期间 cancel(_nonstream_once 返回 None)→ (None, True),不再续试。"""
|
||
loop, calls = _make_loop([_resp(BAD)], [None, _resp(GOOD)])
|
||
resp, cancelled = loop._stream_llm()
|
||
self.assertTrue(cancelled)
|
||
self.assertIsNone(resp)
|
||
self.assertEqual(calls, {"stream": 1, "nonstream": 1})
|
||
|
||
def test_nonstream_once_polls_cancel(self):
|
||
"""真 _nonstream_once:llm.chat 阻塞时点停止,应在 poll 拍内返回 None,
|
||
不等调用整个返回(此前同步阻塞,分钟级生成期间停止按钮无效)。"""
|
||
loop = object.__new__(AgentLoop)
|
||
loop.caps = SimpleNamespace(default_reasoning_effort=None)
|
||
loop.executor = SimpleNamespace(schemas=lambda: [])
|
||
loop.cancel_check = lambda: True
|
||
block = threading.Event()
|
||
|
||
def slow_chat(**kwargs):
|
||
block.wait(30)
|
||
return _resp(GOOD)
|
||
|
||
loop.llm = SimpleNamespace(chat=slow_chat)
|
||
t0 = time.monotonic()
|
||
resp = loop._nonstream_once([])
|
||
elapsed = time.monotonic() - t0
|
||
block.set() # 释放弃养线程,别让它拖住测试进程退出
|
||
self.assertIsNone(resp)
|
||
self.assertLess(elapsed, 5)
|
||
|
||
def test_nonstream_once_returns_response_when_not_cancelled(self):
|
||
"""不取消时行为与原同步版等价:拿到 response 并补 emit 整段 text。"""
|
||
loop = object.__new__(AgentLoop)
|
||
loop.caps = SimpleNamespace(default_reasoning_effort=None)
|
||
loop.executor = SimpleNamespace(schemas=lambda: [])
|
||
loop.cancel_check = None
|
||
loop.events = []
|
||
loop._emit = loop.events.append
|
||
good = _resp(GOOD)
|
||
good.choices[0].message.content = "hello"
|
||
loop.llm = SimpleNamespace(chat=lambda **kwargs: good)
|
||
resp = loop._nonstream_once([])
|
||
self.assertIs(resp, good)
|
||
texts = [e for e in loop.events if e.get("type") == "text"]
|
||
self.assertEqual(len(texts), 1)
|
||
self.assertEqual(texts[0]["delta"], "hello")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|