feat: expose parent_tool_use_id to support subagent tracking #138 (#166)

This commit is contained in:
Dan Siwiec
2025-09-17 14:47:09 -07:00
committed by GitHub
parent fd98d12f94
commit 607921dcb0
3 changed files with 47 additions and 3 deletions

View File

@@ -46,6 +46,7 @@ def parse_message(data: dict[str, Any]) -> Message:
match message_type:
case "user":
try:
parent_tool_use_id = data.get("parent_tool_use_id")
if isinstance(data["message"]["content"], list):
user_content_blocks: list[ContentBlock] = []
for block in data["message"]["content"]:
@@ -70,8 +71,14 @@ def parse_message(data: dict[str, Any]) -> Message:
is_error=block.get("is_error"),
)
)
return UserMessage(content=user_content_blocks)
return UserMessage(content=data["message"]["content"])
return UserMessage(
content=user_content_blocks,
parent_tool_use_id=parent_tool_use_id,
)
return UserMessage(
content=data["message"]["content"],
parent_tool_use_id=parent_tool_use_id,
)
except KeyError as e:
raise MessageParseError(
f"Missing required field in user message: {e}", data
@@ -109,7 +116,9 @@ def parse_message(data: dict[str, Any]) -> Message:
)
return AssistantMessage(
content=content_blocks, model=data["message"]["model"]
content=content_blocks,
model=data["message"]["model"],
parent_tool_use_id=data.get("parent_tool_use_id"),
)
except KeyError as e:
raise MessageParseError(

View File

@@ -230,6 +230,7 @@ class UserMessage:
"""User message."""
content: str | list[ContentBlock]
parent_tool_use_id: str | None = None
@dataclass
@@ -238,6 +239,7 @@ class AssistantMessage:
content: list[ContentBlock]
model: str
parent_tool_use_id: str | None = None
@dataclass

View File

@@ -130,6 +130,17 @@ class TestMessageParser:
assert isinstance(message.content[2], ToolResultBlock)
assert isinstance(message.content[3], TextBlock)
def test_parse_user_message_inside_subagent(self):
"""Test parsing a valid user message."""
data = {
"type": "user",
"message": {"content": [{"type": "text", "text": "Hello"}]},
"parent_tool_use_id": "toolu_01Xrwd5Y13sEHtzScxR77So8",
}
message = parse_message(data)
assert isinstance(message, UserMessage)
assert message.parent_tool_use_id == "toolu_01Xrwd5Y13sEHtzScxR77So8"
def test_parse_valid_assistant_message(self):
"""Test parsing a valid assistant message."""
data = {
@@ -185,6 +196,28 @@ class TestMessageParser:
assert isinstance(message, SystemMessage)
assert message.subtype == "start"
def test_parse_assistant_message_inside_subagent(self):
"""Test parsing a valid assistant message."""
data = {
"type": "assistant",
"message": {
"content": [
{"type": "text", "text": "Hello"},
{
"type": "tool_use",
"id": "tool_123",
"name": "Read",
"input": {"file_path": "/test.txt"},
},
],
"model": "claude-opus-4-1-20250805",
},
"parent_tool_use_id": "toolu_01Xrwd5Y13sEHtzScxR77So8",
}
message = parse_message(data)
assert isinstance(message, AssistantMessage)
assert message.parent_tool_use_id == "toolu_01Xrwd5Y13sEHtzScxR77So8"
def test_parse_valid_result_message(self):
"""Test parsing a valid result message."""
data = {