From 75daccb3662d059d178fd0f12026bb97f29f2452 Mon Sep 17 00:00:00 2001 From: Mathieu Virbel Date: Wed, 30 Apr 2025 09:46:28 -0600 Subject: [PATCH] feat(cubbix): add --no-shell in combination with --run to not drop a shell and exit when the command is done --- README.md | 6 ++ cubbi/cli.py | 85 +++++++++++++++++++----- cubbi/container.py | 108 +++++++++++++++++++++++++++++-- cubbi/images/goose/cubbi-init.sh | 9 +++ 4 files changed, 186 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index c11c31c..a1f62d1 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,9 @@ cubbix # Create a session and run an initial command before the shell starts cubbix --run "ls -l" +# Create a session, run a command, and exit (no shell prompt) +cubbix --run "ls -l" --no-shell + # List all active sessions cubbi session list @@ -111,6 +114,9 @@ cubbix https://github.com/username/repo --mcp github # Shorthand with an initial command cubbix . --run "apt-get update && apt-get install -y my-package" +# Execute a command and exit without starting a shell +cubbix . --run "python script.py" --no-shell + # Enable SSH server in the container cubbix --ssh ``` diff --git a/cubbi/cli.py b/cubbi/cli.py index 4b043ea..8c9d36b 100644 --- a/cubbi/cli.py +++ b/cubbi/cli.py @@ -147,6 +147,11 @@ def create_session( "--run", help="Command to execute inside the container before starting the shell", ), + no_shell: bool = typer.Option( + False, + "--no-shell", + help="Close container after '--run' command finishes (only valid with --run)", + ), no_connect: bool = typer.Option( False, "--no-connect", help="Don't automatically connect to the session" ), @@ -167,6 +172,9 @@ def create_session( None, "--provider", "-p", help="Provider to use" ), ssh: bool = typer.Option(False, "--ssh", help="Start SSH server in the container"), + verbose: bool = typer.Option( + False, "--verbose", "-v", help="Enable verbose logging" + ), ) -> None: """Create a new Cubbi session @@ -264,6 +272,12 @@ def create_session( if path_or_url and os.path.isdir(os.path.expanduser(path_or_url)): mount_local = True + # Check if --no-shell is used without --run + if no_shell and not run_command: + console.print( + "[yellow]Warning: --no-shell is ignored without --run[/yellow]" + ) + session = container_manager.create_session( image_name=image_name, project=path_or_url, @@ -275,6 +289,7 @@ def create_session( networks=all_networks, mcp=all_mcps, run_command=run_command, + no_shell=no_shell, uid=target_uid, gid=target_gid, ssh=ssh, @@ -292,24 +307,62 @@ def create_session( for container_port, host_port in session.ports.items(): console.print(f" {container_port} -> {host_port}") - # Auto-connect based on user config, unless overridden by --no-connect flag + # Auto-connect based on user config, unless overridden by --no-connect flag or --no-shell auto_connect = user_config.get("defaults.connect", True) - # Connect if auto_connect is enabled and --no-connect wasn't used. - # The --run command no longer prevents connection. - should_connect = not no_connect and auto_connect - if should_connect: - container_manager.connect_session(session.id) + + # When --no-shell is used with --run, show logs instead of connecting + if no_shell and run_command: + console.print( + "[yellow]Executing command and waiting for completion...[/yellow]" + ) + console.print("Container will exit after command completes.") + console.print("[bold]Command logs:[/bold]") + # Stream logs from the container until it exits + container_manager.get_session_logs(session.id, follow=True) + # At this point the command and container should have finished + + # Clean up the session entry to avoid leaving stale entries + with console.status("Cleaning up session..."): + # Give a short delay to ensure container has fully exited + import time + + time.sleep(1) + # Remove the session from session manager + session_manager.remove_session(session.id) + try: + # Also try to remove the container to ensure no resources are left behind + container = container_manager.client.containers.get( + session.container_id + ) + if container.status != "running": + container.remove(force=False) + except Exception as e: + # Container might already be gone or in the process of exiting + # This is fine, just log it + if verbose: + console.print(f"[yellow]Note: {e}[/yellow]") + + console.print( + "[green]Command execution complete. Container has exited.[/green]" + ) + console.print("[green]Session has been cleaned up.[/green]") else: - # Explain why connection was skipped - if no_connect: - console.print("\nConnection skipped due to --no-connect.") - console.print( - f"Connect manually with:\n cubbi session connect {session.id}" - ) - elif not auto_connect: - console.print( - f"\nAuto-connect disabled. Connect with:\n cubbi session connect {session.id}" - ) + # Connect if auto_connect is enabled and --no-connect wasn't used. + # The --run command no longer prevents connection. + should_connect = not no_connect and auto_connect + if should_connect: + container_manager.connect_session(session.id) + else: + # Explain why connection was skipped + if no_connect: + console.print("\nConnection skipped due to --no-connect.") + console.print( + f"Connect manually with:\n cubbi session connect {session.id}" + ) + elif not auto_connect: + console.print( + f"\nAuto-connect disabled. Connect with:\n cubbi session connect {session.id}" + ) else: console.print("[red]Failed to create session[/red]") diff --git a/cubbi/container.py b/cubbi/container.py index 2fe156b..ab6c42f 100644 --- a/cubbi/container.py +++ b/cubbi/container.py @@ -147,6 +147,7 @@ class ContainerManager: networks: Optional[List[str]] = None, mcp: Optional[List[str]] = None, run_command: Optional[str] = None, + no_shell: bool = False, uid: Optional[int] = None, gid: Optional[int] = None, model: Optional[str] = None, @@ -164,10 +165,13 @@ class ContainerManager: mount_local: Whether to mount the specified local directory to /app (ignored if project is None) volumes: Optional additional volumes to mount (dict of {host_path: {"bind": container_path, "mode": mode}}) run_command: Optional command to execute before starting the shell + no_shell: Whether to close the container after run_command completes (requires run_command) networks: Optional list of additional Docker networks to connect to mcp: Optional list of MCP server names to attach to the session uid: Optional user ID for the container process gid: Optional group ID for the container process + model: Optional model to use + provider: Optional provider to use ssh: Whether to start the SSH server in the container (default: False) """ try: @@ -470,7 +474,15 @@ class ContainerManager: if run_command: # Set environment variable for cubbi-init.sh to pick up env_vars["CUBBI_RUN_COMMAND"] = run_command - # Set the container's command to be the final shell + + # If no_shell is true, set CUBBI_NO_SHELL environment variable + if no_shell: + env_vars["CUBBI_NO_SHELL"] = "true" + logger.info( + "Setting CUBBI_NO_SHELL=true, container will exit after run command" + ) + + # Set the container's command to be the final shell (or exit if no_shell is true) container_command = [target_shell] logger.info( f"Setting CUBBI_RUN_COMMAND and targeting shell {target_shell}" @@ -818,8 +830,49 @@ class ContainerManager: if session.id == session_id and session.container_id: container = self.client.containers.get(session.container_id) if follow: - for line in container.logs(stream=True, follow=True): - print(line.decode().strip()) + # For streamed logs, we'll buffer by line to avoid character-by-character output + import io + from typing import Iterator + + def process_log_stream( + stream: Iterator[bytes], + ) -> Iterator[str]: + buffer = io.StringIO() + for chunk in stream: + chunk_str = chunk.decode("utf-8", errors="replace") + buffer.write(chunk_str) + + # Process complete lines + while True: + line = buffer.getvalue() + newline_pos = line.find("\n") + if newline_pos == -1: + break + + # Extract complete line and yield it + complete_line = line[:newline_pos].rstrip() + yield complete_line + + # Update buffer to contain only the remaining content + new_buffer = io.StringIO() + new_buffer.write(line[newline_pos + 1 :]) + buffer = new_buffer + + # Don't forget to yield any remaining content at the end + final_content = buffer.getvalue().strip() + if final_content: + yield final_content + + try: + # Process the log stream line by line + for line in process_log_stream( + container.logs(stream=True, follow=True) + ): + print(line) + except KeyboardInterrupt: + # Handle Ctrl+C gracefully + print("\nStopped following logs.") + return None else: return container.logs().decode() @@ -862,9 +915,52 @@ class ContainerManager: f"Following initialization logs for session {session_id}..." ) print("Press Ctrl+C to stop following") - container.exec_run( - "tail -f /init.log", stream=True, demux=True, tty=True - ) + + import io + + def process_exec_stream(stream): + buffer = io.StringIO() + for chunk_type, chunk_bytes in stream: + if chunk_type != 1: # Skip stderr (type 2) + continue + + chunk_str = chunk_bytes.decode( + "utf-8", errors="replace" + ) + buffer.write(chunk_str) + + # Process complete lines + while True: + line = buffer.getvalue() + newline_pos = line.find("\n") + if newline_pos == -1: + break + + # Extract complete line and yield it + complete_line = line[:newline_pos].rstrip() + yield complete_line + + # Update buffer to contain only the remaining content + new_buffer = io.StringIO() + new_buffer.write(line[newline_pos + 1 :]) + buffer = new_buffer + + # Don't forget to yield any remaining content at the end + final_content = buffer.getvalue().strip() + if final_content: + yield final_content + + try: + exec_result = container.exec_run( + "tail -f /init.log", stream=True, demux=True + ) + + # Process the exec stream line by line + for line in process_exec_stream(exec_result[1]): + print(line) + except KeyboardInterrupt: + print("\nStopped following logs.") + return None else: exit_code, output = container.exec_run("cat /init.log") diff --git a/cubbi/images/goose/cubbi-init.sh b/cubbi/images/goose/cubbi-init.sh index b6b6a10..800ebaf 100755 --- a/cubbi/images/goose/cubbi-init.sh +++ b/cubbi/images/goose/cubbi-init.sh @@ -170,6 +170,15 @@ if [ -n "$CUBBI_RUN_COMMAND" ]; then gosu cubbi sh -c "$CUBBI_RUN_COMMAND"; # Run user command as cubbi COMMAND_EXIT_CODE=$?; echo "--- Initial command finished (exit code: $COMMAND_EXIT_CODE) ---"; + + # If CUBBI_NO_SHELL is set, exit instead of starting a shell + if [ "$CUBBI_NO_SHELL" = "true" ]; then + echo "--- CUBBI_NO_SHELL=true, exiting container without starting shell ---"; + # Mark initialization as complete before exiting + echo "=== Cubbi Initialization completed at $(date) ===" + echo "INIT_COMPLETE=true" > /init.status + exit $COMMAND_EXIT_CODE; + fi; fi; # Mark initialization as complete