From 80d3cc63967f192cd200603069b187ca851390ab Mon Sep 17 00:00:00 2001 From: u7-u7 <116776533+u7-u7@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:58:03 +0800 Subject: [PATCH] fix: preserve complete messages in sub-agent stream traces --- ms_agent/tools/agent_tool.py | 14 ++++- ms_agent/utils/stream_writer.py | 37 ++++++++---- tests/agent/test_llm_io.py | 31 ++++++++++ tests/utils/test_stream_writer.py | 94 +++++++++++++++++++++++++++++++ 4 files changed, 163 insertions(+), 13 deletions(-) create mode 100644 tests/utils/test_stream_writer.py diff --git a/ms_agent/tools/agent_tool.py b/ms_agent/tools/agent_tool.py index 84b97e270..174e40497 100644 --- a/ms_agent/tools/agent_tool.py +++ b/ms_agent/tools/agent_tool.py @@ -1226,8 +1226,18 @@ def _emit_stream_event(event: Dict[str, Any]) -> None: else: runner = _run_and_collect + async def run_with_stream_cleanup(): + try: + return await runner() + except BaseException as exc: + if _writer is not None: + error = ('cancelled' if isinstance( + exc, asyncio.CancelledError) else str(exc)) + _writer.on_error(error) + raise + if not self._enable_stats: - result = await runner() + result = await run_with_stream_cleanup() if not spec.run_in_process: self._save_transcript(result, runtime_agent_tag) if _writer is not None: @@ -1241,7 +1251,7 @@ def _emit_stream_event(event: Dict[str, Any]) -> None: status = 'completed' result = None try: - result = await runner() + result = await run_with_stream_cleanup() if not spec.run_in_process: self._save_transcript(result, runtime_agent_tag) if _writer is not None: diff --git a/ms_agent/utils/stream_writer.py b/ms_agent/utils/stream_writer.py index 18f74ef8f..39431a67b 100644 --- a/ms_agent/utils/stream_writer.py +++ b/ms_agent/utils/stream_writer.py @@ -3,7 +3,7 @@ SubAgentStreamWriter — incremental JSONL writer for sub-agent execution progress. When a parent agent calls a sub-agent via AgentTool, the sub-agent produces messages -incrementally. This writer appends new messages to a JSONL file on every chunk event, +incrementally. This writer appends completed messages to a JSONL file, allowing external tools (e.g. ``tail -f``) or the parent agent itself to watch the sub-agent's progress in real time. @@ -38,7 +38,8 @@ class SubAgentStreamWriter: Each instance owns exactly one file. It deduplicates history by tracking how many messages have already been written (``_last_written_count``), so calling ``on_chunk(full_history)`` multiple times is safe — only the newly - appended messages are written. + appended messages are written. The last message can still be growing across + chunks; it is held until a following message arrives or the run finishes. All public methods are safe to call from multiple threads. """ @@ -48,6 +49,7 @@ def __init__(self, output_dir: str, call_id: str, tool_name: str) -> None: self._tool_name: str = tool_name self._lock = threading.Lock() self._last_written_count: int = 0 + self._pending_message: Any = None self._closed: bool = False self._agent_tag: Optional[str] = None self._file = None # opened lazily in on_start @@ -90,7 +92,7 @@ def on_start(self, agent_tag: Optional[str]) -> None: self._file = None def on_chunk(self, history: Any) -> None: - """Append only new messages from *history* since the last call. + """Append completed messages and retain the latest streaming tail. Args: history: The full accumulated message list returned by a streaming @@ -103,14 +105,10 @@ def on_chunk(self, history: Any) -> None: with self._lock: if self._closed or self._file is None: return - for msg in messages[self._last_written_count:]: - self._write_line({ - 'type': 'message', - 'index': self._last_written_count, - 'message': _msg_to_dict(msg), - 'ts': _now_iso(), - }) - self._last_written_count += 1 + for msg in messages[self._last_written_count:-1]: + self._write_message(msg) + if len(messages) > self._last_written_count: + self._pending_message = messages[-1] def on_end(self, history: Any) -> None: """Flush any remaining messages, write footer record, then close. @@ -123,6 +121,7 @@ def on_end(self, history: Any) -> None: with self._lock: if self._closed: return + self._flush_pending() self._closed = True if self._file is not None: try: @@ -152,6 +151,7 @@ def on_error(self, error: str) -> None: with self._lock: if self._closed: return + self._flush_pending() self._closed = True if self._file is not None: try: @@ -173,6 +173,21 @@ def on_error(self, error: str) -> None: # ── private helpers ──────────────────────────────────────────────────── + def _write_message(self, msg: Any) -> None: + self._write_line({ + 'type': 'message', + 'index': self._last_written_count, + 'message': _msg_to_dict(msg), + 'ts': _now_iso(), + }) + self._last_written_count += 1 + + def _flush_pending(self) -> None: + # Caller holds the lock, including during cancellation/error shutdown. + if self._pending_message is not None and self._file is not None: + self._write_message(self._pending_message) + self._pending_message = None + def _write_line(self, record: Dict[str, Any]) -> None: """Serialize *record* as JSON and append a newline. diff --git a/tests/agent/test_llm_io.py b/tests/agent/test_llm_io.py index 08b2c25d1..faa5468fe 100644 --- a/tests/agent/test_llm_io.py +++ b/tests/agent/test_llm_io.py @@ -382,3 +382,34 @@ def test_interrupt_preserves_shared_or_released_connection(http_version, closed) io.interrupt_stream(stream) network.get_extra_info.assert_not_called() stream.close.assert_called_once() + + +@run_async +async def test_stream_trace_preserves_multi_chunk_model_response(tmp_path): + from ms_agent.utils.stream_writer import SubAgentStreamWriter + + events = [] + for delta in [{'content': 'H'}, {'content': 'ello'}, {}]: + events.append('data: ' + json.dumps({ + 'id': 'trace', 'object': 'chat.completion.chunk', 'created': 0, + 'model': 'probe', 'choices': [{'index': 0, 'delta': delta, + 'finish_reason': None if delta else 'stop'}], + }) + '\n\n') + events.append('data: [DONE]\n\n') + agent = make_agent(lambda request: httpx.Response( + 200, content=''.join(events), + headers={'Content-Type': 'text/event-stream'})) + writer = SubAgentStreamWriter(str(tmp_path), 'trace', 'worker') + writer.on_start('worker') + try: + async for history in LLMAgent.step.__wrapped__( + agent, [Message(role='user', content='Say hello')]): + writer.on_chunk(history) + writer.on_end(history) + finally: + agent.llm.client.close() + with open(writer.stream_path, encoding='utf-8') as stream: + rows = [json.loads(line) for line in stream] + messages = [row['message'] for row in rows if row['type'] == 'message'] + assert messages[-1]['content'] == 'Hello' + assert rows[-1]['total_messages'] == len(history) diff --git a/tests/utils/test_stream_writer.py b/tests/utils/test_stream_writer.py new file mode 100644 index 000000000..f194df90e --- /dev/null +++ b/tests/utils/test_stream_writer.py @@ -0,0 +1,94 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Sub-agent traces must retain completed messages, not first-token snapshots.""" +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +from omegaconf import OmegaConf + +from ms_agent.llm.utils import Message +from ms_agent.tools.agent_tool import AgentTool +from ms_agent.utils.stream_writer import SubAgentStreamWriter + + +def records(writer): + return [json.loads(line) for line in + Path(writer.stream_path).read_text(encoding='utf-8').splitlines()] + + +@pytest.mark.parametrize('mutate', [False, True]) +def test_cumulative_tail_is_written_once_when_complete(tmp_path, mutate): + writer = SubAgentStreamWriter(str(tmp_path), 'call', 'worker') + writer.on_start('worker') + history = [Message(role='user', content='question'), + Message(role='assistant', content='H')] + writer.on_chunk(history) + if mutate: + history[-1].content = 'Hello' + else: + history[-1] = Message(role='assistant', content='Hello') + writer.on_chunk(history) + writer.on_end(history) + writer.on_end(history) + rows = records(writer) + messages = [row for row in rows if row['type'] == 'message'] + assert [row['index'] for row in messages] == [0, 1] + assert [row['message']['content'] for row in messages] == ['question', 'Hello'] + assert rows[-1]['total_messages'] == 2 + assert len([row for row in rows if row['type'] == 'footer']) == 1 + + +def test_tool_arguments_finalized_before_next_message(tmp_path): + writer = SubAgentStreamWriter(str(tmp_path), 'call', 'worker') + writer.on_start('worker') + history = [{'role': 'assistant', 'content': '', 'tool_calls': [ + {'id': 'call1', 'arguments': '{"path":'}]}] + writer.on_chunk(history) + history[0]['tool_calls'][0]['arguments'] = '{"path":"report.md"}' + history.append({'role': 'tool', 'content': 'written'}) + writer.on_chunk(history) + # The completed tool call is available before the agent finishes. + assert records(writer)[1]['message']['tool_calls'][0]['arguments'] == ( + '{"path":"report.md"}') + writer.on_end(history) + assert records(writer)[-1]['total_messages'] == 2 + + +@pytest.mark.parametrize('enable_stats', [False, True]) +@pytest.mark.parametrize('failure', [None, RuntimeError('model disconnected'), + asyncio.CancelledError()]) +def test_agent_tool_stream_file_tracks_final_or_interrupted_text( + tmp_path, failure, enable_stats): + class StreamingAgent: + tag = 'worker' + + async def run(self, payload, stream): + assert stream + + async def chunks(): + history = [Message(role='assistant', content='H')] + yield history + history[-1] = Message(role='assistant', content='Hello') + yield history + if failure is not None: + raise failure + return chunks() + + tool = AgentTool(OmegaConf.create({ + 'output_dir': str(tmp_path), 'agent_stream_file': True, 'tools': {}})) + tool._enable_stats = enable_stats + spec = SimpleNamespace(tool_name='worker', server_name='agents', run_in_thread=False, + run_in_process=False) + call = tool._run_agent(StreamingAgent(), 'question', spec, call_id='call') + if failure is None: + asyncio.run(call) + else: + with pytest.raises(type(failure)): + asyncio.run(call) + path = tmp_path / '.ms_agent/subagents/call.stream.jsonl' + rows = [json.loads(line) for line in path.read_text().splitlines()] + assert rows[1]['message']['content'] == 'Hello' + assert rows[-1]['type'] == 'footer' + assert rows[-1]['status'] == ('complete' if failure is None else 'error')