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

ccproxy/services/claude_sdk_service.py

@@ -9,19 +9,20 @@ from claude_code_sdk import ClaudeCodeOptions
 from ccproxy.auth.manager import AuthManager
 from ccproxy.claude_sdk.client import ClaudeSDKClient
 from ccproxy.claude_sdk.converter import MessageConverter
+from ccproxy.claude_sdk.exceptions import StreamTimeoutError
+from ccproxy.claude_sdk.manager import SessionManager
 from ccproxy.claude_sdk.options import OptionsHandler
 from ccproxy.claude_sdk.streaming import ClaudeStreamProcessor
 from ccproxy.config.claude import SDKMessageMode
 from ccproxy.config.settings import Settings
 from ccproxy.core.errors import (
-    AuthenticationError,
     ClaudeProxyError,
     ServiceUnavailableError,
 )
 from ccproxy.models import claude_sdk as sdk_models
+from ccproxy.models.claude_sdk import SDKMessage, create_sdk_message
 from ccproxy.models.messages import MessageResponse
-from ccproxy.observability.access_logger import log_request_access
-from ccproxy.observability.context import RequestContext, request_context
+from ccproxy.observability.context import RequestContext
 from ccproxy.observability.metrics import PrometheusMetrics
 from ccproxy.utils.model_mapping import map_model_to_claude
 from ccproxy.utils.simple_request_logger import write_request_log
@@ -45,6 +46,7 @@ class ClaudeSDKService:
         auth_manager: AuthManager | None = None,
         metrics: PrometheusMetrics | None = None,
         settings: Settings | None = None,
+        session_manager: SessionManager | None = None,
     ) -> None:
         """
         Initialize Claude SDK service.
@@ -54,8 +56,11 @@ class ClaudeSDKService:
             auth_manager: Authentication manager (optional)
             metrics: Prometheus metrics instance (optional)
             settings: Application settings (optional)
+            session_manager: Session manager for dependency injection (optional)
         """
-        self.sdk_client = sdk_client or ClaudeSDKClient()
+        self.sdk_client = sdk_client or ClaudeSDKClient(
+            settings=settings, session_manager=session_manager
+        )
         self.auth_manager = auth_manager
         self.metrics = metrics
         self.settings = settings
@@ -66,14 +71,120 @@ class ClaudeSDKService:
             metrics=self.metrics,
         )
 
+    def _convert_messages_to_sdk_message(
+        self, messages: list[dict[str, Any]], session_id: str | None = None
+    ) -> "SDKMessage":
+        """Convert list of Anthropic messages to single SDKMessage.
+
+        Takes the last user message from the list and converts it to SDKMessage format.
+
+        Args:
+            messages: List of Anthropic API messages
+            session_id: Optional session ID for conversation continuity
+
+        Returns:
+            SDKMessage ready to send to Claude SDK
+        """
+        # Find the last user message
+        last_user_message = None
+        for msg in reversed(messages):
+            if msg.get("role") == "user":
+                last_user_message = msg
+                break
+
+        if not last_user_message:
+            raise ClaudeProxyError(
+                message="No user message found in messages list",
+                error_type="invalid_request_error",
+                status_code=400,
+            )
+
+        # Extract text content from the message
+        content = last_user_message.get("content", "")
+        if isinstance(content, list):
+            # Extract text from content blocks
+            text_parts = []
+            for block in content:
+                if isinstance(block, dict) and block.get("type") == "text":
+                    text_parts.append(block.get("text", ""))
+            content = "\n".join(text_parts)
+        elif not isinstance(content, str):
+            content = str(content)
+
+        return create_sdk_message(content=content, session_id=session_id)
+
+    async def _capture_session_metadata(
+        self,
+        ctx: RequestContext,
+        session_id: str | None,
+        options: "ClaudeCodeOptions",
+    ) -> None:
+        """Capture session metadata for access logging.
+
+        Args:
+            ctx: Request context to add metadata to
+            session_id: Optional session ID
+            options: Claude Code options
+        """
+        if (
+            session_id
+            and hasattr(self.sdk_client, "_session_manager")
+            and self.sdk_client._session_manager
+        ):
+            try:
+                session_client = (
+                    await self.sdk_client._session_manager.get_session_client(
+                        session_id, options
+                    )
+                )
+                if session_client:
+                    # Determine if session pool is enabled
+                    session_pool_enabled = (
+                        hasattr(self.sdk_client._session_manager, "session_pool")
+                        and self.sdk_client._session_manager.session_pool is not None
+                        and hasattr(
+                            self.sdk_client._session_manager.session_pool, "config"
+                        )
+                        and self.sdk_client._session_manager.session_pool.config.enabled
+                    )
+
+                    # Add session metadata to context
+                    ctx.add_metadata(
+                        session_type="session_pool"
+                        if session_pool_enabled
+                        else "direct",
+                        session_status=session_client.status.value,
+                        session_age_seconds=session_client.metrics.age_seconds,
+                        session_message_count=session_client.metrics.message_count,
+                        session_client_id=session_client.client_id,
+                        session_pool_enabled=session_pool_enabled,
+                        session_idle_seconds=session_client.metrics.idle_seconds,
+                        session_error_count=session_client.metrics.error_count,
+                        session_is_new=session_client.is_newly_created,
+                    )
+            except Exception as e:
+                logger.warning(
+                    "failed_to_capture_session_metadata",
+                    session_id=session_id,
+                    error=str(e),
+                )
+        else:
+            # Add basic session metadata for direct connections (no session pool)
+            ctx.add_metadata(
+                session_type="direct",
+                session_pool_enabled=False,
+                session_is_new=True,  # Direct connections are always new
+            )
+
     async def create_completion(
         self,
+        request_context: RequestContext,
         messages: list[dict[str, Any]],
         model: str,
         temperature: float | None = None,
         max_tokens: int | None = None,
         stream: bool = False,
-        user_id: str | None = None,
+        session_id: str | None = None,
         **kwargs: Any,
     ) -> MessageResponse | AsyncIterator[dict[str, Any]]:
         """
@@ -85,7 +196,8 @@ class ClaudeSDKService:
             temperature: Temperature for response generation
             max_tokens: Maximum tokens in response
             stream: Whether to stream responses
-            user_id: User identifier for auth/metrics
+            session_id: Optional session ID for Claude SDK integration
+            request_context: Existing request context to use instead of creating new one
             **kwargs: Additional arguments
 
         Returns:
@@ -96,20 +208,6 @@ class ClaudeSDKService:
             ServiceUnavailableError: If service is unavailable
         """
 
-        # Validate authentication if auth manager is configured
-        if self.auth_manager and user_id:
-            try:
-                await self._validate_user_auth(user_id)
-            except Exception as e:
-                logger.error(
-                    "authentication_failed",
-                    user_id=user_id,
-                    error=str(e),
-                    error_type=type(e).__name__,
-                    exc_info=True,
-                )
-                raise
-
         # Extract system message and create options
         system_message = self.options_handler.extract_system_message(messages)
 
@@ -121,74 +219,55 @@ class ClaudeSDKService:
             temperature=temperature,
             max_tokens=max_tokens,
             system_message=system_message,
+            session_id=session_id,
             **kwargs,
         )
 
-        # Convert messages to prompt format
-        prompt = self.message_converter.format_messages_to_prompt(messages)
-
-        # Generate request ID for correlation
-        from uuid import uuid4
-
-        request_id = str(uuid4())
+        # Messages will be converted to SDK format in the client layer
 
-        # Use request context for observability
-        endpoint = "messages"  # Claude SDK uses messages endpoint
-        async with request_context(
-            method="POST",
-            path=f"/sdk/v1/{endpoint}",
-            endpoint=endpoint,
-            model=model,
-            streaming=stream,
-            service_type="claude_sdk_service",
-            metrics=self.metrics,  # Pass metrics for active request tracking
-        ) as ctx:
-            try:
-                # Log SDK request parameters
-                timestamp = ctx.get_log_timestamp_prefix() if ctx else None
-                await self._log_sdk_request(
-                    request_id, prompt, options, model, stream, timestamp
-                )
+        # Use existing context, but update metadata for this service (preserve original service_type)
+        ctx = request_context
+        metadata = {
+            "endpoint": "messages",
+            "model": model,
+            "streaming": stream,
+        }
+        if session_id:
+            metadata["session_id"] = session_id
+        ctx.add_metadata(**metadata)
+        # Use existing request ID from context
+        request_id = ctx.request_id
 
-                if stream:
-                    # For streaming, return the async iterator directly
-                    # Pass context to streaming method
-                    return self._stream_completion(
-                        prompt, options, model, request_id, ctx, timestamp
-                    )
-                else:
-                    result = await self._complete_non_streaming(
-                        prompt, options, model, request_id, ctx, timestamp
-                    )
-                    return result
+        try:
+            # Log SDK request parameters
+            timestamp = ctx.get_log_timestamp_prefix() if ctx else None
+            await self._log_sdk_request(
+                request_id, messages, options, model, stream, session_id, timestamp
+            )
 
-            except AuthenticationError as e:
-                logger.error(
-                    "authentication_failed",
-                    user_id=user_id,
-                    error=str(e),
-                    error_type=type(e).__name__,
-                    exc_info=True,
+            if stream:
+                # For streaming, return the async iterator directly
+                # Access logging will be handled by the stream processor when ResultMessage is received
+                return self._stream_completion(
+                    ctx, messages, options, model, session_id, timestamp
                 )
-                raise
-            except (ClaudeProxyError, ServiceUnavailableError) as e:
-                # Log error via access logger (includes metrics)
-                await log_request_access(
-                    context=ctx,
-                    method="POST",
-                    error_message=str(e),
-                    metrics=self.metrics,
-                    error_type=type(e).__name__,
+            else:
+                result = await self._complete_non_streaming(
+                    ctx, messages, options, model, session_id, timestamp
                 )
-                raise
+                return result
+        except (ClaudeProxyError, ServiceUnavailableError) as e:
+            # Add error info to context for automatic access logging
+            ctx.add_metadata(error_message=str(e), error_type=type(e).__name__)
+            raise
 
     async def _complete_non_streaming(
         self,
-        prompt: str,
+        ctx: RequestContext,
+        messages: list[dict[str, Any]],
         options: "ClaudeCodeOptions",
         model: str,
-        request_id: str | None = None,
-        ctx: RequestContext | None = None,
+        session_id: str | None = None,
         timestamp: str | None = None,
     ) -> MessageResponse:
         """
@@ -198,7 +277,6 @@ class ClaudeSDKService:
             prompt: The formatted prompt
             options: Claude SDK options
             model: The model being used
-            request_id: The request ID for metrics correlation
 
         Returns:
             Response in Anthropic format
@@ -206,18 +284,31 @@ class ClaudeSDKService:
         Raises:
             ClaudeProxyError: If completion fails
         """
-        # SDK request already logged in create_completion
+        request_id = ctx.request_id
+        logger.debug("claude_sdk_completion_start", request_id=request_id)
 
-        messages = [
-            m
-            async for m in self.sdk_client.query_completion(prompt, options, request_id)
-        ]
+        # Convert messages to single SDKMessage
+        sdk_message = self._convert_messages_to_sdk_message(messages, session_id)
+
+        # Get stream handle
+        stream_handle = await self.sdk_client.query_completion(
+            sdk_message, options, request_id, session_id
+        )
+
+        # Capture session metadata for access logging
+        await self._capture_session_metadata(ctx, session_id, options)
+
+        # Create a listener and collect all messages
+        sdk_messages = []
+        async for m in stream_handle.create_listener():
+            sdk_messages.append(m)
 
         result_message = next(
-            (m for m in messages if isinstance(m, sdk_models.ResultMessage)), None
+            (m for m in sdk_messages if isinstance(m, sdk_models.ResultMessage)), None
         )
         assistant_message = next(
-            (m for m in messages if isinstance(m, sdk_models.AssistantMessage)), None
+            (m for m in sdk_messages if isinstance(m, sdk_models.AssistantMessage)),
+            None,
         )
 
         if result_message is None:
@@ -249,7 +340,7 @@ class ClaudeSDKService:
         # Add other message types to the content block
         all_messages = [
             m
-            for m in messages
+            for m in sdk_messages
             if not isinstance(m, sdk_models.AssistantMessage | sdk_models.ResultMessage)
         ]
 
@@ -306,18 +397,18 @@ class ClaudeSDKService:
             request_id=request_id,
         )
 
-        if ctx:
-            ctx.add_metadata(
-                status_code=200,
-                tokens_input=usage.input_tokens,
-                tokens_output=usage.output_tokens,
-                cache_read_tokens=usage.cache_read_input_tokens,
-                cache_write_tokens=usage.cache_creation_input_tokens,
-                cost_usd=cost_usd,
-            )
-            await log_request_access(
-                context=ctx, status_code=200, method="POST", metrics=self.metrics
-            )
+        ctx.add_metadata(
+            status_code=200,
+            tokens_input=usage.input_tokens,
+            tokens_output=usage.output_tokens,
+            cache_read_tokens=usage.cache_read_input_tokens,
+            cache_write_tokens=usage.cache_creation_input_tokens,
+            cost_usd=cost_usd,
+            session_id=result_message.session_id,
+            num_turns=result_message.num_turns,
+        )
+        # Add success status to context for automatic access logging
+        ctx.add_metadata(status_code=200)
 
         # Log SDK response
         if request_id:
@@ -327,11 +418,11 @@ class ClaudeSDKService:
 
     async def _stream_completion(
         self,
-        prompt: str,
+        ctx: RequestContext,
+        messages: list[dict[str, Any]],
         options: "ClaudeCodeOptions",
         model: str,
-        request_id: str | None = None,
-        ctx: RequestContext | None = None,
+        session_id: str | None = None,
         timestamp: str | None = None,
     ) -> AsyncIterator[dict[str, Any]]:
         """
@@ -341,12 +432,12 @@ class ClaudeSDKService:
             prompt: The formatted prompt
             options: Claude SDK options
             model: The model being used
-            request_id: Optional request ID for logging
             ctx: Optional request context for metrics
 
         Yields:
             Response chunks in Anthropic format
         """
+        request_id = ctx.request_id
         sdk_message_mode = (
             self.settings.claude.sdk_message_mode
             if self.settings
@@ -354,66 +445,167 @@ class ClaudeSDKService:
         )
         pretty_format = self.settings.claude.pretty_format if self.settings else True
 
-        sdk_stream = self.sdk_client.query_completion(prompt, options, request_id)
+        # Convert messages to single SDKMessage
+        sdk_message = self._convert_messages_to_sdk_message(messages, session_id)
 
-        async for chunk in self.stream_processor.process_stream(
-            sdk_stream=sdk_stream,
-            model=model,
-            request_id=request_id,
-            ctx=ctx,
-            sdk_message_mode=sdk_message_mode,
-            pretty_format=pretty_format,
+        # Get stream handle instead of direct iterator
+        stream_handle = await self.sdk_client.query_completion(
+            sdk_message, options, request_id, session_id
+        )
+
+        # Store handle in session client if available for cleanup
+        if (
+            session_id
+            and hasattr(self.sdk_client, "_session_manager")
+            and self.sdk_client._session_manager
         ):
-            # Log streaming chunk
-            if request_id:
-                await self._log_sdk_streaming_chunk(request_id, chunk, timestamp)
-            yield chunk
+            try:
+                session_client = (
+                    await self.sdk_client._session_manager.get_session_client(
+                        session_id, options
+                    )
+                )
+                if session_client:
+                    session_client.active_stream_handle = stream_handle
+            except Exception as e:
+                logger.warning(
+                    "failed_to_store_stream_handle",
+                    session_id=session_id,
+                    error=str(e),
+                )
 
-    async def _validate_user_auth(self, user_id: str) -> None:
-        """
-        Validate user authentication.
+        # Capture session metadata for access logging
+        await self._capture_session_metadata(ctx, session_id, options)
 
-        Args:
-            user_id: User identifier
+        # Create a listener for this stream
+        sdk_stream = stream_handle.create_listener()
 
-        Raises:
-            AuthenticationError: If authentication fails
-        """
-        if not self.auth_manager:
-            return
-        logger.debug("user_auth_validation_start", user_id=user_id)
+        try:
+            async for chunk in self.stream_processor.process_stream(
+                sdk_stream=sdk_stream,
+                model=model,
+                request_id=request_id,
+                ctx=ctx,
+                sdk_message_mode=sdk_message_mode,
+                pretty_format=pretty_format,
+            ):
+                # Log streaming chunk
+                if request_id:
+                    await self._log_sdk_streaming_chunk(request_id, chunk, timestamp)
+                yield chunk
+        except GeneratorExit:
+            # Client disconnected - log and re-raise to propagate to create_listener()
+            logger.info(
+                "claude_sdk_service_client_disconnected",
+                request_id=request_id,
+                session_id=session_id,
+                message="Client disconnected from SDK service stream, propagating to stream handle",
+            )
+            # CRITICAL: Re-raise GeneratorExit to trigger interrupt in create_listener()
+            raise
+        except StreamTimeoutError as e:
+            # Send error events to the client
+            logger.error(
+                "stream_timeout_error",
+                message=str(e),
+                session_id=e.session_id,
+                timeout_seconds=e.timeout_seconds,
+                request_id=request_id,
+            )
+
+            # Create a unique message ID for the error response
+            from uuid import uuid4
+
+            error_message_id = f"msg_error_{uuid4()}"
+
+            # Yield message_start event
+            yield {
+                "type": "message_start",
+                "message": {
+                    "id": error_message_id,
+                    "type": "message",
+                    "role": "assistant",
+                    "model": model,
+                    "content": [],
+                    "stop_reason": "error",
+                    "stop_sequence": None,
+                    "usage": {"input_tokens": 0, "output_tokens": 0},
+                },
+            }
+
+            # Yield content_block_start for error message
+            yield {
+                "type": "content_block_start",
+                "index": 0,
+                "content_block": {"type": "text", "text": ""},
+            }
+
+            # Yield error text delta
+            error_text = f"Error: {e}"
+            yield {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": error_text},
+            }
+
+            # Yield content_block_stop
+            yield {
+                "type": "content_block_stop",
+                "index": 0,
+            }
+
+            # Yield message_delta with stop reason
+            yield {
+                "type": "message_delta",
+                "delta": {"stop_reason": "error", "stop_sequence": None},
+                "usage": {"output_tokens": len(error_text.split())},
+            }
+
+            # Yield message_stop
+            yield {
+                "type": "message_stop",
+            }
+
+            # Update context with error status
+            ctx.add_metadata(
+                status_code=504,  # Gateway Timeout
+                error_message=str(e),
+                error_type="stream_timeout",
+                session_id=e.session_id,
+            )
 
     async def _log_sdk_request(
         self,
         request_id: str,
-        prompt: str,
+        messages: list[dict[str, Any]],
         options: "ClaudeCodeOptions",
         model: str,
         stream: bool,
+        session_id: str | None = None,
         timestamp: str | None = None,
     ) -> None:
         """Log SDK input parameters as JSON dump.
 
         Args:
             request_id: Request identifier
-            prompt: The formatted prompt
+            messages: List of Anthropic API messages
             options: Claude SDK options
             model: The model being used
             stream: Whether streaming is enabled
+            session_id: Optional session ID for Claude SDK integration
             timestamp: Optional timestamp prefix
         """
         # timestamp is already provided from context, no need for fallback
 
         # JSON dump of the parameters passed to SDK completion
         sdk_request_data = {
-            "prompt": prompt,
-            "options": options.model_dump()
-            if hasattr(options, "model_dump")
-            else str(options),
-            "model": model,
+            "messages": messages,
+            "options": options,
             "stream": stream,
             "request_id": request_id,
         }
+        if session_id:
+            sdk_request_data["session_id"] = session_id
 
         await write_request_log(
             request_id=request_id,
@@ -497,6 +689,17 @@ class ClaudeSDKService:
             )
             return False
 
+    async def interrupt_session(self, session_id: str) -> bool:
+        """Interrupt a Claude session due to client disconnection.
+
+        Args:
+            session_id: The session ID to interrupt
+
+        Returns:
+            True if session was found and interrupted, False otherwise
+        """
+        return await self.sdk_client.interrupt_session(session_id)
+
     async def close(self) -> None:
         """Close the service and cleanup resources."""
         await self.sdk_client.close()