From 18e5899d33e637d50bb3012e11aed90800d001c2 Mon Sep 17 00:00:00 2001
From: unclecode <unclecode@kidocode.com>
Date: Tue, 17 Jun 2025 15:59:04 +0800
Subject: [PATCH] Fix KeyError: 'cost_usd' for Max subscription users

This fix addresses issue #2 where Max subscription users encounter a KeyError
when the API response doesn't include the 'cost_usd' field.

Changes:
- Modified ResultMessage dataclass to make cost_usd and total_cost_usd optional (float | None)
- Updated client.py to use .get() with None default for cost-related fields
- Maintains backward compatibility for users with cost data

This ensures the SDK works for all subscription types, including Max subscriptions
that don't have access to cost information.

By: Unclecode <https://github.com/unclecode>
---
 src/claude_code_sdk/_internal/client.py |  5 +++--
 src/claude_code_sdk/types.py            |  4 ++--
 test_fix.py                             | 25 +++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 test_fix.py

diff --git a/src/claude_code_sdk/_internal/client.py b/src/claude_code_sdk/_internal/client.py
index 9e8ceac..ab50bef 100644
--- a/src/claude_code_sdk/_internal/client.py
+++ b/src/claude_code_sdk/_internal/client.py
@@ -84,15 +84,16 @@ class InternalClient:
 
             case "result":
                 # Map total_cost to total_cost_usd for consistency
+                # Handle missing cost_usd for Max subscription users
                 return ResultMessage(
                     subtype=data["subtype"],
-                    cost_usd=data["cost_usd"],
+                    cost_usd=data.get("cost_usd", None),
                     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", None),
                     usage=data.get("usage"),
                     result=data.get("result"),
                 )
diff --git a/src/claude_code_sdk/types.py b/src/claude_code_sdk/types.py
index a80510b..066fcbb 100644
--- a/src/claude_code_sdk/types.py
+++ b/src/claude_code_sdk/types.py
@@ -75,13 +75,13 @@ class ResultMessage:
     """Result message with cost and usage information."""
 
     subtype: str
-    cost_usd: float
+    cost_usd: float | None
     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
     usage: dict[str, Any] | None = None
     result: str | None = None
 
diff --git a/test_fix.py b/test_fix.py
new file mode 100644
index 0000000..ef8064c
--- /dev/null
+++ b/test_fix.py
@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+"""Test script to verify the cost_usd fix works."""
+
+import anyio
+from claude_code_sdk import process_query, ClaudeCodeOptions
+
+async def main():
+    """Test the SDK with a simple query."""
+    print("Testing claude-code-sdk with cost_usd fix...")
+    
+    try:
+        # Simple test query
+        async for message in process_query(
+            "What is 2+2?",
+            ClaudeCodeOptions()
+        ):
+            print(f"Message type: {type(message).__name__}")
+            print(f"Message: {message}")
+            print("-" * 40)
+    except Exception as e:
+        print(f"Error: {type(e).__name__}: {e}")
+        raise
+
+if __name__ == "__main__":
+    anyio.run(main)
\ No newline at end of file