diff --git a/src/claude_code_sdk/_internal/client.py b/src/claude_code_sdk/_internal/client.py index 9e8ceac..ef1070d 100644 --- a/src/claude_code_sdk/_internal/client.py +++ b/src/claude_code_sdk/_internal/client.py @@ -47,11 +47,9 @@ class InternalClient: match data["type"]: case "user": - # Extract just the content from the nested structure return UserMessage(content=data["message"]["content"]) case "assistant": - # Parse content blocks content_blocks: list[ContentBlock] = [] for block in data["message"]["content"]: match block["type"]: @@ -79,20 +77,18 @@ class InternalClient: case "system": return SystemMessage( subtype=data["subtype"], - data=data, # Pass through all data + data=data, ) case "result": - # Map total_cost to total_cost_usd for consistency return ResultMessage( subtype=data["subtype"], - cost_usd=data["cost_usd"], duration_ms=data["duration_ms"], duration_api_ms=data["duration_api_ms"], is_error=data["is_error"], num_turns=data["num_turns"], session_id=data["session_id"], - total_cost_usd=data["total_cost"], + total_cost_usd=data.get("total_cost_usd"), usage=data.get("usage"), result=data.get("result"), ) diff --git a/src/claude_code_sdk/_internal/transport/subprocess_cli.py b/src/claude_code_sdk/_internal/transport/subprocess_cli.py index 564bc94..f63732a 100644 --- a/src/claude_code_sdk/_internal/transport/subprocess_cli.py +++ b/src/claude_code_sdk/_internal/transport/subprocess_cli.py @@ -190,7 +190,11 @@ class SubprocessCLITransport(Transport): try: data = json.loads(line_str) - yield data + try: + yield data + except GeneratorExit: + # Handle generator cleanup gracefully + return except json.JSONDecodeError as e: if line_str.startswith("{") or line_str.startswith("["): raise SDKJSONDecodeError(line_str, e) from e @@ -198,8 +202,6 @@ class SubprocessCLITransport(Transport): except anyio.ClosedResourceError: pass - finally: - tg.cancel_scope.cancel() await self._process.wait() if self._process.returncode is not None and self._process.returncode != 0: diff --git a/src/claude_code_sdk/types.py b/src/claude_code_sdk/types.py index a80510b..21ae8ff 100644 --- a/src/claude_code_sdk/types.py +++ b/src/claude_code_sdk/types.py @@ -75,13 +75,12 @@ class ResultMessage: """Result message with cost and usage information.""" subtype: str - cost_usd: float duration_ms: int duration_api_ms: int is_error: bool num_turns: int session_id: str - total_cost_usd: float + total_cost_usd: float | None = None usage: dict[str, Any] | None = None result: str | None = None diff --git a/tests/test_client.py b/tests/test_client.py index 8ade785..3282ea1 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -88,13 +88,12 @@ class TestQueryFunction: yield { "type": "result", "subtype": "success", - "cost_usd": 0.001, "duration_ms": 1000, "duration_api_ms": 800, "is_error": False, "num_turns": 1, "session_id": "test-session", - "total_cost": 0.001, + "total_cost_usd": 0.001, } mock_transport.receive_messages = mock_receive diff --git a/tests/test_integration.py b/tests/test_integration.py index 0798569..a185335 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -43,13 +43,12 @@ class TestIntegration: yield { "type": "result", "subtype": "success", - "cost_usd": 0.001, "duration_ms": 1000, "duration_api_ms": 800, "is_error": False, "num_turns": 1, "session_id": "test-session", - "total_cost": 0.001, + "total_cost_usd": 0.001, } mock_transport.receive_messages = mock_receive @@ -71,7 +70,7 @@ class TestIntegration: # Check result message assert isinstance(messages[1], ResultMessage) - assert messages[1].cost_usd == 0.001 + assert messages[1].total_cost_usd == 0.001 assert messages[1].session_id == "test-session" anyio.run(_test) @@ -109,13 +108,12 @@ class TestIntegration: yield { "type": "result", "subtype": "success", - "cost_usd": 0.002, "duration_ms": 1500, "duration_api_ms": 1200, "is_error": False, "num_turns": 1, "session_id": "test-session-2", - "total_cost": 0.002, + "total_cost_usd": 0.002, } mock_transport.receive_messages = mock_receive diff --git a/tests/test_types.py b/tests/test_types.py index 7336204..6046292 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -45,7 +45,6 @@ class TestMessageTypes: """Test creating a ResultMessage.""" msg = ResultMessage( subtype="success", - cost_usd=0.01, duration_ms=1500, duration_api_ms=1200, is_error=False, @@ -54,7 +53,7 @@ class TestMessageTypes: total_cost_usd=0.01, ) assert msg.subtype == "success" - assert msg.cost_usd == 0.01 + assert msg.total_cost_usd == 0.01 assert msg.session_id == "session-123"