Merge branch 'port_ui' into 'v0.4.0'

Make UI same as in v3

See merge request aliasrobotics/alias_research/cai!147

src/cai/cli.py

@@ -100,6 +100,7 @@ Usage Examples:
 
 import os
 import sys
+import time
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner, AsyncOpenAI
@@ -114,7 +115,7 @@ from cai.util import create_agent_streaming_context, update_agent_streaming_cont
 from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command
 from cai.repl.ui.keybindings import create_key_bindings
 from cai.repl.ui.logging import setup_session_logging
-from cai.repl.ui.banner import display_banner
+from cai.repl.ui.banner import display_banner, display_quick_guide
 from cai.repl.ui.prompt import get_user_input
 from cai.repl.ui.toolbar import get_toolbar_with_refresh
 
@@ -132,6 +133,10 @@ load_dotenv()
 #
 # set_default_openai_client(external_client)
 
+# Global variables for timing tracking
+global START_TIME
+START_TIME = time.time() 
+
 set_tracing_disabled(True)
 
 # llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini')
@@ -167,7 +172,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
     """
     agent = starting_agent
     turn_count = 0
-
+    ACTIVE_TIME = 0
+    idle_time = 0
     console = Console()
 
     # Initialize command completer and key bindings
@@ -180,7 +186,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
 
     # Display banner
     display_banner(console)
-
+    display_quick_guide(console)
     # Function to get the short name of the agent for display
     def get_agent_short_name(agent):
         if hasattr(agent, 'name'):
@@ -201,6 +207,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
 
     while turn_count < max_turns:
         try:
+            idle_start_time = time.time()
             # Get user input with command completion and history
             user_input = get_user_input(
                 command_completer,
@@ -209,8 +216,59 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
                 get_toolbar_with_refresh,
                 current_text
             )
+            idle_time += time.time() - idle_start_time
+           
         except KeyboardInterrupt:
+            def format_time(seconds):
+                mins, secs = divmod(int(seconds), 60)
+                hours, mins = divmod(mins, 60)
+                return f"{hours:02d}:{mins:02d}:{secs:02d}"
+
+            Total = time.time() - START_TIME
+            idle_time += time.time() - idle_start_time
+            try:
+                active_time = Total - idle_time
+
+                metrics = {
+                    "session_time": format_time(Total),
+                    "active_time": format_time(active_time),
+                    "idle_time": format_time(idle_time),
+                    "llm_time": "0.0s",  # Placeholder, update if available
+                    "llm_percentage": 0.0,  # Placeholder, update if available
+                }
+                logging_path = None  # Set this if you have a log file path
+
+                content = []
+                content.append(f"Session Time: {metrics['session_time']}")
+                content.append(f"Active Time: {metrics['active_time']}")
+                content.append(f"Idle Time: {metrics['idle_time']}")
+                if logging_path:
+                    content.append(f"Log available at: {logging_path}")
+                
+                def print_session_summary(console, metrics, logging_path=None):
+                    """
+                    Print a session summary panel using Rich.
+                    """
+                    from rich.panel import Panel
+                    from rich.text import Text
+                    from rich.box import ROUNDED
+                    from rich.console import Group
+
+                    time_panel = Panel(
+                        Group(*[Text(line) for line in content]),
+                        border_style="blue",
+                        box=ROUNDED,
+                        padding=(0, 1),
+                        title="[bold]Session Summary[/bold]",
+                        title_align="left"
+                    )
+                    console.print(time_panel)
+
+                print_session_summary(console, metrics, logging_path)
+            except Exception:
+                pass
             break
+      
         try:
             # Handle special commands
             if user_input.startswith('/') or user_input.startswith('$'):
@@ -355,10 +413,10 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
                 asyncio.run(process_streamed_response())
             else:
                 # Use non-streamed response
-                console.print("[dim]Thinking...[/dim]")
                 response = asyncio.run(Runner.run(agent, user_input))
                 #console.print(f"Agent: {response.final_output}") # NOTE: this line is commented to avoid duplicate output
             turn_count += 1
+
         except KeyboardInterrupt:
             if stream:
                 # Ensure streaming context is cleaned up on keyboard interrupt

src/cai/repl/ui/banner.py

@@ -253,3 +253,107 @@ def display_welcome_tips(console: Console):
         title="Quick Tips",
         border_style="blue"
     ))
+
+
+def display_quick_guide(console: Console):
+    """Display the quick guide."""
+    # Display help panel instead
+    from rich.panel import Panel
+    from rich.text import Text
+    from rich.columns import Columns
+    from rich.console import Group  # <-- Fix: import Group
+
+    help_text = Text.assemble(
+        ("CAI Command Reference", "bold cyan underline"), "\n\n",
+        ("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", "dim"), "\n",
+        ("WORKSPACE", "bold yellow"), "\n",
+        ("  CAI>/ws set [NAME]", "green"), " - Set current workspace directory\n\n",
+        ("AGENT MANAGEMENT", "bold yellow"), "\n",
+        ("  CAI>/agent [NAME]", "green"), " - Switch to specific agent by name\n",
+        ("  CAI>/agent 1 2 3", "green"), " - Switch to agent by position number\n",
+        ("  CAI>/agent", "green"), " - Display list of all available agents\n\n",
+        ("MODEL SELECTION", "bold yellow"), "\n",
+        ("  CAI>/model [NAME]", "green"), " - Change to a different model by name\n",
+        ("  CAI>/model 1", "green"), " - Change model by position number\n",
+        ("  CAI>/model", "green"), " - Show all available models\n\n",
+        ("INPUT & EXECUTION", "bold yellow"), "\n",
+        ("  ESC + ENTER", "green"), " - Enter multi-line input mode\n",
+        ("  CAI>/shell or CAI> $", "green"), " - Run system shell commands\n",
+        ("  CAI>hi, cybersecurity AI", "green"), " - Any text without commands will be sent as a prompt\n",
+        ("  CAI>/help", "green"), " - Display complete command reference\n",
+        ("  CAI>/flush or CAI> /clear", "green"), " - Clear the conversation history\n\n",
+        ("UTILITY COMMANDS", "bold yellow"), "\n",
+        ("  CAI>/mcp", "green"), " - Load additional tools with MCP server to an agent\n",
+        ("  CAI>/virt", "green"), " - Show all available virtualized environments\n",
+        ("  CAI>/flush", "green"), " - Flush context/message list\n",
+        ("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", "dim"), "\n",
+    )
+    
+    # Get current environment variable values
+    current_model = os.getenv('CAI_MODEL', "qwen2.5:14b")
+    current_agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent")
+    
+    config_text = Text.assemble(
+        ("Quick Start Configuration", "bold cyan underline"), "\n\n",
+        ("1. Configure .env file with your settings", "yellow"), "\n",
+        ("2. Select an agent: ", "yellow"), f"by default: CAI_AGENT_TYPE={current_agent_type}\n",
+        ("3. Select a model: ", "yellow"), f"by default: CAI_MODEL={current_model}\n\n",
+        ("Basic Usage:", "bold yellow"), "\n",
+        ("  1. CAI> /model", "green"), " - View all available models first\n",
+        ("  2. CAI> /agent", "green"), " - View all available agents first\n",
+        ("  3. CAI> /model deepseek/deepseek-chat", "green"), " - Then select your preferred model\n",
+        ("  4. CAI> /agent 16", "green"), " - Then select your preferred agent\n",
+        ("  5. CAI> Scan 192.168.1.1", "green"), " - Example prompt for target scan\n\n",
+        ("  /help", "green"), " - Display complete command reference\n\n",
+        ("Common Environment Variables:", "bold yellow"), "\n",
+        ("  CAI_MODEL", "green"), f" - Model to use (default: {current_model})\n",
+        ("  CAI_AGENT_TYPE", "green"), f" - Agent type (default: {current_agent_type})\n",
+        ("  CAI_DEBUG", "green"), f" - Debug level (default: {os.getenv('CAI_DEBUG', '1')})\n",
+        ("  CAI_MAX_TURNS", "green"), f" - Max conversation turns (default: {os.getenv('CAI_MAX_TURNS', 'inf')})\n",
+        ("  CAI_TRACING", "green"), f" - Enable tracing (default: {os.getenv('CAI_TRACING', 'true')})\n",
+    )
+    
+    # Create additional tips panels
+    ollama_tip = Panel(
+        "To use Ollama models, configure OLLAMA_API_BASE\n"
+        "before startup.\n\n"
+        "Default: host.docker.internal:8000/v1",
+        title="[bold yellow]Ollama Configuration[/bold yellow]",
+        border_style="yellow",
+        padding=(1, 2),
+        title_align="center"
+    )
+    
+    context_tip = Panel(
+        "As security exercises progress, LLM quality may\n"
+        "degrade, especially if progress stalls.\n\n"
+        "It's often better to clear the context window\n"
+        "or restart CAI rather than waiting until\n"
+        "context usage reaches 100%.\n\n"
+        "When context exceeds 80%, follow these steps:\n"
+        "1. CAI> Dump your memory and findings in current scenario in findings.txt\n"
+        "2. CAI> /flush\n"
+        "3. CAI> Analyze findings.txt, and continue exercise with target: ...",
+        title="[bold yellow]Performance Tip[/bold yellow]",
+        border_style="yellow",
+        padding=(1, 2),
+        title_align="center"
+    )
+    
+    # Combine tips into a group
+    # tips_group = Group(ollama_tip, context_tip)
+    tips_group = Group(context_tip)
+    
+    # Create a three-column panel layout
+    console.print(Panel(
+        Columns(
+            [help_text, config_text, tips_group],
+            column_first=True,
+            expand=True,
+            align="center"
+        ),
+        title="[bold]CAI Quick Guide[/bold]",
+        border_style="blue",
+        padding=(1, 2),
+        title_align="center"
+    ))

src/cai/sdk/agents/models/openai_chatcompletions.py

@@ -1632,17 +1632,31 @@ class _Converter:
                                     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', '')),
-                            }
+                        # Always create a token_info dictionary, even if some values are zero
+                        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', '')),
+                        }
+                        
+                        # Calculate costs using standard cost model
+                        if model_instance and hasattr(model_instance, 'model'):
+                            from cai.util import calculate_model_cost
+                            model_name = str(model_instance.model)
+                            token_info['interaction_cost'] = calculate_model_cost(
+                                model_name, 
+                                token_info['interaction_input_tokens'], 
+                                token_info['interaction_output_tokens']
+                            )
+                            token_info['total_cost'] = calculate_model_cost(
+                                model_name,
+                                token_info['total_input_tokens'],
+                                token_info['total_output_tokens']
+                            )
                         
                         # Use the cli_print_tool_output function with actual token values
                         cli_print_tool_output(

src/cai/tools/common.py

@@ -12,6 +12,14 @@ import time
 import uuid
 import sys
 from wasabi import color  # pylint: disable=import-error
+from cai.util  import format_time
+
+
+# Instead of direct import
+try:
+    from cai.cli import START_TIME
+except ImportError:
+    START_TIME = None
 
 # Global dictionary to store active sessions
 ACTIVE_SESSIONS = {}
@@ -228,10 +236,10 @@ def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None
         return f"Error executing CTF command: {str(e)}"
 
 
-def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None):
+def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None):
     # If streaming is enabled and we have a call_id
     if stream and call_id:
-        return _run_local_streamed(command, call_id, timeout)
+        return _run_local_streamed(command, call_id, timeout, tool_name)
     
     try:
         # nosec B602 - shell=True is required for command chaining
@@ -265,7 +273,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None):
         return error_msg
 
 
-def _run_local_streamed(command, call_id, timeout=100):
+def _run_local_streamed(command, call_id, timeout=100, tool_name=None):
     """Run a local command with streaming output to the Tool output panel"""
     try:
         # Try to import Rich for nice display
@@ -293,23 +301,43 @@ def _run_local_streamed(command, call_id, timeout=100):
             bufsize=1
         )
         
+        # If tool_name is not provided, derive it from the command
+        if tool_name is None:
+            # Just use the first command as the tool name
+            tool_name = command.strip().split()[0] + "_command"
+        
         # 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 ""
             
+            # Format clean arguments, following the same rules as cli_print_tool_output
+            arg_parts = []
+            if cmd:
+                arg_parts.append(f"command={cmd}")
+            if args and args.strip():  # Only add args if non-empty
+                arg_parts.append(f"args={args}")
+            args_str = ", ".join(arg_parts)
+            
             header = Text()
             header.append(tool_name, style="#00BCD4")
             header.append("(", 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(args_str, style="yellow")
             header.append(")", style="yellow")
-            
+            tool_time = 0 
+            start_time = time.time()
+            total_time = time.time() - START_TIME
+            timing_info = []
+            if total_time:
+                timing_info.append(f"Total: {format_time(total_time)}")
+            if tool_time:
+                timing_info.append(f"Tool: {format_time(tool_time)}")
+            if timing_info:
+                header.append(f" [{' | '.join(timing_info)}]", style="cyan")
+
             content = Text()
-            content.append(f"Executing: {command}\n\n", style="green")
             
             panel = Panel(
                 Text.assemble(header, "\n\n", content),
@@ -323,15 +351,35 @@ def _run_local_streamed(command, call_id, timeout=100):
             # Start Live display
             with Live(panel, console=console, refresh_per_second=4) as live:
                 # Stream stdout in real-time
+                start_time = time.time()
                 for line in iter(process.stdout.readline, ''):
                     if not line:
                         break
-                    
+
                     # Add to output collection
                     output_buffer.append(line)
-                    
+
                     # Update content with new line
                     content.append(line, style="bright_white")
+
+                    # Update tool_time and header with new timing info
+                    tool_time = time.time() - start_time
+                    total_time = time.time() - START_TIME 
+                    # Remove any previous timing info from header (rebuild header)
+                    timing_info = []
+                    if total_time:
+                        timing_info.append(f"Total: {format_time(total_time)}")
+                    if tool_time:
+                        timing_info.append(f"Tool: {format_time(tool_time)}")
+                    # Rebuild header to update timing
+                    header = Text()
+                    header.append(tool_name, style="#00BCD4")
+                    header.append("(", style="yellow")
+                    header.append(args_str, style="yellow")
+                    header.append(")", style="yellow")
+                    if timing_info:
+                        header.append(f" [{' | '.join(timing_info)}]", style="cyan")
+
                     panel = Panel(
                         Text.assemble(header, "\n\n", content),
                         title="[bold green]Tool Execution[/bold green]",
@@ -341,7 +389,6 @@ def _run_local_streamed(command, call_id, timeout=100):
                         box=ROUNDED
                     )
                     live.update(panel)
-                
                 # Check if process is done
                 process.stdout.close()
                 return_code = process.wait(timeout=timeout)
@@ -363,12 +410,9 @@ def _run_local_streamed(command, call_id, timeout=100):
                     live.update(panel)
                 
                 # Add completion message
-                completion_status = "Completed" if return_code == 0 else f"Failed (code {return_code})"
-                content.append(f"\nCommand {completion_status}", style="green")
                 panel = Panel(
                     Text.assemble(header, "\n\n", content),
                     title="[bold green]Tool Execution[/bold green]",
-                    subtitle=f"[bold green]{completion_status}[/bold green]",
                     border_style="green",
                     padding=(1, 2),
                     box=ROUNDED
@@ -383,10 +427,17 @@ def _run_local_streamed(command, call_id, timeout=100):
             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": ""}
+            
+            # Create a dictionary with only non-empty values (following the same rules)
+            tool_args = {}
+            if cmd:
+                tool_args["command"] = cmd
+            if args and args.strip():
+                tool_args["args"] = args
+            # Note: Omitted empty values and async_mode=False as it's default
             
             # Initial notification - just once
-            cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id)
+            cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id)
             
             # Buffer for collecting output 
             buffer_size = 0
@@ -404,7 +455,7 @@ def _run_local_streamed(command, call_id, timeout=100):
                 # Only update the output periodically to reduce panel refresh rate
                 if buffer_size >= update_interval:
                     current_output = ''.join(output_buffer)
-                    cli_print_tool_output("generic_linux_command", tool_args, current_output, call_id=call_id)
+                    cli_print_tool_output(tool_name, tool_args, current_output, call_id=call_id)
                     buffer_size = 0
             
             # Check if process is done
@@ -421,7 +472,7 @@ def _run_local_streamed(command, call_id, timeout=100):
             if return_code != 0:
                 final_output += f"\nCommand exited with code {return_code}"
                 
-            cli_print_tool_output("generic_linux_command", tool_args, final_output, call_id=call_id)
+            cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id)
         
         # Return the full output
         return ''.join(output_buffer)
@@ -434,7 +485,7 @@ def _run_local_streamed(command, call_id, timeout=100):
         # Update tool output panel with timeout message
         if not rich_available:
             tool_args = {"command": command}
-            cli_print_tool_output("generic_linux_command", tool_args, final_output, call_id=call_id)
+            cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id)
         
         return final_output
         
@@ -445,14 +496,14 @@ def _run_local_streamed(command, call_id, timeout=100):
         # Update tool output panel with error message if simple streaming
         if not rich_available:
             tool_args = {"command": command}
-            cli_print_tool_output("generic_linux_command", tool_args, error_msg, call_id=call_id)
+            cli_print_tool_output(tool_name, tool_args, error_msg, call_id=call_id)
         
         return error_msg
 
 
 def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arguments # noqa: E501
                 async_mode=False, session_id=None,
-                timeout=100, stream=False, call_id=None):
+                timeout=100, stream=False, call_id=None, tool_name=None):
     """
     Run command either in CTF container or on the local attacker machine
 
@@ -465,6 +516,8 @@ def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arg
         timeout: Command timeout in seconds
         stream: Whether to stream output in real-time
         call_id: Unique ID for the command execution (for streaming)
+        tool_name: Name of the tool being executed (for display in streaming output).
+                  If None, the tool name will be derived from the command.
 
     Returns:
         str: Command output, status message, or session ID
@@ -497,4 +550,4 @@ def run_command(command, ctf=None, stdout=False,  # pylint: disable=too-many-arg
     # Otherwise, run command normally
     if ctf:
         return _run_ctf(ctf, command, stdout, timeout, stream, call_id)
-    return _run_local(command, stdout, timeout, stream, call_id)
+    return _run_local(command, stdout, timeout, stream, call_id, tool_name)

src/cai/tools/reconnaissance/generic_linux_command.py

@@ -99,10 +99,11 @@ def generic_linux_command(command: str = "",
     if stream and not call_id:
         call_id = str(uuid.uuid4())[:8]
 
-    # Run the command with the appropriate parameters
+    # Run the command with the appropriate parameters - pass the actual tool name!
     result = run_command(full_command, ctf=ctf,
                        async_mode=async_mode, session_id=session_id,
-                       timeout=timeout, stream=stream, call_id=call_id)
+                       timeout=timeout, stream=stream, call_id=call_id,
+                       tool_name="generic_linux_command")
     
     # 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

src/cai/util.py

@@ -19,7 +19,13 @@ from datetime import datetime
 import atexit
 from dataclasses import dataclass, field
 from typing import Dict, Optional
+import time
 
+# Instead of direct import
+try:
+    from cai.cli import START_TIME
+except ImportError:
+    START_TIME = None
 
 # Shared stats tracking object to maintain consistent costs across calls
 @dataclass
@@ -521,7 +527,22 @@ def get_model_name(model):
         return model
     # If not a string, use environment variable
     return os.environ.get('CAI_MODEL', 'qwen2.5:72b')
-
+ # Helper function to format time in a human-readable way
+def format_time(seconds):
+            if seconds is None:
+                return "N/A"
+            
+            if seconds < 60:
+                return f"{seconds:.1f}s"
+            elif seconds < 3600:
+                minutes = int(seconds / 60)
+                seconds_remainder = seconds % 60
+                return f"{minutes}m {seconds_remainder:.1f}s"
+            else:
+                hours = int(seconds / 3600)
+                minutes = int((seconds % 3600) / 60)
+                return f"{hours}h {minutes}m"
+        
 def get_model_pricing(model_name):
     """
     Get pricing information for a model, using the CostTracker's implementation.
@@ -831,6 +852,13 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug,  # pyli
             text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})",
                         style="bold blue")
         text.append("]", style="dim")
+    elif not parsed_message: 
+        # When parsed_message is empty, only include timestamp and model info
+        text.append(f"Agent: {agent_name} ", style="bold green")
+        text.append(f"[{timestamp}", style="dim")
+        if model:
+            text.append(f" ({model})", style="bold magenta")
+        text.append("]", style="dim")
     else:
         text.append(f"[{counter}] ", style="bold cyan")
         text.append(f"Agent: {agent_name} ", style="bold green")
@@ -861,7 +889,9 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug,  # pyli
             interaction_cost,
             total_cost
         )
-        text.append(tokens_text)
+        # Only append token information if there is a parsed message
+        if parsed_message:
+            text.append(tokens_text)
 
     panel = Panel(
         text,
@@ -873,7 +903,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug,  # pyli
                else "[bold]Agent Interaction[/bold]"),
         title_align="left"
     )
-    console.print("\n")
+    #console.print("\n")
     console.print(panel)
     
     # If there are tool panels, print them after the main message panel
@@ -1052,7 +1082,6 @@ def finish_agent_streaming(context, final_stats=None):
     context["live"].update(final_panel)
     
     # Ensure updates are displayed before stopping
-    import time
     time.sleep(0.5)
     
     # Stop the live display
@@ -1069,12 +1098,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
         output: The output of the tool
         call_id: Optional call ID for streaming updates
         execution_info: Optional execution information
-        token_info: Optional token information
+        token_info: Optional token information with keys:
+            - interaction_input_tokens, interaction_output_tokens, interaction_reasoning_tokens
+            - total_input_tokens, total_output_tokens, total_reasoning_tokens
+            - model: model name string
+            - interaction_cost, total_cost: optional cost values
     """
     # If it's an empty output, don't print anything
     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'
@@ -1112,33 +1146,130 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
         from rich.panel import Panel
         from rich.text import Text
         from rich.box import ROUNDED
+        from rich.console import Group
         
         # Create a console for output
-        console = Console()
+        console = Console(theme=theme)
         
-        # Format arguments as a string if they are a dictionary
+        # Format arguments for display
+        # Parse JSON string if args is a string
+        if isinstance(args, str) and args.strip().startswith('{'):
+            try:
+                import json
+                args = json.loads(args)
+            except:
+                # Keep as is if not valid JSON
+                pass
+        
+        # Format arguments as a clean string
         if isinstance(args, dict):
-            args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()])
+            # Only include non-empty values and exclude async_mode=false
+            arg_parts = []
+            for key, value in args.items():
+                # Skip empty values
+                if value == "" or value == {} or value is None:
+                    continue
+                # Skip async_mode=false (default)
+                if key == "async_mode" and value is False:
+                    continue
+                # Format the value
+                if isinstance(value, str):
+                    arg_parts.append(f"{key}={value}")
+                else:
+                    arg_parts.append(f"{key}={value}")
+            args_str = ", ".join(arg_parts)
         else:
             args_str = str(args)
         
+        # Get session timing information
+        try:
+            from cai.cli import START_TIME
+            total_time = time.time() - START_TIME if START_TIME else None
+        except ImportError:
+            total_time = None
+            
+        
+        # Extract execution timing info
+        
+        tool_time = None
+        status = None
+        if execution_info:
+            # Prefer 'tool_time' if present, else fallback to 'time_taken'
+
+            tool_time = execution_info.get('tool_time') 
+            status = execution_info.get('status', 'completed')
+        
+        # Create header for all panel displays (both streaming and non-streaming)
+        header = Text()
+        header.append(tool_name, style="#00BCD4")
+        header.append("(", style="yellow")
+        header.append(args_str, style="yellow")
+        header.append(")", style="yellow")
+        
+        # Add timing information directly in the header
+        timing_info = []
+        if total_time:
+            timing_info.append(f"Total: {format_time(total_time)}")
+        if tool_time:
+            timing_info.append(f"Tool: {format_time(tool_time)}")
+        if timing_info:
+            header.append(f" [{' | '.join(timing_info)}]", style="cyan")
+           
+        # Add completion status if available - REMOVED, just showing timing now
+        # if status:
+        #     if status == 'completed':
+        #         header.append(f" [Completed]", style="green")
+        #     elif status == 'running':
+        #         header.append(f" [Running]", style="yellow")
+        #     elif status == 'error':
+        #         header.append(f" [Error]", style="red")
+        #     elif status == 'timeout':
+        #         header.append(f" [Timeout]", style="red")
+        #     else:
+        #         header.append(f" [{status.title()}]", style="dim")
+        
         # For streaming mode with call_id, use Rich Live display
         if call_id:
-            # Create the header text
-            header = Text()
-            header.append(tool_name, style="#00BCD4")
-            header.append("(", style="yellow")
-            header.append(args_str, style="yellow")
-            header.append(")", style="yellow")
+            # Create token information if available
+            token_content = None
+            if token_info:
+                model = token_info.get('model', '')
+                interaction_input_tokens = token_info.get('interaction_input_tokens', 0)
+                interaction_output_tokens = token_info.get('interaction_output_tokens', 0)
+                interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0)
+                total_input_tokens = token_info.get('total_input_tokens', 0)
+                total_output_tokens = token_info.get('total_output_tokens', 0)
+                total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0)
+                
+                if (interaction_input_tokens > 0 or total_input_tokens > 0):
+                    token_text = _create_token_display(
+                        interaction_input_tokens,
+                        interaction_output_tokens,
+                        interaction_reasoning_tokens,
+                        total_input_tokens,
+                        total_output_tokens,
+                        total_reasoning_tokens,
+                        model,
+                        token_info.get('interaction_cost'),
+                        token_info.get('total_cost')
+                    )
+                    token_content = Text("\n\n")
+                    token_content.append(token_text)
             
             # Create content text with the output
             content = Text(output)
             
-            # Create the panel for display
+            # Create the panel for display, including token info if available
+            panel_content = [header, Text("\n\n"), content]
+            if token_content:
+                panel_content.append(token_content)
+                
+            # Create title - simple title with no timing info
+            title = "[bold blue]Tool Output[/bold blue]"
+            
             panel = Panel(
-                Text.assemble(header, "\n\n", content),
-                title=f"[bold blue]Tool Output[/bold blue]", #(ID: {call_id})[/bold green]",
-                #subtitle="[bold green]Live Update[/bold green]",
+                Text.assemble(*panel_content),
+                title=title,
                 border_style="blue",
                 padding=(1, 2),
                 box=ROUNDED,
@@ -1149,52 +1280,59 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
             console.print(panel)
             return
             
-        # For non-streaming output, show a standard panel
-        tool_call = f"{tool_name}({args_str})"
-        
-        # Add execution info if available
-        subtitle = "[bold green]Completed[/bold green]"
-        if execution_info:
-            time_taken = execution_info.get('time_taken', 0)
-            status = execution_info.get('status', 'completed')
-            
-            if time_taken:
-                subtitle = f"[bold green]{status.title()} in {time_taken:.2f}s[/bold green]"
-            else:
-                subtitle = f"[bold green]{status.title()}[/bold green]"
-        
-        # Create content sections
-        sections = []
-        
-        # Add token info if available
+        # For non-streaming output, also use a blue panel with the same format
+        # Create token information if available
+        token_text = None
         if token_info:
-            tokens_in = token_info.get('input_tokens', 0)
-            tokens_out = token_info.get('output_tokens', 0)
-            cost = token_info.get('cost', 0)
+            model = token_info.get('model', '')
+            interaction_input_tokens = token_info.get('interaction_input_tokens', 0)
+            interaction_output_tokens = token_info.get('interaction_output_tokens', 0)
+            interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0)
+            total_input_tokens = token_info.get('total_input_tokens', 0)
+            total_output_tokens = token_info.get('total_output_tokens', 0)
+            total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0)
+            interaction_cost = token_info.get('interaction_cost')
+            total_cost = token_info.get('total_cost')
             
-            token_text = Text()
-            if tokens_in or tokens_out:
-                token_text.append(f"Tokens: {tokens_in} in, {tokens_out} out", style="cyan")
-            if cost:
-                token_text.append(f"\nCost: ${cost:.6f}", style="cyan")
-                
-            if token_text:
-                sections.append(token_text)
+            # Generate token display with CostTracker
+            if (interaction_input_tokens > 0 or total_input_tokens > 0):
+                token_text = _create_token_display(
+                    interaction_input_tokens,
+                    interaction_output_tokens,
+                    interaction_reasoning_tokens,
+                    total_input_tokens,
+                    total_output_tokens,
+                    total_reasoning_tokens,
+                    model,
+                    interaction_cost,
+                    total_cost
+                )
         
-        # Add the main output
+        # Now create the panel content, starting with the header
+        panel_content = [header, Text("\n")]
+        
+        # Add token display if available
+        if token_text:
+            panel_content.append(token_text)
+            panel_content.append(Text("\n"))  # Add spacing after token display
+            
+        # Add the output
         if output:
-            output_text = Text()
-            if sections:  # Add a separator if we have previous sections
-                output_text.append("\n")
-            output_text.append(output)
-            sections.append(output_text)
+            output_text = Text(output)
+            panel_content.append(output_text)
         
-        # Create the panel with all sections
+        # If no content was added but we have output, add it directly
+        if len(panel_content) == 2 and output:  # Only header and newline
+            panel_content.append(Text(output))
+        
+        # Create title - simple title with no timing info
+        title = "[bold blue]Tool Output[/bold blue]"
+        
+        # Create the final panel - always blue now
         panel = Panel(
-            Text.assemble(*sections),
-            title=f"[bold green]Tool Output: {tool_call}[/bold green]",
-            subtitle=subtitle,
-            border_style="green",
+            Group(*panel_content),
+            title=title,
+            border_style="blue",
             padding=(1, 2),
             box=ROUNDED
         )
@@ -1204,64 +1342,112 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
         
     except ImportError:
         # Fall back to simple formatting if Rich is not available
-        # Format arguments as a string if they are a dictionary
+        # Format arguments in the cleaner format
+        # Parse JSON string if args is a string
+        if isinstance(args, str) and args.strip().startswith('{'):
+            try:
+                import json
+                args = json.loads(args)
+            except:
+                # Keep as is if not valid JSON
+                pass
+        
+        # Format arguments as a clean string
         if isinstance(args, dict):
-            args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()])
+            # Only include non-empty values and exclude async_mode=false
+            arg_parts = []
+            for key, value in args.items():
+                # Skip empty values
+                if value == "" or value == {} or value is None:
+                    continue
+                # Skip async_mode=false (default)
+                if key == "async_mode" and value is False:
+                    continue
+                # Format the value
+                if isinstance(value, str):
+                    arg_parts.append(f"{key}={value}")
+                else:
+                    arg_parts.append(f"{key}={value}")
+            args_str = ", ".join(arg_parts)
         else:
             args_str = str(args)
-        
-        # Simplify output presentation for streaming mode
-        if call_id:
-            # This is a streaming update, so we need to overwrite previous output
-            # We'll use a basic format that's better suited for streaming
+        # Get session timing information
+        try:
+            from cai.cli import START_TIME
+            total_time = time.time() - START_TIME if START_TIME else None
+        except ImportError:
+            total_time = None
             
-            # Get terminal width for better formatting
-            try:
-                term_width = os.get_terminal_size().columns
-            except:  # pylint: disable=bare-except
-                term_width = 80
-                
-            # Create a header for the tool output
-            header = f"{tool_name}({args_str})"
-            header = header[:term_width-4]
-            
-            # Clear the screen for the tool output (alternative approach)
-            print(f"\r{header}")
-            
-            # Print the content
-            # For streaming updates, we'll use a simple format
-            print(output)
-            
-            # Force flush to ensure output is displayed immediately
-            sys.stdout.flush()
-            return
         
         # For non-streaming output, use the original formatting
         tool_call = f"{tool_name}({args_str})"
         
-        # If there's execution info, add it to the output
+        # Get tool execution time if available
+        tool_time_str = ""
+        execution_status = ""
         if execution_info:
             time_taken = execution_info.get('time_taken', 0)
             status = execution_info.get('status', 'completed')
             
             # Add execution info to the tool call display
             if time_taken:
-                tool_call += f" [{status} in {time_taken:.2f}s]"
+                tool_time_str = f"Tool: {format_time(time_taken)}"
+                execution_status = f" [{status} in {time_taken:.2f}s]"
             else:
-                tool_call += f" [{status}]"
+                execution_status = f" [{status}]"
         
-        print(color(f"Tool Output: {tool_call}", fg="blue"))
-        
-        # If we have token info, display it
-        if token_info:
-            tokens_in = token_info.get('input_tokens', 0)
-            tokens_out = token_info.get('output_tokens', 0)
-            cost = token_info.get('cost', 0)
+        # Create timing display string
+        timing_info = []
+        if total_time:
+            timing_info.append(f"Total: {format_time(total_time)}")
+        if tool_time:
+            timing_info.append(f"Tool: {format_time(tool_time)}")
+        if timing_info:
+            header.append(f" [{' | '.join(timing_info)}]", style="cyan")
             
-            if tokens_in or tokens_out:
-                print(color(f"  Tokens: {tokens_in} in, {tokens_out} out", fg="cyan"))
-            if cost:
-                print(color(f"  Cost: ${cost:.6f}", fg="cyan"))
+        timing_display = f" [{' | '.join(timing_info)}]" if timing_info else ""
+        
+        # Show tool name, args, execution status and timing display
+        print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue"))
+        
+        # If we have token info, display it using the consistent format from _create_token_display
+        if token_info:
+            model = token_info.get('model', '')
+            interaction_input_tokens = token_info.get('interaction_input_tokens', 0)
+            interaction_output_tokens = token_info.get('interaction_output_tokens', 0)
+            interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0)
+            total_input_tokens = token_info.get('total_input_tokens', 0)
+            total_output_tokens = token_info.get('total_output_tokens', 0)
+            total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0)
+            interaction_cost = token_info.get('interaction_cost')
+            total_cost = token_info.get('total_cost')
+            
+            # If we have complete token information, display it
+            if (interaction_input_tokens > 0 or total_input_tokens > 0):
+                # Manually create formatted output similar to _create_token_display
+                print(color(f"  Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan"))
+                
+                # Calculate or use provided costs
+                current_cost = COST_TRACKER.process_interaction_cost(
+                    model, 
+                    interaction_input_tokens, 
+                    interaction_output_tokens,
+                    interaction_reasoning_tokens,
+                    interaction_cost
+                )
+                total_cost_value = COST_TRACKER.process_total_cost(
+                    model,
+                    total_input_tokens,
+                    total_output_tokens,
+                    total_reasoning_tokens,
+                    total_cost
+                )
+                print(color(f"  Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan"))
+                
+                # Show context usage
+                context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100
+                indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "🟥"
+                print(color(f"  Context: {context_pct:.1f}% {indicator}", fg="cyan"))
         
         # Print the actual output
         print(output)

tests/tools/test_tool_generic_linux_command.py

@@ -1,11 +1,11 @@
-"""
-This module contains tests for the generic Linux command tool functionality.
-It includes tests for executing regular commands and handling invalid commands.
-"""
-
+import os
 import pytest
 import json
 from unittest.mock import MagicMock
+
+# Set test environment variables to avoid OpenAI client initialization errors
+os.environ["OPENAI_API_KEY"] = "test_key_for_ci_environment"
+
 from cai.tools.reconnaissance.generic_linux_command import generic_linux_command  
 
 async def test_generic_linux_command_regular_commands():
@@ -48,4 +48,4 @@ async def test_generic_linux_command_invalid_command():
     result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params))
 
     # Assert that the result indicates the command was not found
-    assert "not found" in result  
+    assert "not found" in result