Skip to content

Commit b68d154

Browse files
dicksontsairushilpatel0
authored andcommitted
Fix json error handling
1 parent bce29e4 commit b68d154

File tree

6 files changed

+159
-28
lines changed

6 files changed

+159
-28
lines changed

src/claude_code_sdk/_errors.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,3 +44,11 @@ def __init__(self, line: str, original_error: Exception):
4444
self.line = line
4545
self.original_error = original_error
4646
super().__init__(f"Failed to decode JSON: {line[:100]}...")
47+
48+
49+
class MessageParseError(ClaudeSDKError):
50+
"""Raised when unable to parse a message from CLI output."""
51+
52+
def __init__(self, message: str, data: dict | None = None):
53+
self.data = data
54+
super().__init__(message)

src/claude_code_sdk/_internal/client.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,8 @@ async def process_query(
4949
await chosen_transport.connect()
5050

5151
async for data in chosen_transport.receive_messages():
52-
message = parse_message(data)
53-
if message:
54-
yield message
52+
yield parse_message(data)
53+
5554

5655
finally:
5756
await chosen_transport.disconnect()

src/claude_code_sdk/_internal/message_parser.py

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import logging
44
from typing import Any
55

6+
from .._errors import MessageParseError
67
from ..types import (
78
AssistantMessage,
89
ContentBlock,
@@ -18,33 +19,37 @@
1819
logger = logging.getLogger(__name__)
1920

2021

21-
def parse_message(data: dict[str, Any]) -> Message | None:
22+
def parse_message(data: dict[str, Any]) -> Message:
2223
"""
2324
Parse message from CLI output into typed Message objects.
2425
2526
Args:
2627
data: Raw message dictionary from CLI output
2728
2829
Returns:
29-
Parsed Message object or None if type is unrecognized or parsing fails
30+
Parsed Message object
31+
32+
Raises:
33+
MessageParseError: If parsing fails or message type is unrecognized
3034
"""
31-
try:
32-
message_type = data.get("type")
33-
if not message_type:
34-
logger.warning("Message missing 'type' field: %s", data)
35-
return None
35+
if not isinstance(data, dict):
36+
raise MessageParseError(
37+
f"Invalid message data type (expected dict, got {type(data).__name__})",
38+
data,
39+
)
3640

37-
except AttributeError:
38-
logger.error("Invalid message data type (expected dict): %s", type(data))
39-
return None
41+
message_type = data.get("type")
42+
if not message_type:
43+
raise MessageParseError("Message missing 'type' field", data)
4044

4145
match message_type:
4246
case "user":
4347
try:
4448
return UserMessage(content=data["message"]["content"])
4549
except KeyError as e:
46-
logger.error("Missing required field in user message: %s", e)
47-
return None
50+
raise MessageParseError(
51+
f"Missing required field in user message: {e}", data
52+
) from e
4853

4954
case "assistant":
5055
try:
@@ -72,8 +77,9 @@ def parse_message(data: dict[str, Any]) -> Message | None:
7277

7378
return AssistantMessage(content=content_blocks)
7479
except KeyError as e:
75-
logger.error("Missing required field in assistant message: %s", e)
76-
return None
80+
raise MessageParseError(
81+
f"Missing required field in assistant message: {e}", data
82+
) from e
7783

7884
case "system":
7985
try:
@@ -82,8 +88,9 @@ def parse_message(data: dict[str, Any]) -> Message | None:
8288
data=data,
8389
)
8490
except KeyError as e:
85-
logger.error("Missing required field in system message: %s", e)
86-
return None
91+
raise MessageParseError(
92+
f"Missing required field in system message: {e}", data
93+
) from e
8794

8895
case "result":
8996
try:
@@ -99,9 +106,9 @@ def parse_message(data: dict[str, Any]) -> Message | None:
99106
result=data.get("result"),
100107
)
101108
except KeyError as e:
102-
logger.error("Missing required field in result message: %s", e)
103-
return None
109+
raise MessageParseError(
110+
f"Missing required field in result message: {e}", data
111+
) from e
104112

105113
case _:
106-
logger.debug("Unknown message type: %s", message_type)
107-
return None
114+
raise MessageParseError(f"Unknown message type: {message_type}", data)

src/claude_code_sdk/_internal/transport/subprocess_cli.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -307,8 +307,9 @@ async def receive_messages(self) -> AsyncIterator[dict[str, Any]]:
307307
except GeneratorExit:
308308
return
309309
except json.JSONDecodeError:
310-
# Don't clear buffer - we might be in the middle of a split JSON message
311-
# The buffer will be cleared when we successfully parse or hit size limit
310+
# We are speculatively decoding the buffer until we get
311+
# a full JSON object. If there is an actual issue, we
312+
# raise an error after _MAX_BUFFER_SIZE.
312313
continue
313314

314315
except anyio.ClosedResourceError:

src/claude_code_sdk/client.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -126,9 +126,7 @@ async def receive_messages(self) -> AsyncIterator[Message]:
126126
from ._internal.message_parser import parse_message
127127

128128
async for data in self._transport.receive_messages():
129-
message = parse_message(data)
130-
if message:
131-
yield message
129+
yield parse_message(data)
132130

133131
async def query(
134132
self, prompt: str | AsyncIterable[dict[str, Any]], session_id: str = "default"

tests/test_message_parser.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
"""Tests for message parser error handling."""
2+
3+
import pytest
4+
5+
from claude_code_sdk._errors import MessageParseError
6+
from claude_code_sdk._internal.message_parser import parse_message
7+
from claude_code_sdk.types import (
8+
AssistantMessage,
9+
ResultMessage,
10+
SystemMessage,
11+
TextBlock,
12+
ToolUseBlock,
13+
UserMessage,
14+
)
15+
16+
17+
class TestMessageParser:
18+
"""Test message parsing with the new exception behavior."""
19+
20+
def test_parse_valid_user_message(self):
21+
"""Test parsing a valid user message."""
22+
data = {"type": "user", "message": {"content": [{"type": "text", "text": "Hello"}]}}
23+
message = parse_message(data)
24+
assert isinstance(message, UserMessage)
25+
26+
def test_parse_valid_assistant_message(self):
27+
"""Test parsing a valid assistant message."""
28+
data = {
29+
"type": "assistant",
30+
"message": {
31+
"content": [
32+
{"type": "text", "text": "Hello"},
33+
{
34+
"type": "tool_use",
35+
"id": "tool_123",
36+
"name": "Read",
37+
"input": {"file_path": "/test.txt"},
38+
},
39+
]
40+
},
41+
}
42+
message = parse_message(data)
43+
assert isinstance(message, AssistantMessage)
44+
assert len(message.content) == 2
45+
assert isinstance(message.content[0], TextBlock)
46+
assert isinstance(message.content[1], ToolUseBlock)
47+
48+
def test_parse_valid_system_message(self):
49+
"""Test parsing a valid system message."""
50+
data = {"type": "system", "subtype": "start"}
51+
message = parse_message(data)
52+
assert isinstance(message, SystemMessage)
53+
assert message.subtype == "start"
54+
55+
def test_parse_valid_result_message(self):
56+
"""Test parsing a valid result message."""
57+
data = {
58+
"type": "result",
59+
"subtype": "success",
60+
"duration_ms": 1000,
61+
"duration_api_ms": 500,
62+
"is_error": False,
63+
"num_turns": 2,
64+
"session_id": "session_123",
65+
}
66+
message = parse_message(data)
67+
assert isinstance(message, ResultMessage)
68+
assert message.subtype == "success"
69+
70+
def test_parse_invalid_data_type(self):
71+
"""Test that non-dict data raises MessageParseError."""
72+
with pytest.raises(MessageParseError) as exc_info:
73+
parse_message("not a dict") # type: ignore
74+
assert "Invalid message data type" in str(exc_info.value)
75+
assert "expected dict, got str" in str(exc_info.value)
76+
77+
def test_parse_missing_type_field(self):
78+
"""Test that missing 'type' field raises MessageParseError."""
79+
with pytest.raises(MessageParseError) as exc_info:
80+
parse_message({"message": {"content": []}})
81+
assert "Message missing 'type' field" in str(exc_info.value)
82+
83+
def test_parse_unknown_message_type(self):
84+
"""Test that unknown message type raises MessageParseError."""
85+
with pytest.raises(MessageParseError) as exc_info:
86+
parse_message({"type": "unknown_type"})
87+
assert "Unknown message type: unknown_type" in str(exc_info.value)
88+
89+
def test_parse_user_message_missing_fields(self):
90+
"""Test that user message with missing fields raises MessageParseError."""
91+
with pytest.raises(MessageParseError) as exc_info:
92+
parse_message({"type": "user"})
93+
assert "Missing required field in user message" in str(exc_info.value)
94+
95+
def test_parse_assistant_message_missing_fields(self):
96+
"""Test that assistant message with missing fields raises MessageParseError."""
97+
with pytest.raises(MessageParseError) as exc_info:
98+
parse_message({"type": "assistant"})
99+
assert "Missing required field in assistant message" in str(exc_info.value)
100+
101+
def test_parse_system_message_missing_fields(self):
102+
"""Test that system message with missing fields raises MessageParseError."""
103+
with pytest.raises(MessageParseError) as exc_info:
104+
parse_message({"type": "system"})
105+
assert "Missing required field in system message" in str(exc_info.value)
106+
107+
def test_parse_result_message_missing_fields(self):
108+
"""Test that result message with missing fields raises MessageParseError."""
109+
with pytest.raises(MessageParseError) as exc_info:
110+
parse_message({"type": "result", "subtype": "success"})
111+
assert "Missing required field in result message" in str(exc_info.value)
112+
113+
def test_message_parse_error_contains_data(self):
114+
"""Test that MessageParseError contains the original data."""
115+
data = {"type": "unknown", "some": "data"}
116+
with pytest.raises(MessageParseError) as exc_info:
117+
parse_message(data)
118+
assert exc_info.value.data == data

0 commit comments

Comments
 (0)