diff --git a/src/cai/repl/commands/virtualization.py b/src/cai/repl/commands/virtualization.py index 6388d2ec..0146be71 100644 --- a/src/cai/repl/commands/virtualization.py +++ b/src/cai/repl/commands/virtualization.py @@ -147,17 +147,30 @@ class DockerManager: Tuple[bool, str]: Success status and output message """ try: - process = subprocess.run( + # Use Popen to stream pull progress line by line. + pull_proc = subprocess.Popen( ["docker", "pull", image_name], - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True, - check=False + bufsize=1, + ) + + # Stream output in real time to the console. + if pull_proc.stdout is not None: + for line in pull_proc.stdout: + console.print(line.rstrip()) + + pull_proc.wait() + + if pull_proc.returncode == 0: + return True, ( + f"Successfully pulled image: {image_name}" + ) + return False, ( + f"Failed to pull image: {image_name} " + f"(code {pull_proc.returncode})" ) - - if process.returncode == 0: - return True, f"Successfully pulled image: {image_name}" - else: - return False, f"Failed to pull image: {process.stderr}" except (subprocess.SubprocessError, FileNotFoundError) as e: return False, f"Error pulling image: {str(e)}" diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index d295451e..f687dded 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -11,6 +11,7 @@ import signal import time import uuid import sys +import shlex from wasabi import color # pylint: disable=import-error from cai.util import format_time @@ -774,18 +775,37 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501 - # Handle Streaming Container Execution - not yet implemented for containers + # Handle Streaming Container Execution if stream: - # For now, display that streaming isn't supported for containers - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "container": container_id[:12]} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for container execution. Running normally...", - call_id=call_id - ) + # Ensure workspace directory exists inside the container first + mkdir_cmd = [ + "docker", "exec", container_id, + "mkdir", "-p", container_workspace + ] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False, + timeout=10 + ) + + # Build docker exec command as a single shell string for streaming + docker_exec_cmd = ( + "docker exec -w " + f"{shlex.quote(container_workspace)} " + f"{shlex.quote(container_id)} sh -c " + f"{shlex.quote(command)}" + ) + + # Re-use the local streaming helper to provide real-time output + return _run_local_streamed( + docker_exec_cmd, + call_id, + timeout, + tool_name, + workspace_dir=_get_workspace_dir() + ) # Handle Synchronous Execution in Container try: diff --git a/tests/others/test_pretty_print.py b/tests/others/test_pretty_print.py index b3d6efef..020f6d61 100644 --- a/tests/others/test_pretty_print.py +++ b/tests/others/test_pretty_print.py @@ -47,7 +47,7 @@ async def test_pretty_run_result_streaming(): RunResultStreaming: - Current agent: Agent(name="test_agent", ...) - Current turn: 1 -- Max turns: 10 +- Max turns: inf - Is complete: True - Final output (str): Hi there @@ -111,7 +111,7 @@ async def test_pretty_run_result_streaming_structured_output(): RunResultStreaming: - Current agent: Agent(name="test_agent", ...) - Current turn: 1 -- Max turns: 10 +- Max turns: inf - Is complete: True - Final output (Foo): { @@ -189,7 +189,7 @@ async def test_pretty_run_result_streaming_list_structured_output(): RunResultStreaming: - Current agent: Agent(name="test_agent", ...) - Current turn: 1 -- Max turns: 10 +- Max turns: inf - Is complete: True - Final output (list): [Foo(bar='Test'), Foo(bar='Test 2')]