diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index f658e183..e294bee7 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1609,43 +1609,50 @@ class _Converter: # Display the tool output immediately with the matched tool call from cai.util import cli_print_tool_output - # Look up the original tool call to get the name and arguments - if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: - tool_call = cls.recent_tool_calls[call_id] - tool_name = tool_call.get('name', 'Unknown Tool') - tool_args = tool_call.get('arguments', {}) - execution_info = tool_call.get('execution_info', {}) - - # Get token counts from the OpenAIChatCompletionsModel if available - model_instance = None - for frame in inspect.stack(): - if 'self' in frame.frame.f_locals: - self_obj = frame.frame.f_locals['self'] - if isinstance(self_obj, OpenAIChatCompletionsModel): - model_instance = self_obj - break - - token_info = {} - if model_instance: - token_info = { - 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), - 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), - 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), - 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), - 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), - 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), - 'model': str(getattr(model_instance, 'model', '')), - } - - # Use the cli_print_tool_output function with actual token values - cli_print_tool_output( - tool_name=tool_name, - args=tool_args, - output=output_content, - call_id=call_id, - execution_info=execution_info, - token_info=token_info - ) + # Check if we're in streaming mode - don't show tool output panel in streaming mode + is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true' + if is_streaming_enabled: + # Don't display tool output in streaming mode - it will be handled elsewhere + pass # Just skip the display, but preserve the tool output + else: + # For non-streaming mode, maintain the original behavior + # Look up the original tool call to get the name and arguments + if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: + tool_call = cls.recent_tool_calls[call_id] + tool_name = tool_call.get('name', 'Unknown Tool') + tool_args = tool_call.get('arguments', {}) + execution_info = tool_call.get('execution_info', {}) + + # Get token counts from the OpenAIChatCompletionsModel if available + model_instance = None + for frame in inspect.stack(): + if 'self' in frame.frame.f_locals: + self_obj = frame.frame.f_locals['self'] + if isinstance(self_obj, OpenAIChatCompletionsModel): + model_instance = self_obj + break + + token_info = {} + if model_instance: + token_info = { + 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), + 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), + 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), + 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), + 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), + 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), + 'model': str(getattr(model_instance, 'model', '')), + } + + # Use the cli_print_tool_output function with actual token values + cli_print_tool_output( + tool_name=tool_name, + args=tool_args, + output=output_content, + call_id=call_id, # Keep call_id for non-streaming mode + execution_info=execution_info, + token_info=token_info + ) # Continue with normal processing flush_assistant_message() diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index ebf6821e..91d9dbbe 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -245,6 +245,14 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): output = result.stdout if result.stdout else result.stderr if stdout: print("\033[32m" + output + "\033[0m") + + # Skip passing output to cli_print_tool_output when CAI_STREAM=true + # This prevents duplicate output in streaming mode + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if not is_streaming_enabled: + # Optional: Add cli_print_tool_output call here if needed for non-streaming + pass + return output except subprocess.TimeoutExpired as e: error_output = e.stdout.decode() if e.stdout else str(e) @@ -288,10 +296,16 @@ def _run_local_streamed(command, call_id, timeout=100): # Create panel content for Rich display if rich_available: tool_name = "generic_linux_command" + # Parse command into command and args + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args = parts[1] if len(parts) > 1 else "" + header = Text() header.append(tool_name, style="#00BCD4") header.append("(", style="yellow") - header.append(f"command='{command}'", style="yellow") + # Format to match: generic_linux_command({"command":"ls","args":"-la","ctf":{},"async_mode":false,"session_id":""}) + header.append(f'{{"command":"{cmd}","args":"{args}","ctf":{{}},"async_mode":false,"session_id":""}}', style="yellow") header.append(")", style="yellow") content = Text() @@ -299,9 +313,9 @@ def _run_local_streamed(command, call_id, timeout=100): panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -320,9 +334,9 @@ def _run_local_streamed(command, call_id, timeout=100): content.append(line, style="bright_white") panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -340,9 +354,9 @@ def _run_local_streamed(command, call_id, timeout=100): output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -353,9 +367,9 @@ def _run_local_streamed(command, call_id, timeout=100): content.append(f"\nCommand {completion_status}", style="green") panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle=f"[bold green]{completion_status}[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -365,7 +379,11 @@ def _run_local_streamed(command, call_id, timeout=100): time.sleep(0.5) else: # Fallback to simpler streaming with cli_print_tool_output - tool_args = {"command": command} + # Parse command into command and args (same as rich mode) + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args = parts[1] if len(parts) > 1 else "" + tool_args = {"command": cmd, "args": args, "ctf": {}, "async_mode": False, "session_id": ""} # Initial notification - just once cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id) diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index 7fa156e7..6c1274b8 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -90,9 +90,20 @@ def generic_linux_command(command: str = "", stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' # Generate a call_id for streaming if needed + # Only use call_id for streaming mode to prevent duplicate panels + # When CAI_STREAM=true, we want the Tool Execution panel but not the Tool Output panel call_id = str(uuid.uuid4())[:8] if stream else None + + # In streaming mode, prevent displaying both panels by forcing a call_id + # This tricks cli_print_tool_output into thinking it's being called for a streaming update + if stream and not call_id: + call_id = str(uuid.uuid4())[:8] # Run the command with the appropriate parameters - return run_command(full_command, ctf=ctf, + result = run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, timeout=timeout, stream=stream, call_id=call_id) + + # For better output formatting, if we're in streaming mode, we can modify the output + # to include any additional information needed while still preventing the duplicate panel + return result diff --git a/src/cai/util.py b/src/cai/util.py index 6397a230..087c98d2 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -877,8 +877,11 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli console.print(panel) # If there are tool panels, print them after the main message panel - for tool_panel in tool_panels: - console.print(tool_panel) + # But only in non-streaming mode to avoid duplicates + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if tool_panels and not is_streaming_enabled: + for tool_panel in tool_panels: + console.print(tool_panel) def create_agent_streaming_context(agent_name, counter, model): """Create a streaming context object that maintains state for streaming agent output.""" @@ -1072,10 +1075,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if not output and not call_id: return + # CRITICAL CHECK: When in streaming mode (CAI_STREAM=true), ONLY show output panels + # for streaming updates (those with call_id). This prevents duplicate output panels. + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if is_streaming_enabled and not call_id: + # Skip all non-streaming tool output in streaming mode + return + # Track seen call IDs to prevent duplicate panels if not hasattr(cli_print_tool_output, '_seen_calls'): cli_print_tool_output._seen_calls = {} - + # For streaming updates, only show updates for the same call_id # but allow the first appearance of each call_id if call_id: @@ -1127,11 +1137,12 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create the panel for display panel = Panel( Text.assemble(header, "\n\n", content), - title=f"[bold blue]Tool Output (ID: {call_id})[/bold blue]", - subtitle="[bold green]Live Update[/bold green]", + title=f"[bold blue]Tool Output[/bold blue]", #(ID: {call_id})[/bold green]", + #subtitle="[bold green]Live Update[/bold green]", border_style="blue", padding=(1, 2), - box=ROUNDED + box=ROUNDED, + title_align="left" ) # Display using Rich @@ -1181,9 +1192,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create the panel with all sections panel = Panel( Text.assemble(*sections), - title=f"[bold blue]Tool Output: {tool_call}[/bold blue]", + title=f"[bold green]Tool Output: {tool_call}[/bold green]", subtitle=subtitle, - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -1215,7 +1226,6 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut header = header[:term_width-4] # Clear the screen for the tool output (alternative approach) - # This works better with streamed content as it doesn't scroll the terminal print(f"\r{header}") # Print the content