fix: preserve complete messages in sub-agent stream traces

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:

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.
 

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)

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')