diff --git a/src/claude_agent_sdk/_internal/sessions.py b/src/claude_agent_sdk/_internal/sessions.py index a49d2c63..874a1784 100644 --- a/src/claude_agent_sdk/_internal/sessions.py +++ b/src/claude_agent_sdk/_internal/sessions.py @@ -833,7 +833,40 @@ def _pick_best(candidates: list[_TranscriptEntry]) -> _TranscriptEntry: chain_cur = by_uuid.get(parent) if parent else None chain.reverse() - return chain + + # Parallel tool calls produce multiple user messages (tool results) that + # share the same parentUuid. The linear walk above follows only one path + # from leaf to root, so sibling tool-result messages are missed. This + # pass finds and inserts those missing siblings in file order. + children_by_parent: dict[str, list[_TranscriptEntry]] = {} + for entry in entries: + p = entry.get("parentUuid") + if p: + children_by_parent.setdefault(p, []).append(entry) + + chain_uuids: set[str] = {e["uuid"] for e in chain} + + expanded: list[_TranscriptEntry] = [] + for entry in chain: + p = entry.get("parentUuid") + if p: + for sib in children_by_parent.get(p, []): + sib_uuid = sib["uuid"] + # Only include user-type siblings (parallel tool results). + # Assistant-type siblings represent conversation branches + # (e.g., retries) and should not be merged. + if ( + sib_uuid not in chain_uuids + and sib.get("type") == "user" + and not sib.get("isSidechain") + and not sib.get("teamName") + and not sib.get("isMeta") + ): + chain_uuids.add(sib_uuid) + expanded.append(sib) + expanded.append(entry) + + return expanded def _is_visible_message(entry: _TranscriptEntry) -> bool: diff --git a/tests/test_sessions.py b/tests/test_sessions.py index 39cb7755..23af7118 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -1062,6 +1062,57 @@ def test_only_progress_entries_returns_empty(self): result = _build_conversation_chain(entries) assert result == [] + def test_parallel_tool_calls_includes_siblings(self): + """Parallel tool results sharing the same parentUuid are all included.""" + entries = [ + {"type": "user", "uuid": "U1", "parentUuid": None}, + {"type": "assistant", "uuid": "A1", "parentUuid": "U1"}, + {"type": "user", "uuid": "U2", "parentUuid": "A1"}, # tool_result 1 + {"type": "user", "uuid": "U3", "parentUuid": "A1"}, # tool_result 2 + {"type": "assistant", "uuid": "A2", "parentUuid": "U3"}, + ] + result = _build_conversation_chain(entries) + uuids = [e["uuid"] for e in result] + assert uuids == ["U1", "A1", "U2", "U3", "A2"] + + def test_parallel_tool_calls_three_siblings(self): + """Three parallel tool results are all included in file order.""" + entries = [ + {"type": "user", "uuid": "U1", "parentUuid": None}, + {"type": "assistant", "uuid": "A1", "parentUuid": "U1"}, + {"type": "user", "uuid": "U2", "parentUuid": "A1"}, # tool_result 1 + {"type": "user", "uuid": "U3", "parentUuid": "A1"}, # tool_result 2 + {"type": "user", "uuid": "U4", "parentUuid": "A1"}, # tool_result 3 + {"type": "assistant", "uuid": "A2", "parentUuid": "U4"}, + ] + result = _build_conversation_chain(entries) + uuids = [e["uuid"] for e in result] + assert uuids == ["U1", "A1", "U2", "U3", "U4", "A2"] + + def test_parallel_siblings_excludes_sidechain(self): + """Sidechain siblings should not be pulled into the main chain.""" + entries = [ + {"type": "user", "uuid": "U1", "parentUuid": None}, + {"type": "assistant", "uuid": "A1", "parentUuid": "U1"}, + {"type": "user", "uuid": "U2", "parentUuid": "A1", "isSidechain": True}, + {"type": "user", "uuid": "U3", "parentUuid": "A1"}, + {"type": "assistant", "uuid": "A2", "parentUuid": "U3"}, + ] + result = _build_conversation_chain(entries) + uuids = [e["uuid"] for e in result] + assert uuids == ["U1", "A1", "U3", "A2"] + + def test_linear_chain_unchanged_with_sibling_fix(self): + """Linear chains (no parallel calls) produce identical results.""" + entries = [ + {"type": "user", "uuid": "a", "parentUuid": None}, + {"type": "assistant", "uuid": "b", "parentUuid": "a"}, + {"type": "user", "uuid": "c", "parentUuid": "b"}, + {"type": "assistant", "uuid": "d", "parentUuid": "c"}, + ] + result = _build_conversation_chain(entries) + assert [e["uuid"] for e in result] == ["a", "b", "c", "d"] + class TestSessionMessageType: """Tests for the SessionMessage dataclass."""