streaming implementation working

Author: jagath-jaikumar

src/rb-api/rb/api/routes/cli.py

@@ -1,71 +1,92 @@
 import inspect
-from typing import Callable, Optional
+from typing import Callable, Generator, Optional
 
 import typer
 from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
 from makefun import with_signature
 from rb.api.models import CommandResult
 from rb.lib.stdout import Capturing  # type: ignore
+from rb.lib.stdout import capture_stdout_as_generator
 
 from rescuebox.main import app as rescuebox_app
 
 cli_router = APIRouter()
 
 
 def static_endpoint(callback: Callable, *args, **kwargs) -> CommandResult:
+    """Execute a CLI command and return the result synchronously"""
     with Capturing() as stdout:
         try:
             result = callback(*args, **kwargs)
-            success = True
-            error = None
+            return CommandResult(result=result, stdout=stdout, success=True, error=None)
         except Exception as e:
-            result = None
-            success = False
-            error = f"Typer CLI aborted {e}"
-    return CommandResult(result=result, stdout=stdout, success=success, error=error)
+            return CommandResult(
+                result=None,
+                stdout=stdout,
+                success=False,
+                error=f"Typer CLI aborted {e}",
+            )
+
+
+def streaming_endpoint(callback: Callable, *args, **kwargs) -> Generator:
+    """Execute a CLI command and stream the results"""
+    line_buffer = []
+    for line in capture_stdout_as_generator(callback, *args, **kwargs):
+        line_buffer.append(line)
+        result = CommandResult(
+            result=line, stdout=line_buffer, success=True, error=None
+        )
+        yield result.model_dump_json()
 
 
 def command_callback(command: typer.models.CommandInfo):
+    """Create a FastAPI endpoint handler for a Typer CLI command"""
     # Get the original callback signature
     original_signature = inspect.signature(command.callback)
 
-    # Modify the signature to include `streaming`
+    # Add streaming parameter to signature
     new_params = list(original_signature.parameters.values())
-    new_params.append(
-        inspect.Parameter(
-            "streaming",
-            inspect.Parameter.KEYWORD_ONLY,
-            default=False,
-            annotation=Optional[bool],
-        )
+    streaming_param = inspect.Parameter(
+        "streaming",
+        inspect.Parameter.KEYWORD_ONLY,
+        default=False,
+        annotation=Optional[bool],
     )
+    new_params.append(streaming_param)
     new_signature = original_signature.replace(parameters=new_params)
 
-    # Create a new function with the modified signature
     @with_signature(new_signature)
     def wrapper(*args, **kwargs):
-        # Extract additional parameters
-        # TODO(Jagath): Implement streaming
-        streaming = kwargs.pop("streaming", False)  # noqa: F841
+        streaming = kwargs.pop("streaming", False)
+
+        if streaming:
+            return StreamingResponse(
+                streaming_endpoint(command.callback, *args, **kwargs)
+            )
 
-        # Call the static endpoint with the wrapped callback and arguments
         result = static_endpoint(command.callback, *args, **kwargs)
         if not result.success:
-            # Return the last 10 lines of stdout if there's an error
             raise HTTPException(
                 status_code=400,
-                detail={"error": result.error, "stdout": result.stdout[-10:]},
+                detail={
+                    "error": result.error,
+                    "stdout": result.stdout[-10:],  # Last 10 lines of output
+                },
             )
         return result
 
+    # Preserve original function metadata
     wrapper.__name__ = command.callback.__name__
     wrapper.__doc__ = command.callback.__doc__
 
     return wrapper
 
 
+# Register routes for each plugin command
 for plugin in rescuebox_app.registered_groups:
     router = APIRouter()
+
     for command in plugin.typer_instance.registered_commands:
         router.add_api_route(
             f"/{command.callback.__name__}/",
@@ -74,4 +95,5 @@ def wrapper(*args, **kwargs):
             name=command.callback.__name__,
             response_model=CommandResult,
         )
+
     cli_router.include_router(router, prefix=f"/{plugin.name}", tags=[plugin.name])

src/rb-lib/rb/lib/stdout.py

@@ -1,5 +1,8 @@
+import contextlib
+import io
 import sys
 from io import StringIO
+from typing import Callable, Generator
 
 
 class Capturing(list):
@@ -12,3 +15,16 @@ def __exit__(self, *args):
         self.extend(self._stringio.getvalue().splitlines())
         del self._stringio  # free up some memory
         sys.stdout = self._stdout
+
+
+def capture_stdout_as_generator(
+    func: Callable, *args, **kwargs
+) -> Generator[str, None, None]:
+    stdout_buffer = io.StringIO()
+    with contextlib.redirect_stdout(stdout_buffer):
+        func(*args, **kwargs)
+
+    # Go to the start of the buffer and yield each line
+    stdout_buffer.seek(0)
+    for line in stdout_buffer:
+        yield line