feat(cubbix): add --no-shell in combination with --run to not drop a shell and exit when the command is done

This commit is contained in:
2025-04-30 09:46:28 -06:00
parent e852fdd1fb
commit 75daccb366
4 changed files with 186 additions and 22 deletions

View File

@@ -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
```

View File

@@ -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,8 +307,46 @@ 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)
# 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:
# 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

View File

@@ -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")

View File

@@ -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