ccproxy-api 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

ccproxy/claude_sdk/stream_worker.py

@@ -0,0 +1,392 @@
+"""Stream worker for consuming Claude SDK messages and distributing via queue."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from ccproxy.claude_sdk.exceptions import StreamTimeoutError
+from ccproxy.claude_sdk.message_queue import MessageQueue
+from ccproxy.models import claude_sdk as sdk_models
+
+
+if TYPE_CHECKING:
+    from ccproxy.claude_sdk.session_client import SessionClient
+    from ccproxy.claude_sdk.stream_handle import StreamHandle
+
+logger = structlog.get_logger(__name__)
+
+
+class WorkerStatus(str, Enum):
+    """Status of the stream worker."""
+
+    IDLE = "idle"
+    STARTING = "starting"
+    RUNNING = "running"
+    DRAINING = "draining"
+    COMPLETED = "completed"
+    ERROR = "error"
+    INTERRUPTED = "interrupted"
+
+
+class StreamWorker:
+    """Worker that consumes messages from Claude SDK and distributes via queue."""
+
+    def __init__(
+        self,
+        worker_id: str,
+        message_iterator: AsyncIterator[Any],
+        session_id: str | None = None,
+        request_id: str | None = None,
+        session_client: SessionClient | None = None,
+        stream_handle: StreamHandle | None = None,
+    ):
+        """Initialize the stream worker.
+
+        Args:
+            worker_id: Unique identifier for this worker
+            message_iterator: Async iterator of SDK messages
+            session_id: Optional session ID for logging
+            request_id: Optional request ID for logging
+            session_client: Optional session client for state management
+            stream_handle: Optional stream handle for message lifecycle tracking
+        """
+        self.worker_id = worker_id
+        self._message_iterator = message_iterator
+        self.session_id = session_id
+        self.request_id = request_id
+        self._session_client = session_client
+        self._stream_handle = stream_handle
+
+        # Worker state
+        self.status = WorkerStatus.IDLE
+        self._message_queue = MessageQueue()
+        self._worker_task: asyncio.Task[None] | None = None
+        self._started_at: float | None = None
+        self._completed_at: float | None = None
+
+        # Statistics
+        self._total_messages = 0
+        self._messages_delivered = 0
+        self._messages_discarded = 0
+        self._last_message_time: float | None = None
+
+    async def start(self) -> None:
+        """Start the worker task."""
+        if self.status != WorkerStatus.IDLE:
+            logger.warning(
+                "stream_worker_already_started",
+                worker_id=self.worker_id,
+                status=self.status,
+            )
+            return
+
+        self.status = WorkerStatus.STARTING
+        self._started_at = time.time()
+
+        # Create worker task
+        self._worker_task = asyncio.create_task(self._run_worker())
+
+        logger.debug(
+            "stream_worker_started",
+            worker_id=self.worker_id,
+            session_id=self.session_id,
+            request_id=self.request_id,
+        )
+
+    async def stop(self, timeout: float = 5.0) -> None:
+        """Stop the worker gracefully.
+
+        Args:
+            timeout: Maximum time to wait for worker to stop
+        """
+        if self._worker_task and not self._worker_task.done():
+            logger.debug(
+                "stream_worker_stopping",
+                worker_id=self.worker_id,
+                timeout=timeout,
+            )
+
+            # Cancel the worker task
+            self._worker_task.cancel()
+
+            try:
+                # Use asyncio.wait instead of wait_for to handle cancelled tasks properly
+                done, pending = await asyncio.wait(
+                    [self._worker_task],
+                    timeout=timeout,
+                    return_when=asyncio.ALL_COMPLETED,
+                )
+
+                if pending:
+                    logger.warning(
+                        "stream_worker_stop_timeout",
+                        worker_id=self.worker_id,
+                        timeout=timeout,
+                    )
+                elif done:
+                    # Task completed (likely with CancelledError)
+                    logger.debug(
+                        "stream_worker_stopped",
+                        worker_id=self.worker_id,
+                        task_cancelled=self._worker_task.cancelled(),
+                    )
+
+            except Exception as e:
+                logger.warning(
+                    "stream_worker_stop_error",
+                    worker_id=self.worker_id,
+                    error=str(e),
+                    error_type=type(e).__name__,
+                )
+
+    async def wait_for_completion(self, timeout: float | None = None) -> bool:
+        """Wait for the worker to complete.
+
+        Args:
+            timeout: Optional timeout in seconds
+
+        Returns:
+            True if completed successfully, False if timed out
+        """
+        if not self._worker_task:
+            return True
+
+        try:
+            if timeout:
+                await asyncio.wait_for(self._worker_task, timeout=timeout)
+            else:
+                await self._worker_task
+            return True
+        except TimeoutError:
+            return False
+
+    def get_message_queue(self) -> MessageQueue:
+        """Get the message queue for creating listeners.
+
+        Returns:
+            The worker's message queue
+        """
+        return self._message_queue
+
+    async def _run_worker(self) -> None:
+        """Main worker loop that consumes messages and distributes them."""
+        try:
+            self.status = WorkerStatus.RUNNING
+
+            logger.debug(
+                "stream_worker_consuming",
+                worker_id=self.worker_id,
+                session_id=self.session_id,
+            )
+
+            async for message in self._message_iterator:
+                self._total_messages += 1
+                self._last_message_time = time.time()
+
+                # Check if we have listeners
+                if await self._message_queue.has_listeners():
+                    # Broadcast to all listeners
+                    delivered_count = await self._message_queue.broadcast(message)
+                    self._messages_delivered += delivered_count
+
+                    logger.debug(
+                        "stream_worker_message_delivered",
+                        worker_id=self.worker_id,
+                        message_type=type(message).__name__,
+                        delivered_to=delivered_count,
+                        total_messages=self._total_messages,
+                    )
+                else:
+                    # No listeners - discard message
+                    self._messages_discarded += 1
+
+                    logger.debug(
+                        "stream_worker_message_discarded",
+                        worker_id=self.worker_id,
+                        message_type=type(message).__name__,
+                        total_messages=self._total_messages,
+                        total_discarded=self._messages_discarded,
+                    )
+
+                # Update stream handle with message lifecycle tracking
+                if self._stream_handle:
+                    # Track all message activity
+                    self._stream_handle.on_message_received(message)
+
+                    # Track first chunk (SystemMessage with init subtype)
+                    if (
+                        isinstance(message, sdk_models.SystemMessage)
+                        and hasattr(message, "subtype")
+                        and message.subtype == "init"
+                    ):
+                        self._stream_handle.on_first_chunk_received()
+
+                    # Track completion (ResultMessage)
+                    elif isinstance(message, sdk_models.ResultMessage):
+                        self._stream_handle.on_completion()
+
+                # Update session client if we have one
+                if self._session_client and isinstance(
+                    message, sdk_models.ResultMessage
+                ):
+                    self._session_client.sdk_session_id = message.session_id
+
+            # Stream completed successfully
+            self.status = WorkerStatus.COMPLETED
+            await self._message_queue.broadcast_complete()
+
+            logger.debug(
+                "stream_worker_completed",
+                worker_id=self.worker_id,
+                total_messages=self._total_messages,
+                messages_delivered=self._messages_delivered,
+                messages_discarded=self._messages_discarded,
+                duration_seconds=time.time() - (self._started_at or 0),
+            )
+
+        except asyncio.CancelledError:
+            # Worker was cancelled
+            self.status = WorkerStatus.INTERRUPTED
+            logger.debug(
+                "stream_worker_cancelled",
+                worker_id=self.worker_id,
+                messages_processed=self._total_messages,
+            )
+            raise
+
+        except StreamTimeoutError as e:
+            # Handle timeout errors gracefully - these are expected for some commands
+            self.status = WorkerStatus.ERROR
+            await self._message_queue.broadcast_error(e)
+
+            logger.debug(
+                "stream_worker_timeout",
+                worker_id=self.worker_id,
+                timeout_message=str(e),
+                messages_processed=self._total_messages,
+                message="Stream worker completed due to timeout - this is expected for some commands",
+            )
+            # Don't re-raise StreamTimeoutError to avoid unhandled task exceptions
+
+        except Exception as e:
+            # Error during processing (other than timeout)
+            self.status = WorkerStatus.ERROR
+            await self._message_queue.broadcast_error(e)
+
+            logger.error(
+                "stream_worker_error",
+                worker_id=self.worker_id,
+                error=str(e),
+                error_type=type(e).__name__,
+                messages_processed=self._total_messages,
+            )
+            raise
+
+        finally:
+            self._completed_at = time.time()
+
+            # Clean up
+            if self._session_client:
+                self._session_client.has_active_stream = False
+
+            # Close the message queue
+            await self._message_queue.close()
+
+    async def drain_remaining(self, timeout: float = 30.0) -> int:
+        """Drain remaining messages without listeners.
+
+        This is useful for ensuring the stream completes properly
+        even after all listeners have disconnected.
+
+        Args:
+            timeout: Maximum time to spend draining
+
+        Returns:
+            Number of messages drained
+        """
+        if self.status not in (WorkerStatus.RUNNING, WorkerStatus.STARTING):
+            return 0
+
+        self.status = WorkerStatus.DRAINING
+        start_time = time.time()
+        drained_count = 0
+
+        logger.debug(
+            "stream_worker_draining",
+            worker_id=self.worker_id,
+            timeout=timeout,
+        )
+
+        try:
+            # Continue consuming but without broadcasting
+            async for message in self._message_iterator:
+                drained_count += 1
+                self._total_messages += 1
+                self._messages_discarded += 1
+
+                logger.debug(
+                    "stream_worker_draining_message",
+                    worker_id=self.worker_id,
+                    message_type=type(message).__name__,
+                    drained_count=drained_count,
+                )
+
+                # Check timeout
+                if time.time() - start_time > timeout:
+                    logger.warning(
+                        "stream_worker_drain_timeout",
+                        worker_id=self.worker_id,
+                        drained_count=drained_count,
+                        timeout=timeout,
+                    )
+                    break
+
+                # Check for completion message
+                if isinstance(message, sdk_models.ResultMessage):
+                    logger.debug(
+                        "stream_worker_drain_complete",
+                        worker_id=self.worker_id,
+                        drained_count=drained_count,
+                    )
+                    break
+
+        except Exception as e:
+            logger.error(
+                "stream_worker_drain_error",
+                worker_id=self.worker_id,
+                error=str(e),
+                drained_count=drained_count,
+            )
+
+        return drained_count
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get worker statistics.
+
+        Returns:
+            Dictionary of worker statistics
+        """
+        runtime = None
+        if self._started_at:
+            end_time = self._completed_at or time.time()
+            runtime = end_time - self._started_at
+
+        queue_stats = self._message_queue.get_stats()
+
+        return {
+            "worker_id": self.worker_id,
+            "status": self.status.value,
+            "session_id": self.session_id,
+            "request_id": self.request_id,
+            "total_messages": self._total_messages,
+            "messages_delivered": self._messages_delivered,
+            "messages_discarded": self._messages_discarded,
+            "runtime_seconds": runtime,
+            "queue_stats": queue_stats,
+        }

ccproxy/claude_sdk/streaming.py

@@ -139,6 +139,15 @@ class ClaudeStreamProcessor:
                             mode=sdk_message_mode.value,
                             request_id=request_id,
                         )
+                        logger.info(
+                            "sdk_tool_use_block",
+                            tool_id=block.id,
+                            tool_name=block.name,
+                            input_keys=list(block.input.keys()) if block.input else [],
+                            block_index=content_block_index,
+                            mode=sdk_message_mode.value,
+                            request_id=request_id,
+                        )
                         chunks = (
                             self.message_converter._create_sdk_content_block_chunks(
                                 sdk_object=block,
@@ -167,6 +176,20 @@ class ClaudeStreamProcessor:
                             mode=sdk_message_mode.value,
                             request_id=request_id,
                         )
+                        logger.info(
+                            "sdk_tool_result_block",
+                            tool_use_id=block.tool_use_id,
+                            is_error=block.is_error,
+                            content_type=type(block.content).__name__
+                            if block.content
+                            else "None",
+                            content_preview=str(block.content)[:100]
+                            if block.content
+                            else None,
+                            block_index=content_block_index,
+                            mode=sdk_message_mode.value,
+                            request_id=request_id,
+                        )
                         chunks = (
                             self.message_converter._create_sdk_content_block_chunks(
                                 sdk_object=block,
@@ -251,17 +274,18 @@ class ClaudeStreamProcessor:
                         yield chunk
                     content_block_index += 1
 
-                # Final message, contains metrics
-                if ctx:
-                    usage_model = message.usage_model
-                    ctx.add_metadata(
-                        status_code=200,
-                        tokens_input=usage_model.input_tokens,
-                        tokens_output=usage_model.output_tokens,
-                        cache_read_tokens=usage_model.cache_read_input_tokens,
-                        cache_write_tokens=usage_model.cache_creation_input_tokens,
-                        cost_usd=message.total_cost_usd,
-                    )
+                    if ctx:
+                        usage_model = message.usage_model
+                        ctx.add_metadata(
+                            status_code=200,
+                            tokens_input=usage_model.input_tokens,
+                            tokens_output=usage_model.output_tokens,
+                            cache_read_tokens=usage_model.cache_read_input_tokens,
+                            cache_write_tokens=usage_model.cache_creation_input_tokens,
+                            cost_usd=message.total_cost_usd,
+                            session_id=message.session_id,
+                            num_turns=message.num_turns,
+                        )
 
                 end_chunks = self.message_converter.create_streaming_end_chunks(
                     stop_reason=message.stop_reason
@@ -282,5 +306,23 @@ class ClaudeStreamProcessor:
                     message_content=str(message)[:200],
                     request_id=request_id,
                 )
+        else:
+            # Stream ended without a ResultMessage - this indicates an error/interruption
+            if ctx and "status_code" not in ctx.metadata:
+                # Set error status if not already set (e.g., by StreamTimeoutError handler)
+                logger.warning(
+                    "stream_ended_without_result_message",
+                    request_id=request_id,
+                    message="Stream ended without ResultMessage, likely interrupted",
+                )
+                ctx.add_metadata(
+                    status_code=499,  # Client Closed Request
+                    error_type="stream_interrupted",
+                    error_message="Stream ended without completion",
+                )
+
+        # Final message, contains metrics
+        # NOTE: Access logging is now handled by StreamingResponseWithLogging
+        # No need for manual access logging here anymore
 
         logger.debug("claude_sdk_stream_processing_completed", request_id=request_id)

ccproxy/cli/commands/serve.py

@@ -35,7 +35,9 @@ from ..options.claude_options import (
     validate_max_thinking_tokens,
     validate_max_turns,
     validate_permission_mode,
+    validate_pool_size,
     validate_sdk_message_mode,
+    validate_system_prompt_injection_mode,
 )
 from ..options.security_options import SecurityOptions, validate_auth_token
 from ..options.server_options import (
@@ -442,6 +444,48 @@ def api(
             rich_help_panel="Claude Settings",
         ),
     ] = None,
+    sdk_pool: Annotated[
+        bool,
+        typer.Option(
+            "--sdk-pool/--no-sdk-pool",
+            help="Enable/disable general Claude SDK client connection pooling",
+            rich_help_panel="Claude Settings",
+        ),
+    ] = False,
+    sdk_pool_size: Annotated[
+        int | None,
+        typer.Option(
+            "--sdk-pool-size",
+            help="Number of clients to maintain in the general pool (1-20)",
+            callback=validate_pool_size,
+            rich_help_panel="Claude Settings",
+        ),
+    ] = None,
+    sdk_session_pool: Annotated[
+        bool,
+        typer.Option(
+            "--sdk-session-pool/--no-sdk-session-pool",
+            help="Enable/disable session-aware Claude SDK client pooling",
+            rich_help_panel="Claude Settings",
+        ),
+    ] = False,
+    system_prompt_injection_mode: Annotated[
+        str | None,
+        typer.Option(
+            "--system-prompt-injection-mode",
+            help="System prompt injection mode: minimal (Claude Code ID only), full (all detected system messages)",
+            callback=validate_system_prompt_injection_mode,
+            rich_help_panel="Claude Settings",
+        ),
+    ] = None,
+    builtin_permissions: Annotated[
+        bool,
+        typer.Option(
+            "--builtin-permissions/--no-builtin-permissions",
+            help="Enable built-in permission handling infrastructure (MCP server and SSE endpoints). When disabled, users can configure custom MCP servers and permission tools.",
+            rich_help_panel="Claude Settings",
+        ),
+    ] = True,
     # Core settings
     docker: Annotated[
         bool,
@@ -526,6 +570,31 @@ def api(
             rich_help_panel="Docker Settings",
         ),
     ] = None,
+    # Network control flags
+    no_network_calls: Annotated[
+        bool,
+        typer.Option(
+            "--no-network-calls",
+            help="Disable all network calls (version checks and pricing updates)",
+            rich_help_panel="Privacy Settings",
+        ),
+    ] = False,
+    disable_version_check: Annotated[
+        bool,
+        typer.Option(
+            "--disable-version-check",
+            help="Disable version update checks (prevents calls to GitHub API)",
+            rich_help_panel="Privacy Settings",
+        ),
+    ] = False,
+    disable_pricing_updates: Annotated[
+        bool,
+        typer.Option(
+            "--disable-pricing-updates",
+            help="Disable pricing data updates (prevents downloads from GitHub)",
+            rich_help_panel="Privacy Settings",
+        ),
+    ] = False,
 ) -> None:
     """
     Start the CCProxy API server.
@@ -573,10 +642,28 @@ def api(
             cwd=cwd,
             permission_prompt_tool_name=permission_prompt_tool_name,
             sdk_message_mode=sdk_message_mode,
+            sdk_pool=sdk_pool,
+            sdk_pool_size=sdk_pool_size,
+            sdk_session_pool=sdk_session_pool,
+            system_prompt_injection_mode=system_prompt_injection_mode,
+            builtin_permissions=builtin_permissions,
         )
 
         security_options = SecurityOptions(auth_token=auth_token)
 
+        # Handle network control flags
+        scheduler_overrides = {}
+        if no_network_calls:
+            # Disable both network features
+            scheduler_overrides["pricing_update_enabled"] = False
+            scheduler_overrides["version_check_enabled"] = False
+        else:
+            # Handle individual flags
+            if disable_pricing_updates:
+                scheduler_overrides["pricing_update_enabled"] = False
+            if disable_version_check:
+                scheduler_overrides["version_check_enabled"] = False
+
         # Extract CLI overrides from structured option containers
         cli_overrides = config_manager.get_cli_overrides_from_args(
             # Server options
@@ -599,8 +686,17 @@ def api(
             permission_prompt_tool_name=claude_options.permission_prompt_tool_name,
             cwd=claude_options.cwd,
             sdk_message_mode=claude_options.sdk_message_mode,
+            sdk_pool=claude_options.sdk_pool,
+            sdk_pool_size=claude_options.sdk_pool_size,
+            sdk_session_pool=claude_options.sdk_session_pool,
+            system_prompt_injection_mode=claude_options.system_prompt_injection_mode,
+            builtin_permissions=claude_options.builtin_permissions,
         )
 
+        # Add scheduler overrides if any
+        if scheduler_overrides:
+            cli_overrides["scheduler"] = scheduler_overrides
+
         # Load settings with CLI overrides
         settings = config_manager.load_settings(
             config_path=config, cli_overrides=cli_overrides

ccproxy/cli/options/claude_options.py

@@ -93,6 +93,38 @@ def validate_sdk_message_mode(
     return value
 
 
+def validate_pool_size(
+    ctx: typer.Context, param: typer.CallbackParam, value: int | None
+) -> int | None:
+    """Validate pool size."""
+    if value is None:
+        return None
+
+    if value < 1:
+        raise typer.BadParameter("Pool size must be at least 1")
+
+    if value > 20:
+        raise typer.BadParameter("Pool size must not exceed 20")
+
+    return value
+
+
+def validate_system_prompt_injection_mode(
+    ctx: typer.Context, param: typer.CallbackParam, value: str | None
+) -> str | None:
+    """Validate system prompt injection mode."""
+    if value is None:
+        return None
+
+    valid_modes = {"minimal", "full"}
+    if value not in valid_modes:
+        raise typer.BadParameter(
+            f"System prompt injection mode must be one of: {', '.join(valid_modes)}"
+        )
+
+    return value
+
+
 # Factory functions removed - use Annotated syntax directly in commands
 
 
@@ -115,6 +147,11 @@ class ClaudeOptions:
         cwd: str | None = None,
         permission_prompt_tool_name: str | None = None,
         sdk_message_mode: str | None = None,
+        sdk_pool: bool = False,
+        sdk_pool_size: int | None = None,
+        sdk_session_pool: bool = False,
+        system_prompt_injection_mode: str | None = None,
+        builtin_permissions: bool = True,
     ):
         """Initialize Claude options.
 
@@ -129,6 +166,11 @@ class ClaudeOptions:
             cwd: Working directory path
             permission_prompt_tool_name: Permission prompt tool name
             sdk_message_mode: SDK message handling mode
+            sdk_pool: Enable general Claude SDK client connection pooling
+            sdk_pool_size: Number of clients to maintain in the general pool
+            sdk_session_pool: Enable session-aware Claude SDK client pooling
+            system_prompt_injection_mode: System prompt injection mode
+            builtin_permissions: Enable built-in permission handling infrastructure
         """
         self.max_thinking_tokens = max_thinking_tokens
         self.allowed_tools = allowed_tools
@@ -140,3 +182,8 @@ class ClaudeOptions:
         self.cwd = cwd
         self.permission_prompt_tool_name = permission_prompt_tool_name
         self.sdk_message_mode = sdk_message_mode
+        self.sdk_pool = sdk_pool
+        self.sdk_pool_size = sdk_pool_size
+        self.sdk_session_pool = sdk_session_pool
+        self.system_prompt_injection_mode = system_prompt_injection_mode
+        self.builtin_permissions = builtin_permissions