address coderabbit review comments

Signed-off-by: Jordan Dubrick <jdubrick@redhat.com>

Author: Jdubrick

src/app/endpoints/stream_interrupt.py

@@ -45,7 +45,19 @@ async def stream_interrupt_endpoint_handler(
         StreamInterruptRegistry, Depends(get_stream_interrupt_registry)
     ],
 ) -> StreamingInterruptResponse:
-    """Interrupt an in-progress streaming query by request identifier."""
+    """Interrupt an in-progress streaming query by request identifier.
+
+    Parameters:
+        interrupt_request: Request payload containing the stream request ID.
+        auth: Auth context tuple resolved from the authentication dependency.
+        registry: Stream interrupt registry dependency used to cancel streams.
+
+    Returns:
+        StreamingInterruptResponse: Confirmation payload when interruption succeeds.
+
+    Raises:
+        HTTPException: If no active stream for the given request ID can be interrupted.
+    """
     user_id, _, _, _ = auth
     request_id = interrupt_request.request_id
     interrupted = registry.cancel_stream(request_id, user_id)

src/app/endpoints/streaming_query.py

@@ -353,6 +353,11 @@ async def generate_response(
             user_id=user_id,
             task=current_task,
         )
+    else:
+        logger.warning(
+            "No current asyncio task for request %s; stream interruption will not be available",
+            request_id,
+        )
 
     stream_completed = False
     try:

src/models/requests.py

@@ -268,7 +268,11 @@ def validate_media_type(self) -> Self:
 
 
 class StreamingInterruptRequest(BaseModel):
-    """Model representing a request to interrupt an active streaming query."""
+    """Model representing a request to interrupt an active streaming query.
+
+    Attributes:
+        request_id: Unique ID of the active streaming request to interrupt.
+    """
 
     request_id: str = Field(
         description="The active streaming request ID to interrupt",
@@ -287,7 +291,17 @@ class StreamingInterruptRequest(BaseModel):
     @field_validator("request_id")
     @classmethod
     def check_request_id(cls, value: str) -> str:
-        """Validate that request identifier matches expected SUID format."""
+        """Validate that request identifier matches expected SUID format.
+
+        Parameters:
+            value: Request identifier submitted by the caller.
+
+        Returns:
+            str: The validated request identifier.
+
+        Raises:
+            ValueError: If the request identifier is not a valid SUID.
+        """
         if not suid.check_suid(value):
             raise ValueError(f"Improper request ID {value}")
         return value

src/utils/stream_interrupts.py

@@ -1,17 +1,22 @@
 """In-memory registry for interrupting active streaming requests."""
 
 import asyncio
-import logging
+from log import get_logger
 from dataclasses import dataclass
 from threading import Lock
-from typing import Optional
+from typing import Any
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 @dataclass
 class ActiveStream:
-    """Represents one active streaming request bound to a user."""
+    """Represents one active streaming request bound to a user.
+
+    Attributes:
+        user_id: Owner of the streaming request.
+        task: Asyncio task producing the stream response.
+    """
 
     user_id: str
     task: asyncio.Task
@@ -28,7 +33,13 @@ def __init__(self) -> None:
     def register_stream(
         self, request_id: str, user_id: str, task: asyncio.Task
     ) -> None:
-        """Register an active stream task for interrupt support."""
+        """Register an active stream task for interrupt support.
+
+        Parameters:
+            request_id: Unique streaming request identifier.
+            user_id: User identifier that owns the stream.
+            task: Asyncio task associated with the stream.
+        """
         with self._lock:
             self._streams[request_id] = ActiveStream(user_id=user_id, task=task)
 
@@ -39,6 +50,10 @@ def cancel_stream(self, request_id: str, user_id: str) -> bool:
         lock so that a concurrent ``deregister_stream`` cannot remove
         the entry between the ownership check and the cancel call.
 
+        Parameters:
+            request_id: Unique streaming request identifier.
+            user_id: User identifier attempting the interruption.
+
         Returns:
             bool: True when cancellation was requested, otherwise False.
         """
@@ -59,12 +74,23 @@ def cancel_stream(self, request_id: str, user_id: str) -> bool:
             return True
 
     def deregister_stream(self, request_id: str) -> None:
-        """Remove stream task from registry once completed/cancelled."""
+        """Remove stream task from registry once completed/cancelled.
+
+        Parameters:
+            request_id: Unique streaming request identifier.
+        """
         with self._lock:
             self._streams.pop(request_id, None)
 
-    def get_stream(self, request_id: str) -> Optional[ActiveStream]:
-        """Get currently registered stream metadata for tests/introspection."""
+    def get_stream(self, request_id: str) -> ActiveStream | None:
+        """Get currently registered stream metadata for tests/introspection.
+
+        Parameters:
+            request_id: Unique streaming request identifier.
+
+        Returns:
+            ActiveStream | None: Registered stream metadata, or None when absent.
+        """
         with self._lock:
             return self._streams.get(request_id)