Merge pull request #60 from anthropics/fix-fastapi-streaming-issue-4

Fix FastAPI SSE streaming compatibility

src/claude_code_sdk/_internal/transport/subprocess_cli.py

@@ -1,6 +1,7 @@
 """Subprocess transport implementation using Claude Code CLI."""
 
 import json
+import logging
 import os
 import shutil
 from collections.abc import AsyncIterator
@@ -17,6 +18,8 @@ from ..._errors import CLIJSONDecodeError as SDKJSONDecodeError
 from ...types import ClaudeCodeOptions
 from . import Transport
 
+logger = logging.getLogger(__name__)
+
 _MAX_BUFFER_SIZE = 1024 * 1024  # 1MB buffer limit
 
 
@@ -170,22 +173,13 @@ class SubprocessCLITransport(Transport):
         if not self._process or not self._stdout_stream:
             raise CLIConnectionError("Not connected")
 
-        stderr_lines = []
-
-        async def read_stderr() -> None:
-            """Read stderr in background."""
-            if self._stderr_stream:
-                try:
-                    async for line in self._stderr_stream:
-                        stderr_lines.append(line.strip())
-                except anyio.ClosedResourceError:
-                    pass
-
-        async with anyio.create_task_group() as tg:
-            tg.start_soon(read_stderr)
+        # Safety constants
+        max_stderr_size = 10 * 1024 * 1024  # 10MB
+        stderr_timeout = 30.0  # 30 seconds
 
         json_buffer = ""
 
+        # Process stdout messages first
         try:
             async for line in self._stdout_stream:
                 line_str = line.strip()
@@ -223,16 +217,60 @@ class SubprocessCLITransport(Transport):
 
         except anyio.ClosedResourceError:
             pass
+        except GeneratorExit:
+            # Client disconnected - still need to clean up
+            pass
 
-        await self._process.wait()
-        if self._process.returncode is not None and self._process.returncode != 0:
-            stderr_output = "\n".join(stderr_lines)
-            if stderr_output and "error" in stderr_output.lower():
+        # Process stderr with safety limits
+        stderr_lines = []
+        stderr_size = 0
+
+        if self._stderr_stream:
+            try:
+                # Use timeout to prevent hanging
+                with anyio.fail_after(stderr_timeout):
+                    async for line in self._stderr_stream:
+                        line_text = line.strip()
+                        line_size = len(line_text)
+
+                        # Enforce memory limit
+                        if stderr_size + line_size > max_stderr_size:
+                            stderr_lines.append(
+                                f"[stderr truncated after {stderr_size} bytes]"
+                            )
+                            # Drain rest of stream without storing
+                            async for _ in self._stderr_stream:
+                                pass
+                            break
+
+                        stderr_lines.append(line_text)
+                        stderr_size += line_size
+
+            except TimeoutError:
+                stderr_lines.append(
+                    f"[stderr collection timed out after {stderr_timeout}s]"
+                )
+            except anyio.ClosedResourceError:
+                pass
+
+        # Check process completion and handle errors
+        try:
+            returncode = await self._process.wait()
+        except Exception:
+            returncode = -1
+
+        stderr_output = "\n".join(stderr_lines) if stderr_lines else ""
+
+        # Use exit code for error detection, not string matching
+        if returncode is not None and returncode != 0:
             raise ProcessError(
-                    "CLI process failed",
-                    exit_code=self._process.returncode,
+                f"Command failed with exit code {returncode}",
+                exit_code=returncode,
                 stderr=stderr_output,
             )
+        elif stderr_output:
+            # Log stderr for debugging but don't fail on non-zero exit
+            logger.debug(f"Process stderr: {stderr_output}")
 
     def is_connected(self) -> bool:
         """Check if subprocess is running."""

tests/test_subprocess_buffering.py

@@ -255,9 +255,8 @@ class TestSubprocessBuffering:
                 async for msg in transport.receive_messages():
                     messages.append(msg)
 
-            assert len(exc_info.value.exceptions) == 1
-            assert isinstance(exc_info.value.exceptions[0], CLIJSONDecodeError)
-            assert "exceeded maximum buffer size" in str(exc_info.value.exceptions[0])
+            assert isinstance(exc_info.value, CLIJSONDecodeError)
+            assert "exceeded maximum buffer size" in str(exc_info.value)
 
         anyio.run(_test)