dreadnode 2.0.8__tar.gz → 2.0.10__tar.gz

{dreadnode-2.0.8 → dreadnode-2.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 2.0.8
+Version: 2.0.10
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://dreadnode.io
 Project-URL: Documentation, https://docs.dreadnode.io
@@ -19,7 +19,7 @@ Requires-Dist: gepa>=0.1.1
 Requires-Dist: httpx<1.0.0,>=0.28.0
 Requires-Dist: jsonpath-ng>=1.7.0
 Requires-Dist: jsonref>=1.1.0
-Requires-Dist: litellm<=1.82.2,>=1.80.11
+Requires-Dist: litellm<=1.83.0,>=1.83.0
 Requires-Dist: logfire<=3.20.0,>=3.5.3
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: mcp<2.0.0,>=1.25.0

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/agent.py

@@ -1,7 +1,7 @@
 import asyncio
 import inspect
 import typing as t
-from contextlib import AsyncExitStack, asynccontextmanager, nullcontext, suppress
+from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from copy import deepcopy
 from pathlib import Path
 from textwrap import dedent
@@ -120,7 +120,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
     )
     judge: Judge[Rubric] | None = Config(default=None, repr=False)
     trajectory: Trajectory = Field(default_factory=Trajectory, exclude=True, repr=False)
-    max_steps: int = Config(default=10, ge=1)
+    max_steps: int = Config(default=1000, ge=1)
     """Maximum number of generation/tool steps before the agent stops."""
     generation_timeout: int | None = Config(default=None)
     """Timeout in seconds for each LLM generation call. None = no timeout."""
@@ -423,7 +423,8 @@ class Agent(Executor[AgentEvent, Trajectory]):
         )
 
         try:
-            summary = await summarize_conversation.bind(summarizer)(  # ty: ignore[unresolved-attribute]  # .bind() added by generator framework at runtime
+            summary = await summarize_conversation(
+                summarizer,
                 "\n".join(str(msg) for msg in to_summarize),
             )
         except Exception:
@@ -585,7 +586,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
                 agent_name=self.name,
                 inputs={"goal": self._current_input},
                 params={
-                    "session_id": str(self.trajectory.session_id),
+                    "session_id": str(traj.session_id),
                     "max_steps": self.max_steps,
                     **({"model": self.model_name} if self.model_name else {}),
                     **({"tools": [tool.name for tool in self.all_tools]} if self.all_tools else {}),
@@ -682,7 +683,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
                         )
 
                     # Check stop conditions INSIDE span
-                    if any(cond(self.trajectory.steps) for cond in self.stop_conditions):
+                    if any(cond(traj.steps) for cond in self.stop_conditions):
                         logger.info("A stop condition was met. Ending run.")
                         gen_end = GenerationEnd(
                             agent_id=self.agent_id,
@@ -871,7 +872,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
                         )
 
                     # Check stop conditions again INSIDE span
-                    if any(cond(self.trajectory.steps) for cond in self.stop_conditions):
+                    if any(cond(traj.steps) for cond in self.stop_conditions):
                         break
 
             except Continue as e:
@@ -898,7 +899,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
             stop_reason = "error"
         elif max_steps_reached:
             stop_reason = "max_steps_reached"
-        elif self.trajectory.steps and isinstance(self.trajectory.steps[-1], AgentStalled):
+        elif traj.steps and isinstance(traj.steps[-1], AgentStalled):
             stop_reason = "stalled"
 
         async for event in self._dispatch(
@@ -960,17 +961,12 @@ class Agent(Executor[AgentEvent, Trajectory]):
                 if hasattr(tool_container, "__aenter__") and hasattr(tool_container, "__aexit__"):
                     await stack.enter_async_context(tool_container)
 
-            # Parent span only on new conversation (not when borrowing external trajectory)
             display_name = self.name or str(self.agent_id)[:8]
-            ctx = (
-                dreadnode.task_span(
-                    f"agent:{display_name}",
-                    type="agent",
-                    label=display_name,
-                    tags=["agent"],
-                )
-                if reset and trajectory is None
-                else nullcontext()
+            ctx = dreadnode.task_span(
+                f"agent:{display_name}",
+                type="agent",
+                label=display_name,
+                tags=["agent"],
             )
 
             with ctx as parent_span:

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/hooks.py

@@ -21,7 +21,6 @@ from dreadnode.agents.events import (
 )
 from dreadnode.agents.reactions import Reaction, Retry
 from dreadnode.core.hook import Hook, hook
-from dreadnode.generators.chat import Chat
 from dreadnode.generators.generator import Generator
 from dreadnode.generators.message import Message, make_compaction_message
 
@@ -159,98 +158,98 @@ CONTEXT_LENGTH_ERROR_PATTERNS = [
 class Summary:
     analysis: str
     summary: str
-    chat: Chat
 
 
-async def summarize_conversation(conversation: str, *, guidance: str = "") -> Summary:  # type: ignore[empty-body]
-    """
-    Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
-    This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
-
-    Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:
-
-    1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
-    - The user's explicit requests and intents
-    - Your approach to addressing the user's requests
-    - Key decisions, technical concepts and code patterns
-    - Specific technical details like paths, usernames, structured objects, and code
-    - Tool interactions performed with a specific focus on intent and outcome
-    - Errors that you ran into and how you fixed them
-    - Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
-
-    2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
-
-    Your summary should include the following sections:
-
-    1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
-    2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
-    3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
-    4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
-    5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
-    6. All user messages: List ALL user messages that are not tool results. These are critical for understanding the users' feedback and changing intent.
-    7. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
-    8. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
-    9. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
-    If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
-
-    Here's an example of how your output should be structured:
+_SUMMARIZATION_SYSTEM_PROMPT = """\
+Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
+This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
 
-    <example>
-    <analysis>
-    [Your thought process, ensuring all points are covered thoroughly and accurately]
-    </analysis>
+Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:
 
-    <summary>
-    1. Primary Request and Intent:
-    [Detailed description]
+1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
+- The user's explicit requests and intents
+- Your approach to addressing the user's requests
+- Key decisions, technical concepts and code patterns
+- Specific technical details like paths, usernames, structured objects, and code
+- Tool interactions performed with a specific focus on intent and outcome
+- Errors that you ran into and how you fixed them
+- Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
 
-    2. Key Technical Details:
+2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
 
-    - [Concept 1]
-    - [Concept 2]
-    - [Object 1]
-    - [...]
+Your summary should include the following sections:
 
-    3. Tool Interactions:
-
-    - [Tool Call 1]: [Description of what the tool did, inputs, outputs]
-    - [Tool Call 2]: ...
-    - [...]
-
-    4. Errors and fixes:
-
-    - [Detailed description of error 1]:
-        - [How you fixed the error]
-        - [User feedback on the error if any]
-    - [...]
-
-    5. Problem Solving:
-    [Description of solved problems and ongoing troubleshooting]
+1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
+2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
+3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
+4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
+5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
+6. All user messages: List ALL user messages that are not tool results. These are critical for understanding the users' feedback and changing intent.
+7. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
+8. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
+9. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
+If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
+"""
 
-    6. All user messages:
+_SUMMARIZATION_USER_TEMPLATE = """\
+Please provide your summary based on the conversation below, following the system instructions and ensuring precision and thoroughness in your response.
 
-    - [Detailed non tool use user message]
-    - [...]
+{guidance_section}<conversation>
+{conversation}
+</conversation>"""
 
-    7. Pending Tasks:
 
-    - [Task 1]
-    - [Task 2]
-    - [...]
+async def summarize_conversation(
+    generator: "str | Generator",
+    conversation: str,
+    *,
+    guidance: str = "",
+) -> Summary:
+    """Run the summarization prompt against the given generator and return a Summary."""
+    import re
+
+    from dreadnode.generators.generator import GenerateParams
+    from dreadnode.generators.generator import get_generator as _get_generator
+
+    if isinstance(generator, str):
+        generator = _get_generator(generator)
+
+    guidance_section = ""
+    if guidance:
+        guidance_section = f"Additional summarization guidance:\n{guidance}\n\n"
+
+    messages = [
+        Message(role="system", content=_SUMMARIZATION_SYSTEM_PROMPT),
+        Message(
+            role="user",
+            content=_SUMMARIZATION_USER_TEMPLATE.format(
+                guidance_section=guidance_section,
+                conversation=conversation,
+            ),
+        ),
+    ]
 
-    8. Current Work:
-    [Precise description of current work]
+    results = await generator.generate_messages(
+        [messages],
+        [generator.params or GenerateParams()],
+    )
 
-    9. Optional Next Step:
-    [Optional Next step to take]
+    result = results[0]
+    if isinstance(result, BaseException):
+        raise result
 
-    </summary>
-    </example>
+    response_text = result.message.content
 
-    Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response.
+    analysis = ""
+    summary_text = response_text
+    analysis_match = re.search(r"<analysis>(.*?)</analysis>", response_text, re.DOTALL)
+    if analysis_match:
+        analysis = analysis_match.group(1).strip()
+    summary_match = re.search(r"<summary>(.*?)</summary>", response_text, re.DOTALL)
+    if summary_match:
+        summary_text = summary_match.group(1).strip()
 
-    There may be additional summarization guidance provided. If so, remember to follow these instructions when creating the above summary.
-    """
+    return Summary(analysis=analysis, summary=summary_text)
 
 
 def _is_context_length_error(error: BaseException) -> bool:
@@ -273,8 +272,8 @@ def find_summarization_boundary(
 
     Walks messages from the start and finds the latest safe split point that
     leaves at least ``min_messages_to_keep`` messages in the "keep" portion.
-    Safe boundaries are after a simple assistant message (no tool calls) or
-    after the last tool result in a consecutive tool-result sequence.
+    A safe boundary is after a simple assistant message (no tool calls) —
+    this is the natural end of a complete conversational turn.
 
     Returns:
         Index splitting ``messages[:boundary]`` (to summarize) from
@@ -288,10 +287,7 @@ def find_summarization_boundary(
         is_simple_assistant = message.role == "assistant" and not getattr(
             message, "tool_calls", None
         )
-        is_last_tool = message.role == "tool" and (
-            i + 1 == len(messages) or messages[i + 1].role != "tool"
-        )
-        if is_simple_assistant or is_last_tool:
+        if is_simple_assistant:
             best_boundary = i + 1
     return best_boundary
 
@@ -371,7 +367,8 @@ def _make_summarize_hook(
         if not to_summarize:
             return None
 
-        summary = await summarize_conversation.bind(summarizer_model)(  # ty: ignore[unresolved-attribute]  # .bind() added by generator framework at runtime
+        summary = await summarize_conversation(
+            summarizer_model,
             "\n".join(str(msg) for msg in to_summarize),
             guidance=guidance,
         )
@@ -398,22 +395,3 @@ def _make_summarize_hook(
 
 # Pre-instantiated with defaults -- this is what the wrapper discovers
 summarize_when_long = _make_summarize_hook()
-
-
-# =============================================================================
-# Default Hooks
-# =============================================================================
-
-
-def default_hooks() -> list:
-    """Core hooks every agent should have."""
-    from dreadnode.agents import stopping
-
-    hooks: list = [
-        tool_metrics(),
-        summarize_when_long,
-        stopping.step_count(50, name="server_step_limit"),
-    ]
-    if backoff_on_ratelimit is not None:
-        hooks.insert(0, backoff_on_ratelimit)
-    return hooks

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/mcp/auth.py

@@ -5,6 +5,7 @@ Provides file-based token storage and helpers for building OAuth providers
 using the MCP SDK's built-in OAuthClientProvider.
 """
 
+import asyncio
 import json
 import os
 import typing as t
@@ -64,7 +65,8 @@ class FileTokenStorage:
     async def get_tokens(self) -> "OAuthToken | None":
         from mcp.shared.auth import OAuthToken
 
-        data = self._get_entry().get("tokens")
+        entry = await asyncio.to_thread(self._get_entry)
+        data = entry.get("tokens")
         if data is None:
             return None
         try:
@@ -74,14 +76,18 @@ class FileTokenStorage:
             return None
 
     async def set_tokens(self, tokens: "OAuthToken") -> None:
-        entry = self._get_entry()
-        entry["tokens"] = tokens.model_dump(mode="json", exclude_none=True)
-        self._set_entry(entry)
+        def _update() -> None:
+            entry = self._get_entry()
+            entry["tokens"] = tokens.model_dump(mode="json", exclude_none=True)
+            self._set_entry(entry)
+
+        await asyncio.to_thread(_update)
 
     async def get_client_info(self) -> "OAuthClientInformationFull | None":
         from mcp.shared.auth import OAuthClientInformationFull
 
-        data = self._get_entry().get("client_info")
+        entry = await asyncio.to_thread(self._get_entry)
+        data = entry.get("client_info")
         if data is None:
             return None
         try:
@@ -91,9 +97,12 @@ class FileTokenStorage:
             return None
 
     async def set_client_info(self, client_info: "OAuthClientInformationFull") -> None:
-        entry = self._get_entry()
-        entry["client_info"] = client_info.model_dump(mode="json", exclude_none=True)
-        self._set_entry(entry)
+        def _update() -> None:
+            entry = self._get_entry()
+            entry["client_info"] = client_info.model_dump(mode="json", exclude_none=True)
+            self._set_entry(entry)
+
+        await asyncio.to_thread(_update)
 
 
 async def _default_redirect_handler(url: str) -> None:

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/mcp/client.py

@@ -254,7 +254,7 @@ class MCPClient:
             error_msg = f"{error_msg}\n\nServer stderr:\n{stderr_tail}"
 
         error_lower = error_msg.lower()
-        if "unauthorized" in error_lower or "401" in error_lower or "oauth" in error_lower:
+        if any(kw in error_lower for kw in ("unauthorized", "forbidden", "401", "403", "oauth")):
             self._status = MCPStatus.NEEDS_AUTH
         else:
             self._status = MCPStatus.FAILED
@@ -363,9 +363,9 @@ class MCPClient:
         """Probe whether a URL supports streamable HTTP (POST).
 
         Returns True if the server likely supports streamable HTTP, False if it
-        clearly doesn't (4xx excluding auth errors). Connectivity errors also
-        return False (server unreachable via POST, try SSE). Auth and SSL errors
-        are allowed to propagate — they would affect SSE equally.
+        clearly doesn't. Connectivity errors also return False (server
+        unreachable via POST, try SSE). Auth and SSL errors are allowed to
+        propagate — they would affect SSE equally.
         """
         try:
             probe_headers = dict(headers) if headers else {}
@@ -377,6 +377,16 @@ class MCPClient:
             )
             async with httpx.AsyncClient(timeout=min(http_timeout, 5), auth=auth) as probe:
                 r = await probe.post(url, content=b"{}", headers=probe_headers)
+                # Auth errors affect both transports equally — propagate so
+                # callers can classify as NEEDS_AUTH instead of falling through.
+                if r.status_code in (401, 403):
+                    r.raise_for_status()
+                if r.status_code == 400 and self._looks_like_jsonrpc_probe_rejection(r):
+                    logger.debug(
+                        "Streamable HTTP probe got JSON-RPC validation error for {}, using streamable-http",
+                        url,
+                    )
+                    return True
                 # 400 = bad request (SSE-only server doesn't understand POST body)
                 # 404 = endpoint not found for POST
                 # 405 = method not allowed (explicitly rejects POST)
@@ -394,6 +404,24 @@ class MCPClient:
 
         return True
 
+    @staticmethod
+    def _looks_like_jsonrpc_probe_rejection(response: httpx.Response) -> bool:
+        """Detect JSON-RPC validation errors from real MCP streamable-http servers."""
+        content_type = response.headers.get("content-type", "").lower()
+        if "application/json" not in content_type:
+            return False
+
+        with contextlib.suppress(ValueError):
+            payload = response.json()
+            if (
+                isinstance(payload, dict)
+                and payload.get("jsonrpc") == "2.0"
+                and isinstance(payload.get("error"), dict)
+            ):
+                return True
+
+        return False
+
     async def _connect_via_streamable_http(self, connection: dict[str, t.Any]) -> "ClientSession":
         """Connect via streamable HTTP, falling back to SSE on failure."""
         from mcp import ClientSession

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/subagent.py

@@ -172,14 +172,18 @@ class SubAgentToolset(Toolset):
         sub_agent.reset()
 
         try:
-            response = await sub_agent.chat(task)
+            trajectory = await sub_agent.run(task)
+            # Get the last assistant message content as the response
+            last_message = trajectory.messages[-1] if trajectory.messages else None
+            response = str(last_message.content) if last_message else ""
+
             result = f"## Sub-agent: {config['name']}\n\n"
             result += f"**Task:** {task}\n\n"
-            result += f"**Steps taken:** {len(sub_agent.trajectory.steps)}\n"
-            result += f"**Tokens used:** {sub_agent.trajectory.usage.total_tokens}\n\n"
+            result += f"**Steps taken:** {len(trajectory.steps)}\n"
+            result += f"**Tokens used:** {trajectory.usage.total_tokens}\n\n"
             result += f"**Result:**\n{response}"
 
-            logger.info(f"Sub-agent completed in {len(sub_agent.trajectory.steps)} steps")
+            logger.info(f"Sub-agent completed in {len(trajectory.steps)} steps")
         except Exception as e:
             logger.error(f"Sub-agent failed: {e}")
             return f"Sub-agent failed: {e}"

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/tools.py

@@ -7,6 +7,7 @@ import re
 import typing as t
 from pathlib import Path
 
+import aiofiles
 import typing_extensions as te
 from pydantic import (
     BaseModel,
@@ -162,21 +163,24 @@ class ToolResponse(XMLModel):
     result: str
 
 
-def offload_tool_output(
+async def offload_tool_output(
     content: str,
     tool_call_id: str,
     tool_name: str,  # noqa: ARG001 - callers pass as keyword arg
     working_dir: Path,
 ) -> tuple[str, Path]:
     """Write tool output to disk and return middle-out summary plus file path."""
-    output_dir = working_dir / ".dreadnode" / "tool-output"
-    output_dir.mkdir(parents=True, exist_ok=True)
+    output_root = working_dir if working_dir.name == ".dreadnode" else working_dir / ".dreadnode"
+    output_dir = output_root / "tool-output"
+    await asyncio.to_thread(output_dir.mkdir, parents=True, exist_ok=True)
 
     safe_id = re.sub(r"[^a-zA-Z0-9_-]", "_", tool_call_id)
     file_path = output_dir / f"{safe_id}.txt"
-    file_path.write_text(content)
+    async with aiofiles.open(file_path, "w") as f:
+        await f.write(content)
 
-    relative_path = file_path.relative_to(working_dir)
+    relative_base = working_dir.parent if working_dir.name == ".dreadnode" else working_dir
+    relative_path = file_path.relative_to(relative_base)
 
     half = OFFLOAD_THRESHOLD // 2
     start = content[:half]
@@ -422,7 +426,9 @@ class Tool(BaseModel, t.Generic[P, R]):
         else:
             with contextlib.suppress(Exception):
                 if type(result) not in [str, int, float, bool]:
-                    result = TypeAdapter(t.Any).dump_json(result).decode(errors="replace")
+                    result = await asyncio.to_thread(
+                        lambda: TypeAdapter(t.Any).dump_json(result).decode(errors="replace")
+                    )
             message.content_parts = [ContentText(text=str(result))]
 
         if self.truncate:
@@ -432,7 +438,7 @@ class Tool(BaseModel, t.Generic[P, R]):
 
         if self.offload and message.content and len(message.content) > OFFLOAD_THRESHOLD:
             working_dir = self.working_dir or Path.cwd()
-            truncated_content, output_file = offload_tool_output(
+            truncated_content, output_file = await offload_tool_output(
                 message.content,
                 tool_call.id,
                 tool_call.name,

{dreadnode-2.0.8 → dreadnode-2.0.10}/dreadnode/agents/trajectory.py

@@ -297,10 +297,18 @@ def trajectory_from_openai_format(
     step_num = 0
 
     for msg in messages:
-        message = message_class(
-            role=msg.get("role", "user"),
-            content=msg.get("content", ""),
-        )
+        kwargs: dict[str, t.Any] = {
+            "role": msg.get("role", "user"),
+            "content": msg.get("content", ""),
+        }
+        # Preserve tool calls (assistant messages) and tool_call_id (tool responses)
+        if msg.get("tool_calls"):
+            kwargs["tool_calls"] = msg["tool_calls"]
+        if msg.get("tool_call_id"):
+            kwargs["tool_call_id"] = msg["tool_call_id"]
+        if msg.get("metadata"):
+            kwargs["metadata"] = msg["metadata"]
+        message = message_class(**kwargs)
         current_messages.append(message)
 
         # Create a step when we see an assistant message (the natural boundary)